arbi-serve — OpenAI-compatible inference server built around TKV.

The engine is designed around ``tkv.runtime.attention.TKVCore``: the
paged uint8 KV slab is the native cache layout, and metadata fields and
CUDAGraph capture are shaped to match the codec without adapters.

Public entry points:
  - :func:`arbi_serve.cli.main` — ``arbi-serve`` console script.
  - :class:`arbi_serve.engine.engine.Engine` — programmatic API.
  - :class:`arbi_serve.config.ServerConfig` — typed config.

``python -m arbi_serve`` entrypoint.

Dispatches ``dump_openapi`` through a light entrypoint so it doesn't pay
the ~1.5 s of ``import arbi_serve.cli`` (which transitively imports
torch, uvicorn, every backend): it only needs FastAPI + the router
modules (no torch, no engine), so routing it through the heavy CLI would
be wasteful. The serve / warmup / calibrate paths go through the full CLI.

Split this card between us and everyone else, before we hold any of it.

Whatever is resident on a device right now belongs to some other process:
nothing here has made a CUDA call yet. That one reading is what later lets
the boot tell its own driver residency from a co-tenant's — and on hosts
where NVML reports pids in a namespace that is not ours, so its
per-process walk can never match ``os.getpid()``, it is the only thing that
can. The KV pool is the card's residual, so a co-tenant booked as ours is
served context lost for the life of the boot.

It has to run HERE: ``import arbi_serve.cli`` pulls torch and the backends,
and something on that path creates the primary context, so any mark taken
after it already contains our own bytes. NVML only — no torch import is
pulled forward. Best-effort: a boot must never depend on it.

Boot-time self-heal broom for the persistent shared caches.

Every arbi-serve on-disk cache writer that lands a *final* entry does so
atomically — ``mkstemp`` (or ``<name>.tmp``) in the same directory, then
``os.replace`` — so a crash mid-write can never leave a half-written
FINAL file that a later boot reads as valid (budget cache, boot manifest,
boot artifacts, k-calibration, compile Mega-Cache, tkv/calibration
bundles all follow this discipline). What a crash (OOM / IMA / SIGKILL —
all of which the fleet triggers constantly) DOES leave behind is the
**orphaned temp file**: ``<key>.<rand>.json.tmp`` / ``<name>.tmp`` that the
``os.replace`` never reached.

Those orphans are harmless in isolation (no reader ever opens a ``*.tmp``
— readers only open the final ``<key>.json``), but the compile caches
persist across boots on the shared ``/cache`` volume, so with a
crash-happy fleet they accumulate without bound and clutter the operator's
``ls`` of the cache root, and a stale ``.arbi_writable_probe`` from a
crashed writability probe lingers too. This broom sweeps them once at
boot.

WHAT THIS DELIBERATELY DOES NOT TOUCH
-------------------------------------
  * **The torch-owned compile caches** (Inductor FS cache
    ``/cache/inductor`` / ``$TORCHINDUCTOR_CACHE_DIR``, ``$TRITON_CACHE_DIR``,
    ``$TORCH_EXTENSIONS_DIR``): these are managed by torch/Triton, which
    write their artifacts atomically and validate them by content hash —
    a partial one misses by key and recompiles. They are also HUGE on a
    shared ``/cache`` (thousands of hashed subdirs; a full ``os.walk``
    can take tens of seconds on NFS), so recursing them at boot would be
    the very boot-stall this broom exists to prevent. We sweep ONLY the
    arbi-serve-owned *leaf* cache dirs, and even those under a depth AND
    wall-clock bound.
  * **Final cache entries** (``<key>.json`` / ``*.pt`` / ``*.bin``): every
    reader already version- + fingerprint-validates them and self-heals a
    corrupt/stale one (returns "miss" → recompute). They are atomic, so
    they are never partial. Sweeping them would throw away valid warm-boot
    state.
  * **Lock files** (``lock`` / ``.ninja_lock``): the torch cpp-extension
    build baton is made self-healing by :mod:`arbi_serve._jit_baton`
    (PID-aware, PID-reuse-aware) and swept by
    :mod:`arbi_serve._cpp_ext_locks` only when the built ``.so`` proves the
    build finished. Age-blind removal here could race a live build, so we
    leave locks to those owners. Triton's own ``filelock`` locks are
    advisory ``flock``s the kernel releases on process death, so a dead
    holder never wedges a later boot regardless.

SAFETY AGAINST CONCURRENT / LIVE WRITES
---------------------------------------
A temp is only removed when it is OLDER than ``min_age_s`` (default 15
min). A live atomic write holds its temp for milliseconds; even a long
calibration/capture sweep writes its final temp in well under a second
once the payload is computed. So the age gate can never race an in-flight
write from THIS boot or a co-resident TP rank booting at the same time. A
``FileNotFoundError`` (another rank swept the same orphan first) is
ignored.

Torch-free, best-effort, and never raises into boot — it is called from
``cli.bootstrap._set_default_env`` right after the compile-cache pin, on
every serve launch path (including each torchrun TP rank).

Resolve the persistent cache roots to sweep, torch-free.

ONLY arbi-serve-owned *leaf* cache dirs — never the torch-managed
Inductor / Triton / cpp-extension trees (they are atomic + hash-keyed on
their own, and huge on a shared ``/cache``: recursing them at boot would
be the boot-stall this broom prevents). Honors the same env overrides the
real cache modules read, then falls back to their default paths. Kept
env-only (no ``runtime_flags`` / torch import) so this stays importable on
the earliest, torch-free boot hook.

Remove orphaned atomic-write temp files under the cache ``roots``.

Only files whose name is provable write scaffolding (``*.tmp`` /
``.arbi_writable_probe``) AND older than ``min_age_s`` are removed.
Final cache entries and lock files are never touched. Returns the number
of files removed. Best-effort — every filesystem error is swallowed, so
the sweep can never raise into boot.

Args:
    roots: directories to sweep. ``None`` → :func:`_default_cache_roots`.
    min_age_s: minimum age (seconds) an orphan must reach before removal,
        so a live/concurrent atomic write is never disturbed.
    now: injectable clock (seconds since epoch) for tests.

Depth- and time-bounded sweep of one cache root. Never raises.

Explicit ``os.scandir`` stack (not ``os.walk``) so the wall-clock budget
is enforced per-ENTRY, not per-directory — a huge NFS dir with thousands
of subdirs can never blow the budget. Each directory's files are handled
before its subdirs are pushed, so the flat depth-0 orphans (the common,
load-bearing case) are always swept before the budget could expire.
``deadline`` is a shared ``time.monotonic()`` cutoff across all roots.

Pin the compile-cache env vars to the persistent ``/cache`` volume,
BEFORE torch can self-stamp them ephemeral.

WHY THIS IS ITS OWN (TORCH-FREE) MODULE, CALLED FROM ``arbi_serve/__init__``
---------------------------------------------------------------------------
torch 2.12's ``torch._inductor.runtime.cache_dir_utils.cache_dir()`` has a
WRITE side effect: when ``$TORCHINDUCTOR_CACHE_DIR`` is unset, it stamps
``os.environ["TORCHINDUCTOR_CACHE_DIR"] = <tempdir>/torchinductor_<user>``
(``/tmp/torchinductor_root`` in the container). And ``import torch._dynamo``
calls it at MODULE level (``torch/_dynamo/package.py`` constructs the global
``DynamoCache = DiskDynamoCache(os.path.join(cache_dir(), "dynamo"))``),
which ``arbi_serve/__init__`` reaches transitively the moment it imports
``_custom_ops`` (``@torch._dynamo.assume_constant_result``).

The old pin lived only in ``cli.bootstrap._set_default_env()`` — which runs
inside ``main()``, i.e. AFTER ``from arbi_serve.cli import main`` has already
imported the package and its torch side-effect imports. By then torch had
self-stamped the env, so the bootstrap ``setdefault`` was a no-op and
``configure_inductor_caches``'s resolver honored the stamped ``/tmp`` path as
if it were an operator override. Net effect: the Inductor/Triton compile
cache landed in container-ephemeral ``/tmp``, every boot re-ran codegen +
combo-kernel benchmarking + autotune, and whatever kernels that lottery
produced were baked into the cudagraph capture.

The fix: :func:`pin_compile_cache_env` is torch-free and is called from
``arbi_serve/__init__.py`` BEFORE its ``import torch`` — the one spot that
precedes every torch import on every entry path (``python -m arbi_serve``
imports the package ``__init__`` before ``__main__`` runs; torchrun TP ranks,
the console script, and programmatic ``Engine(cfg)`` embedders all import the
package first). ``cli.bootstrap._set_default_env`` still calls it (harmless
re-run) so light subcommands and any exotic entry keep the pin.

Defense in depth: torch's self-stamped value is byte-predictable
(:func:`torch_ephemeral_inductor_default` mirrors torch's
``default_cache_dir()``), so when we see EXACTLY that value in the env we
treat it as "unset by any operator" and overwrite it — recovering even on
paths where torch got imported before arbi_serve (e.g. an embedder that
imports torch first). A genuinely operator-chosen dir never matches it.

The exact path torch's ``default_cache_dir()`` self-stamps into
``$TORCHINDUCTOR_CACHE_DIR`` — ``<tempdir>/torchinductor_<user>``.

Mirrors ``torch._inductor.runtime.cache_dir_utils.default_cache_dir``
byte-for-byte (same ``getpass.getuser()`` sanitization) WITHOUT importing
torch, so callers can distinguish "torch stamped its ephemeral default"
from "an operator deliberately exported a cache dir". fbcode's
``/var/tmp`` variant is irrelevant here (we never run fbcode builds).

Fleet-shared torch-extensions dir, keyed per torch/python/CUDA build.

``root`` (the persistent cache base, ``/cache`` in production) is shared
across hosts and venvs that do NOT all run the same torch or Python — and
torch only appends its own ``py{ver}_cu{ver}`` folder to the DEFAULT cache
root; a set ``TORCH_EXTENSIONS_DIR`` is used verbatim. Without the key, one
venv's ``.so`` (wrong torch/Python ABI) is served to another. The key
helper is torch-free, so this module stays torch-free.

Versioned, uid-namespaced Inductor/Triton compile-cache root.

``<root>/arbi-serve/inductor-<torch-py-cu>/u<uid>``. Keyed by the same
torch/Python/CUDA ``cpp_ext_cache_key`` as the torch-extensions dir: a
compiled Inductor artifact (and the Triton cubins it references) is bound
to the exact torch build that emitted it, so a torch bump on a SHARED
``root`` must land in a fresh subtree — otherwise a warm boot could load
a mismatched artifact that faults at runtime. Uid-namespaced for the same
cross-user-poison safety as :func:`default_torch_extensions_dir` and the
engine's own ``_resolve_inductor_cache_dir`` (a root-owned artifact from a
bare-box container run must never be served to another uid). The key
helper is torch-free, so this module stays torch-free.

Create ``path`` (parents included) and probe it is actually writable.

torch's ``cache_dir()`` does ``os.makedirs(<env value>)`` and RAISES on
failure, so pinning ``$TORCHINDUCTOR_CACHE_DIR`` to a dir this uid cannot
create would turn a cache miss into an import-time crash. Observed in the
wild: a prior ROOT container run left ``/cache/arbi-serve/inductor-<key>``
as ``drwxr-xr-x root`` on the shared volume, so a later non-root run could
not create its ``u<uid>`` subtree. To keep the shared volume multi-uid
(the whole point of the uid namespacing), when WE create ``shared_parent``
we best-effort chmod it sticky+world-writable (``01777``, like ``/cache``
itself on the fleet) so OTHER uids can still create their own subtrees
later. Returns True only when ``path`` exists and a write probe succeeds.

The probe name is PER-PROCESS and its unlink tolerates a missing file.
Every rank of a TP>1 boot runs this concurrently against the SAME shared
``/cache``, so a fixed probe name races: rank0 unlinks the file rank1 just
opened, rank1's own ``unlink`` raises ``ENOENT``, and rank1 concludes
``/cache`` is unwritable. It then falls back to ``~/.cache/arbi`` — an
EPHEMERAL source — and ``assert_warm_cache_or_die`` kills the boot with a
message blaming a missing bind-mount that is in fact correctly mounted.
Observed on a TP2 27B bench: rank0 logged ``source=/cache volume`` while
rank1 died on ``WarmCacheMiss`` in the same container.

One compile cache's boot-time status.

``present`` is the HONEST per-cache predictor of this boot's cost: True iff
the dir already holds artifacts, so this cache is REUSED; False ⇒ it is
(re)COMPILED / JIT'd this boot. A binary whole-boot warm/cold flag hides
that some caches hit while others miss — this is per-component.

Refuse a declared phase the metrics layer does not know.

A typo here is invisible at runtime — the panel simply colours the phase as
uncached — so it is caught at import instead.

``(compat_key, root_source, components)`` from the last pin, or ``None``.

``compat_key`` (``torch{ver}-cu{cuda}-py{maj}.{min}``) is the build key
every cache entry is namespaced / hashed under, so incompatible artifacts
can never be loaded as-if-compatible — a version mismatch is a cache MISS
(recompile), not poison. This is the poisoning guard, and it holds on cold,
warm, and hot-swap boots because the key travels with each entry (the path
namespacing is additional hygiene). ``components`` lists each cache's
per-boot REUSE/COMPILE status so a partial-hit boot is fully visible.

Append ``components`` to the stamped report, replacing same-named ones.

The report is built torch-free at import-pin time, so caches whose status
only a torch-aware caller can resolve (the persisted compile blobs, whose
path is keyed on the torch/GPU build) register themselves here before the
boot logs the report and :func:`assert_warm_cache_or_die` reads it.
No-op when nothing has been pinned yet.

A boot that must have been WARM re-compiled — raised by
:func:`assert_warm_cache_or_die` under ``ARBI_ASSERT_WARM_CACHE``.

Two triggers, both root-cause bugs, never model behaviour:
  * the resolved cache root is EPHEMERAL (``/cache`` not bind-mounted); or
  * a cache that a PRIOR boot published for this build key is now MISSING.

Per-(root, build-key, uid) marker recording which caches have EVER been
published on this key. Lives ON the persistent cache root, so it survives
exactly when the caches it vouches for do — a wiped ``/cache`` takes the
baseline with it and the next boot is (correctly) treated as first-ever.

Fail the boot LOUDLY if a warm compile-cache was silently missed.

No-op unless ``enabled`` (``ARBI_ASSERT_WARM_CACHE``). Given the same
``report`` the boot log already surfaced, it enforces two invariants and
otherwise records this boot's published caches as the warm baseline:

  1. **Persistent root.** ``cache_root``'s ``source`` must be the ``/cache``
     volume or an explicit ``$ARBI_CACHE_DIR`` — never an ephemeral
     ``~/.cache`` / ``/tmp`` fallback. An ephemeral root means ``/cache``
     was not bind-mounted, the exact footgun this insurance exists for.
  2. **No warm→cold regression.** Any managed cache that a PRIOR boot
     published for this build key MUST still be present. If it vanished the
     mount was lost or the cache evicted, and this boot is silently paying
     the full re-compile — refuse. A first-ever boot on a new build key has
     no baseline, so it is allowed to compile and publish.

Raises :class:`WarmCacheMiss` on violation; the message names exactly what
to fix (provision/bind ``/cache``).

True iff ``root/name`` is a compiled artifact rather than scaffolding.

An artifact must carry bytes: every managed cache's entries are compiled
output (cubin / ``.so`` / codegen / tune table), so a zero-length file is a
marker or a truncated write, not something a boot can reuse.

True iff ``path`` holds at least one compiled artifact (bounded scan).

``present`` decides whether a cache is REUSED this boot, and
:func:`assert_warm_cache_or_die` persists every present cache into the warm
baseline — so only a real artifact may set it (:func:`_is_cache_artifact`).
A stranded lock or probe vouching for a cache that was never published
would kill the next honest boot on a warm→cold "regression" that never
happened.

A warm Inductor cache is a deep tree; finding ONE artifact is enough to
classify the boot warm, so the walk short-circuits on the first hit and
caps its breadth to stay cheap at import time.

Best writable cache base, in priority order. ``(root, source)``.

``root`` is ``None`` only when not even a stable ``/tmp`` dir is writable
(torch's own default then stands). Priority:

  1. ``$ARBI_CACHE_DIR`` — explicit operator override.
  2. ``/cache`` — the production volume (image default; a ``-v`` mount or
     the baked anonymous volume).
  3. ``~/.cache/arbi`` — per-user; the persistent home for bare-metal /
     non-container runs where ``/cache`` does not exist.
  4. ``/tmp/arbi-cache`` — a STABLE last resort (a fixed path, unlike
     torch's per-invocation default, so a host still reuses it in-life).

NFS roots (``/mnt/k8scache``) are deliberately NOT auto-selected: cold
Operators wanting the fleet-shared cache point ``ARBI_CACHE_DIR`` at it.

Pin the JIT/compile cache roots to the persistent ``/cache`` volume.

Sets (respecting genuine operator overrides — see the ephemeral-stamp
carve-out below):

  * ``TORCH_EXTENSIONS_DIR`` — JIT'd cpp-extension kernels (tkv et al.),
    version-keyed; one nvcc compile per kernel version warms every host.
  * ``TKV_CACHE_DIR`` — tkv autotune tune-tables + kernel cache (the
    registry default is per-box ``~/.cache/tkv``; the warm-tune sweep is
    ~9 min per new shape fingerprint — fleet-share it).
  * ``CUTE_DSL_CACHE_DIR`` — the arbi-prefill CuTeDSL JIT cache; unpinned it
    recompiles the prefill kernels every boot (nothing bakes them).
  * ``TORCHINDUCTOR_CACHE_DIR`` — the Inductor FS cache root, pinned to a
    torch/CUDA/py-versioned, uid-namespaced subdir of ``/cache`` so warm
    boots REUSE compiled artifacts (FX-graph cache, autotune results,
    combo-kernel benchmark picks) instead of re-JITing into a container-
    ephemeral ``/tmp/torchinductor_<user>``.
  * ``TRITON_CACHE_DIR`` — pinned UNDER the same persisted root: torch
    2.12's ``triton_cache_dir()`` does NOT track
    ``$TORCHINDUCTOR_CACHE_DIR`` and the persisted FX-graph bundle
    references its cubins by ABSOLUTE path, so co-locating them is what
    makes the warm-boot static-launcher reload resolve instead of
    missing + re-JITing ("Cubin file saved by TritonBundler not found").
  * ``TORCHINDUCTOR_FX_GRAPH_CACHE`` / ``TORCHINDUCTOR_AUTOGRAD_CACHE`` /
    ``TORCHINDUCTOR_BUNDLE_TRITON_INTO_FX_GRAPH_CACHE`` — enable the
    on-disk caches + bundle the cubins INTO the FX-graph entry, pinned so
    an upstream default flip can't silently reintroduce the re-JIT.

Ephemeral-stamp carve-out: a pre-existing ``$TORCHINDUCTOR_CACHE_DIR``
normally wins (operator override) — EXCEPT when it equals torch's
self-stamped ephemeral default (see module docstring), which no operator
meaningfully chooses; that value is REPLACED with the persistent pin.

Never silently ephemeral: resolves the best writable persistent root
(:func:`_resolve_cache_root`) and stamps :data:`_CACHE_STATUS` so the boot
path can log it LOUDLY. Torch-free, idempotent, and safe to call multiple
times.

Cross-process serialization + stale-lock hygiene for torch JIT builds.

Two tools, both invoked from the global ``cpp_extension.load`` wrapper in
:mod:`arbi_serve._prebake_loader`, so they cover **every** JIT extension
(exl3, AWQ Marlin, xgrammar, and any future one) at one chokepoint:

1. :func:`extension_serial_lock` — an ``flock``-based EXCLUSIVE lock held
   around the **entire** ``load()`` call (version check, ninja build, and
   the ``dlopen``). This is the fix for the multi-rank JIT build race:
   torch's ``FileBaton`` is a presence-file lock with heuristic staleness
   recovery, and the recovery path is a TOCTOU — a contender that judged
   an *abandoned* lock stale can ``os.remove`` a *fresh* lock a new
   builder re-created at the same path, at which point two ranks build
   and link the same ``.so`` concurrently and a third dlopens the
   partially-linked file. ``flock`` has none of that: the kernel owns
   the lock state and releases it when the holder dies (crash, OOM,
   SIGKILL), so there is no staleness heuristic and nothing to steal.
   One rank builds; the others enter
   torch's ``load()`` only after the ``.so`` is complete and dlopen a
   finished file.

2. :func:`clear_stale_extension_lock` — removes a dead ``FileBaton`` lock
   file left by a killed build (torch never releases it on process death,
   so every later ``load()`` of that extension would spin in
   ``FileBaton.wait()`` forever — even when the ``.so`` is already fully
   built). Called INSIDE the serial lock, so for arbi processes the
   "concurrent live builder" hazard is structurally gone; the ``.so``
   present + min-age guards remain for non-arbi contenders.

Hold an exclusive kernel-managed ``flock`` for JIT extension ``name``.

Serializes concurrent arbi processes (TP ranks, co-tenant boots on a
shared ``TORCH_EXTENSIONS_DIR``) through the whole build-or-load of one
extension. The lock file is a dedicated zero-byte sentinel next to the
extension build dirs — NOT torch's ``lock`` FileBaton file (torch
manages that one itself; ours wraps it entirely).

Robustness properties (all inherited from ``flock`` semantics):
  - released by the kernel when the holder exits or is killed — a
    SIGKILLed build never wedges the next boot;
  - no staleness heuristic, hence no steal path and no TOCTOU;
  - re-entrant across ``load()`` calls in one process (each call opens
    its own fd; nested holds of the SAME extension don't occur — one
    ``load()`` per extension at a time per process).

Falls back to lock-free (a warning, not a crash) only when the lock
file cannot be created at all (read-only extensions root — e.g. the
baked ``/opt/cache-baked`` dir in the slim image, where no build can
happen anyway and the prebaked ``.so`` short-circuits before this).

Remove a stale ``lock`` / ``.ninja_lock`` for JIT extension ``name``.

Covers BOTH build-dir layouts torch uses: ``<root>/<name>`` (when
``TORCH_EXTENSIONS_DIR`` is set — the rig layout) and
``<root>/<pyver_accel>/<name>`` (the default root layout). The
original glob only matched the second, so the rig layout was silently
unguarded.

Best-effort and safe: only acts on a build dir whose ``<name>.so`` is
already present (build finished → any lock is stale) and whose lock is
older than ``min_age_s``. Callers in arbi hold
:func:`extension_serial_lock`, so a concurrent ARBI builder cannot be
racing this; the two guards protect against non-arbi processes. The
``.so``-present guard is deliberately conservative — during a REBUILD
(stale ``.so`` still on disk, new build in progress) a >min_age lock
could be live, which is exactly why this must only ever run under the
serial lock.

arbi-serve custom-op registry for the torch.compile pipeline.

Every kernel arbi-serve calls on the hot path is registered here as a
``torch.library`` custom op with a ``register_fake`` impl. Inductor
needs the fakes to plan allocations through these "black-box" kernels.

The wired ops route their production hot paths through
``torch.ops.arbi_serve.*``. Each ``op_func`` body dispatches into
the underlying kernel call the backend makes directly. The fakes
remain — Inductor consults them at trace time, the real-impl runs at
eager time.

The wired ops (real-impl filled in):
    * ``mla_attention``            — MLA explicit-decompress reference path.
    * ``tkv_attention``           — TKVCore.forward dispatch.
    * ``gdn_attention_v2``         — FLA delta-rule call site (decode +
      verify per-step + chunk prefill), slab-mutating. Takes the whole
      recurrent slab + ``state_indices`` and does gather + FLA-kernel +
      scatter inside the opaque op; declares
      ``mutates_args=("recurrent_state_slab",)`` so Inductor respects
      the in-place semantics. Returns only ``core_attn_out`` (final
      state is the slab mutation). The conv update, projections, and
      snapshot writes stay in :class:`GDNBlock`'s Python forward — only
      the FLA kernel call goes through the op. Production path under
      ``ARBI_COMPILE_ON=1``.
    * ``gdn_packed_decode_v2``     — vendored vLLM packed-decode kernel
      wrapped as a slab-mutating custom op. Same pattern as
      ``gdn_attention_v2`` but wraps
      ``fused_recurrent_gated_delta_rule_packed_decode`` (the headline
      decode kernel — single Triton launch fusing gate + sigmoid +
      softplus + recurrence + slab gather/scatter).  Without this op,
      compile-on falls through to the slower ``_forward_decode_fla``
      path because Dynamo refuses to trace the kernel's
      ``triton.knobs`` env-read.  Declares
      ``mutates_args=("recurrent_state_slab", "out")`` — the kernel
      writes both the slab (recurrence) and the output buffer
      (decode-step output).
    * ``gdn_decode_conv_update``   — indexed single-token Triton GDN
      causal-conv update; its snapshot sibling also publishes the MTP
      rollback slab in-launch.
    * ``short_conv_step``          — :func:`causal_conv1d_fn` /
      :func:`causal_conv1d_update` call site in :class:`ShortConvBlock`.
    * ``mtp_verify_greedy``        — argmax-stability verify path
      (pure-tensor: ``main_logits`` argmax + match against
      ``draft_tokens`` + first-mismatch reduction).
    * ``mtp_sample_residual``      — Leviathan-2023 / Chen-2023
      rejection-sampling kernel core. The per-row ``SamplingParams``
      chain (temperature/top_p/top_k/min_p) is applied to
      ``main_logits`` BEFORE the op call by
      :func:`arbi_serve.spec_decode.rejection_sampler.GraphSafeRejectionSampler.sample_batched`;
      the op sees pure-tensor ``main_probs (K+1, B, V)`` /
      ``draft_probs (K, B, V)`` / ``uniforms (K, B)`` /
      ``gumbel_noise (K+1, B, V)``. Inductor sees a stable boundary;
      the SamplingParams Python loop stays outside any compiled region.
    * ``mtp_sample_residual_philox`` — CUDA sibling of
      ``mtp_sample_residual``: same accept math and output contract, but
      the recovery/bonus Gumbel noise is generated INSIDE a single-pass
      Triton kernel from counter-based Philox keyed on the on-device
      step seed (``seed (1,) int64``) instead of a materialized
      ``gumbel_noise`` operand.

Surfaces registered here
------------------------

Attention (``arbi_serve::*``):
    * ``tkv_attention``           — turbo-attn decode + bypass + prefill
                                     (variant captured by ``k_bits`` / ``v_bits``
                                     kwargs; one entry, vLLM-style).
    * ``tkv_bypass_attention``     — raw-bf16 turbo-attn over the fused
                                     paged slab (split-K paged flash-decode
                                     + Turbo prefill); per-layer dispatch.
    * ``mla_attention``            — DeepSeek-style MLA paged attention.
    * ``gdn_attention_v2``         — FLA gated delta-rule (chunk + decode),
                                     slab-mutating; production hot path
                                     under compile-on.
    * ``gdn_packed_decode_v2``     — vendored vLLM packed-decode kernel
                                     (slab-mutating op for compile-on hot path).
    * ``gdn_decode_conv_update`` / ``gdn_decode_conv_update_snapshot`` —
      Triton GDN decode conv update, optionally with in-launch MTP snapshot.
    * ``short_conv_step``          — LFM2 ShortConv mixer.

Spec-decode:
    * ``mtp_verify_greedy``        — argmax-stability verify path.
    * ``mtp_verify_greedy_tree``   — the greedy accept over a TREE-shaped
                                     verify block: a path walk, not a
                                     prefix scan.
    * ``mtp_verify_stochastic_tree`` — the lossless multi-candidate accept
                                     over the same block: one uniform per
                                     depth, inverse-CDF over the sibling
                                     set, residual draw at the frontier.
    * ``mtp_sample_residual``      — Leviathan-2023 residual sampler.
    * ``mtp_sample_residual_philox`` — counter-based (Philox) residual
                                     sampler; CUDA-only.

Already-registered ops we DO NOT redefine here:
    * ``arbi_fp8::triton_fused_mm`` — registered (with fake) in
      ``arbi_serve/weight_quant/fp8/triton_gemm/triton_fused.py`` at module
      import. The fp8 ``triton_gemm`` package is the registration site;
      this module just re-asserts that the op is callable via
      ``torch.ops.arbi_fp8.triton_fused_mm`` for the test surface.
    * ``arbi_serve_awq_marlin::*``  — registered by C++
      (``TORCH_LIBRARY_EXPAND``) in
      ``arbi_serve/weight_quant/awq/marlin/csrc/torch_bindings.cpp``. The
      C++ side ships *no* fake impl — :func:`register_marlin_fakes_if_loaded`
      attaches Python-side fakes idempotently after the JIT load. Called
      from ``arbi_serve.weight_quant.awq.marlin.loader._ensure_loaded`` so
      the fakes appear the first time anything actually loads Marlin
      (CPU smoke / vanilla bf16 runs never trigger this).
    * ``_C::*`` (fp4-Marlin) — the vendored stable-ABI NVFP4 W4A16
      extension (``_C::gptq_marlin_repack`` / ``_C::marlin_gemm``;
      ``arbi_serve/weight_quant/nvfp4/marlin/csrc/bindings.cpp``). Same
      deferred-fake pattern via
      :func:`register_nvfp4_marlin_fakes_if_loaded`, called from
      ``arbi_serve.weight_quant.nvfp4.marlin.loader._ensure_loaded``.

xgrammar logits processor (``apply_token_bitmask_inplace``) is
out-of-scope: it is a third-party C++ op (registered by the xgrammar
package), not an arbi-serve op. It runs *after* the compiled forward
returns logits — Inductor never needs to see through it. It lives
outside any ``@support_torch_compile`` region (the sampler pipeline).
We document it as excluded here rather than register a no-op fake for
an op we don't own.

EXL3 trellis GEMM is similarly excluded: ``exllamav3_ext.BC_LinearEXL3``
is a C++ kernel registered by exllamav3 itself, not arbi-serve. The
``EXL3LinearBase._exl3_forward`` Python path is what arbi-serve owns,
and the kernel call inside it is ``self._inner.forward(...)`` — not a
``torch.ops.*`` op. Wrapping that as an arbi-serve custom op would
require teaching ``exllamav3_ext`` about meta tensors, which is out of
scope here.

Registration discipline
-----------------------

Each ``arbi_serve::*`` op declares only **torch.library-compatible types**
on its signature: Tensors, ``int``, ``float``, ``bool``, ``str``,
``Optional[Tensor]``, and lists of those. Backend-specific Python
objects (``TQRunState``, ``AttnPagedKVMeta``) are broken down into
their constituent tensors at the call site.
Layer-specific state is threaded through a side-channel
(``get_attention_context(layer_name)``).

Ops that are not yet wired register an ``op_func`` real-impl that
raises ``NotImplementedError`` — the production runtime does not call
``torch.ops.arbi_serve.*`` for those. Tracing (Inductor /
``torch.compile``) only needs the ``register_fake`` impl, which is
what Inductor consults to plan allocations through the black box. See
``tests/test_compile_op_fakes.py`` for the trace-time verification.

Package layout
--------------

The single source of truth for the registration guards lives in
:mod:`arbi_serve._custom_ops._registry` (``_LIB``, ``_REGISTERED``, the
per-layer dispatch registry, the CSR thread-local side-channel, and the
Marlin-fakes guard) — every family module imports it, never re-creates
it. The family modules are :mod:`attention`, :mod:`recurrent`
(:mod:`recurrent_gdn` + :mod:`recurrent_ssm`), :mod:`spec_decode`,
:mod:`norm`, and :mod:`marlin`. :func:`register_all`
runs at package import preserving the original registration order.

Register every arbi_serve::* op schema + fake. Idempotent.

Order is load-bearing for the ``_REGISTERED`` de-dupe guard:
attention (tkv, tkv-bypass, mla), then the recurrent family
(GDN v2/v3 + packed-decode + prefill-conv-flat, then mamba2 +
short-conv), then spec-decode (verify-greedy, sample-residual),
then norm (fused-add-rms-norm), then the copy-engine TP reduce pair,
then a best-effort Marlin attempt.

Single source of truth for the arbi-serve custom-op registry.

This module owns the process-global registration state shared by every
family module in the :mod:`arbi_serve._custom_ops` package:

  * ``_LIB``                 — the ``arbi_serve`` torch.library FRAGMENT.
  * ``_REGISTERED``          — name-set guard so re-imports / re-runs of
                               ``register_all`` don't re-register and trip
                               ``RuntimeError: Tried to register an operator
                               with a different schema``.
  * ``_PER_LAYER_DISPATCH``  — per-layer dispatch registry keyed by
                               ``(op_name, layer_idx)``.

These are defined here and imported (never re-created) by the family
modules — a duplicated guard desyncs and re-registration raises "different
schema".

Register a per-layer dispatch target for a custom op.

Backends call this from their ``make_attn_op`` so the op_func
real-impl can find the right per-layer kernel-host object.

``owner`` scopes the key (main model vs external drafter). Registering a
key already owned by a DIFFERENT scope raises — this is the cross-owner
clobber (a drafter op overwriting a main-model layer, or vice versa) that
silently broke the main model's compiled forward. Same-owner overwrite is
allowed (model rebuild / stable-VA pool members re-registering).

Drop every ``(op_name, layer_idx)`` entry whose registered target is
(by identity) one of ``targets``. Returns the number of entries dropped.

A per-kind teardown (:func:`arbi_serve.engine.active.teardown_kind`)
drops its kind's ops from ``eng.attn_ops``, but this process-global
registry — a SEPARATE strong-ref path, keyed by ``(op_name,
layer_idx)`` rather than owned by ``eng`` — still points at the SAME
op objects until the next ``build_active()`` overwrites those keys
(same-owner overwrite is allowed; see :func:`register_layer_dispatch`).
Between teardown and that overwrite, an op parked here alone is enough
to keep it (and anything it holds — e.g. a lazily-built kernel-host
object) reachable, which is exactly the seam
``NamedPoolRegistry.release_empty_pool`` polices: a torn-down pool must
be genuinely empty, not merely unreferenced by ``eng``. Callers should
invoke this for a kind's dropped ops in the SAME teardown call that
drops them from ``eng.attn_ops``, so this registry never outlives the
kind's actual state. Identity (``is``), not equality — dispatch targets
are stateful objects, not values.

Shallow copy of the current per-layer dispatch registry.

Stable-VA residency snapshots each model's dispatch right after that
model builds (before the next pool member's build overwrites the
shared ``(op_name, layer_idx)`` keys with its block objects), then
:func:`install_layer_dispatch` re-installs the woken model's snapshot
on every switch. Without this the eager custom-op path
(:func:`get_layer_dispatch`) dispatches a woken model's attention to
the last-built member's block — whose per-model buffers are parked /
at a different VA — an illegal-memory-access on the first eager
multi-sequence prefill (the captured/compiled path is unaffected: it
routes per-call meta via ``publish_call_meta``, not this registry).

Run TKV paged attention for layer ``layer_idx``, writing into
``output`` and ``kv_cache`` in place.

Resolves the per-layer :class:`TkvAttnOp` and dispatches to its
kernel-call helper, which composes the compress/bypass/decode paths
for the given ``(k_bits, v_bits)`` variant.

Run raw-bf16 turbo-attn attention for layer ``layer_idx`` over the
fused paged KV slab, writing into ``output`` and ``kv_cache`` in place.

Scatters new K/V into the fused ``[K|V]`` bf16 slot, then attends
*in place*: decode (one query per row) via the split-K paged
flash-decode (:class:`BypassDecodeLoader`, q pre-scaled), prefill
via the Turbo prefill kernel (:class:`BypassLoader`). The per-layer
:class:`TkvBypassAttnOp` holds the per-step CSR page-metadata
triplet (published via ``publish_call_meta``); per-layer state is
routed via :func:`get_layer_dispatch`.

Run MLA attention for layer ``layer_idx``,
writing into ``output`` and ``kv_cache`` in place.

Dispatches through the per-layer :class:`MLAAttnOp`, which owns the
codec and ``kv_b_proj`` reference for the decompress path.

Every PER-STEP value is an argument: the cache slab, the four
scheduling tensors, and the host ints that route the step
(``max_query_len`` / ``mtp_block_m``) or size its launches
(``max_seq_len`` / ``total_kv``). The op resolves only per-LAYER
state (codec, ``kv_b_proj`` halves) through
:func:`get_layer_dispatch` — the same discipline
``gdn_attention_v2`` follows, and what makes the op traceable:
a Dynamo trace of this call cannot silently freeze a per-step
Python attribute, and a captured graph replays the values its
buffers hold.

Register Python-side fakes for the Marlin C++ ops, if loaded.

Idempotent — safe to call from
:func:`arbi_serve.weight_quant.awq.marlin.loader._ensure_loaded`
and from anywhere else after JIT load. Returns True if the fakes
are now in place; False if the extension hasn't been loaded yet
(caller can re-try later).

The C++ side defines the schemas via ``TORCH_LIBRARY_EXPAND``; we
only need the Python ``register_fake`` impls here. Marlin's output
layout is fixed by the kernel and not arbi-serve-specific.

Register Python-side fakes for the vendored fp4-Marlin (NVFP4
W4A16) C++ ops, if loaded.

Same idempotent contract as :func:`register_marlin_fakes_if_loaded`,
but for the stable-ABI extension that registers into the ``_C``
namespace (``_C::gptq_marlin_repack`` / ``_C::marlin_gemm``; see
:mod:`arbi_serve.weight_quant.nvfp4.marlin.loader`). The W4A16 MLP /
lm_head projections call ``marlin_gemm`` inside the compiled decode
graph, so Inductor needs these meta impls to plan through it.

Returns True if the fakes are now in place; False if the extension
isn't loaded yet (caller can re-try later).

Fuse the pre-norm residual add and RMSNorm into one Triton
launch.

Updates ``residual`` in place to ``residual + x`` and returns the
RMSNorm of the new residual cast back to ``x``'s dtype.

Recurrent / hybrid custom ops (GDN + Mamba2 + ShortConv).

The family is split into :mod:`recurrent_gdn` (GDN attention v1/v2/v3,
packed-decode, prefill-conv) and :mod:`recurrent_ssm` (Mamba2 selective
scan + ShortConv step). This module re-aggregates them and exposes
:func:`register_recurrent_ops`.

GDN (Gated Delta Net) recurrent custom ops.

The GDN attention variants (v2/v3), the vendored packed-decode
kernel wrapper, decode conv update, and prefill rolling-buffer conv ops. See
:mod:`arbi_serve._custom_ops` for the package-level docstring.

The registration closures are split across sibling modules: each
``torch.library`` op is defined and registered in exactly one place,
imported once here, and fired once by :func:`register_gdn_ops`:

    * :mod:`recurrent_gdn_attention`     — ``gdn_attention`` v2/v3.
    * :mod:`recurrent_gdn_verify`        — ``gdn_verify_replay_save`` /
      ``gdn_verify_chained`` / ``gdn_decode_fused``.
    * :mod:`recurrent_gdn_packed_decode` — ``gdn_packed_decode_v2`` (+
      snapshot sibling).
    * :mod:`recurrent_gdn_decode_conv`   — ``gdn_decode_conv_update``.
    * :mod:`recurrent_gdn_prefill_conv`  — ``gdn_prefill_conv_flat``
      (+ the ``causal_conv1d`` opt-in gate).

``_gdn_conv_causal_conv1d_enabled`` is re-exported here so it stays
importable from ``arbi_serve._custom_ops.recurrent_gdn`` (a test surface).

GDN gated-delta-rule attention custom ops (v2 / v3).

The two slab-mutating FLA delta-rule call sites: ``gdn_attention_v2``
(slab gather/scatter inside the opaque op) and ``gdn_attention_v3``
(v2 plus the MTP rollback snapshot in the same launch). Registration is
fired by :func:`arbi_serve._custom_ops.recurrent_gdn.register_gdn_ops`.

Register ``arbi_serve::gdn_attention_v2`` — slab-mutating variant.

An op that declares ``mutates_args=()`` cannot express the slab
gather (``state_view.recurrent_state.index_select(0,
state_indices)``) before the kernel call and the slab scatter
(``state_view.recurrent_state.index_copy_(0, state_indices,
final_state)``) after it — both would sit in plain Python that
Inductor traces. Inductor reasons functionally about the
surrounding chain and freely re-orders / re-inplaces around the
"redundant" ``index_copy_``, producing wrong decode logits under
``ARBI_COMPILE_ON=1``. The whole slab + state-indices are instead
passed into the op, which declares
``mutates_args={"recurrent_state_slab"}`` so Inductor respects the
in-place semantics.

Slab layout. The arbi-serve recurrent slab is allocated as
``(N, HV, V, K)`` (V outer, K inner) — see
``arbi_serve/cache/recurrent_pool.py``. FLA's
``chunk_gated_delta_rule`` /
``fused_recurrent_gated_delta_rule`` consume initial / final state
in ``(B, HV, K, V)`` (K outer). The op handles the K↔V transpose
at the slab boundary inside the opaque region — callers pass the
raw slab and the op gathers, transposes, runs the kernel, and
scatters the K↔V-transposed final state back.

The snapshot write (``snap_recurrent_state[t][slab_row]``) for
MTP rollback is intentionally left outside the op — it reads from
the slab after the op call (``slab.index_select(0, idx)``) so the
snapshot runs as a separate kernel after the op's mutation
completes; Inductor sees the data dep and orders it correctly.

Signature:

  Inputs:
    q, k, v, beta, g       — FLA delta-rule inputs.
    recurrent_state_slab   — the WHOLE slab, ``(N, HV, V, K)``.
                             Mutated in place at rows
                             ``state_indices``.
    state_indices          — int64 row indices into the slab,
                             shape ``(B,)``.
    cu_seqlens             — varlen offsets for prefill mode==0;
                             unused (placeholder ok) for mode==1.
    mode                   — 0 = chunk_prefill, 1 = fused_recurrent.
    layer_idx              — per-layer dispatch key.
    use_qk_l2norm_in_kernel — passthrough to the FLA kernel.

  Output:
    core_attn_out — same shape/dtype as ``v`` (the FLA kernel's
                    attention output). The "final state" is not
                    returned — it has been written into the slab.

Register ``arbi_serve::gdn_attention_v3`` — slab-mutating variant
that also writes the MTP rollback snapshot in the same launch.

The MTP verify
path (:meth:`GDNBlock._forward_verify_fla` ARBI_GDN_FUSED_-
RECURRENT=0 branch) writes a per-t-step slab snapshot for rollback
(``snap_rec_full[t][slab_row] = state after committing tokens
[0..t]``). Since ``gdn_attention_v2`` only mutates the slab and
returns ``core_attn_out``, the snapshot has to round-trip through a
``slab.index_select(0, state_indices)`` read after the op — extra
``index_select`` per t-step (T launches per layer × ~18 GDN layers
per verify step).

v3 fuses the snapshot write into the op real-impl: the kernel's
final state is written directly to both the slab and the snapshot
tensor in the same dispatch.  The schema declares
``mutates_args=("recurrent_state_slab", "snapshot_out")`` so
Inductor respects both in-place mutations.

Slab vs snapshot layout. Both are ``(N, HV, V, K)`` (V outer, K
inner) — see ``cache/recurrent_pool.py`` and ``RecurrentStatePool.
attach_mtp_snapshot_buffers``.  ``snapshot_out`` is ``snap_rec_-
full[t]`` — a per-step slice of the snapshot buffer, shape
``(N, HV, V, K)``.  The op ``index_copy_``'s into rows
``state_indices`` of ``snapshot_out`` (same indices as the slab
write).

MTP-only.  The decode and prefill hot paths still use v2 (no
snapshot to write — decode writes ``snap_rec_full[0]`` once after
the op call, well outside the per-t loop).  Only the verify per-t
loop benefits from v3.

Signature: extends v2 with ``snapshot_out`` as an additional
mutate-arg.  The ``snapshot_out_dtype_cast`` is implicit — the op
reads ``snapshot_out.dtype`` and casts ``final_state_vk`` to it
before scattering.

Run the GDN delta-rule kernel with slab gather/scatter inside
the opaque op (production hot path under ``ARBI_COMPILE_ON=1``).

Mutates ``recurrent_state_slab`` in place at rows ``state_indices``
and returns only ``core_attn_out`` (the final state is the slab
mutation).

Run the GDN delta-rule kernel and also write the MTP rollback
snapshot in the same opaque op (MTP verify per-t-step path).

Mutates both ``recurrent_state_slab`` and ``snapshot_out`` in place
at rows ``state_indices``; returns only ``core_attn_out``.

GDN vendored packed-decode custom ops.

The vendored vLLM packed-decode kernel wrapped as slab-mutating custom
ops — ``gdn_packed_decode_v2`` (no snapshot) and its sibling
``gdn_packed_decode_v2_snapshot`` (writes the MTP rollback snapshot in
the same launch). Registration is fired by
:func:`arbi_serve._custom_ops.recurrent_gdn.register_gdn_ops`.

Register the vendored vLLM packed-decode kernel as slab-mutating
custom ops — ``arbi_serve::gdn_packed_decode_v2`` (no snapshot) and
its sibling ``arbi_serve::gdn_packed_decode_v2_snapshot`` (writes the
MTP rollback snapshot in the same launch).

Two ops, not one: ``snapshot_out`` is optional on the seed-decode
path (``None`` when MTP is off), but a declared mutated arg that is
also optional+trailing is a torch footgun — when a caller omits it,
Dynamo's ADInplaceOrView version-bump indexes ``args[idx]`` past the
end of the (shorter) boxed positional call and raises ``IndexError:
tuple index out of range`` (only the eager CUDA path dodged it; the
whole engine under compile/cudagraphs did not). Splitting into two
ops gives each a fully static ``mutates_args`` whose every entry
always has a positional slot, so both trace cleanly. The call site
routes on whether a snapshot buffer is present.

``arbi_serve.kernels.fused_recurrent_packed_decode.fused_recurrent_-
gated_delta_rule_packed_decode`` is the headline GDN decode kernel
(single Triton launch fusing gate + sigmoid + softplus + the
recurrence + slab gather/scatter — see the kernel file's preamble
for the win breakdown).  Production hot-path eager runs it directly,
but under ``torch.compile`` Dynamo refuses to trace into the
``triton.knobs`` env-read inside the JIT shell, falling through to
the slower ``_forward_decode_fla`` path under trace unless the call
site gates on ``not torch.compiler.is_compiling()``.

Wrapping the kernel call in a ``torch.library.custom_op`` gives
Dynamo a single opaque boundary — the kernel's internal
``triton.knobs`` read happens inside the real-impl, never in the
traced path.  The op declares
``mutates_args=("recurrent_state_slab", "out")`` because the
kernel mutates both the slab (in-place state update at row
``ssm_state_indices[i]``) and the pre-allocated ``out`` tensor (the
decode-step output written via ``tl.store``).  Inductor respects the
in-place semantics and won't reorder around either tensor.

Slab layout. ``recurrent_state_slab`` is ``(N, HV, V, K)`` (V outer,
K inner — see ``arbi_serve/cache/recurrent_pool.py`` line 398 and
the kernel's ``transpose_state=True`` branch at fused_recurrent_-
packed_decode.py line 164-168).  The op always passes
``transpose_state=True`` (production layout); the eager-only
``transpose_state=False`` testing path stays direct (no compile
coverage needed).

Signature:

  Inputs:
    mixed_qkv              — (B, conv_dim_local), post-conv [Q‖K‖V]
                             packed.
    a, b                   — (B, HV) raw projections.
    A_log, dt_bias         — (HV,) per-head decay + bias.
    recurrent_state_slab   — the WHOLE slab, (N, HV, V, K).
                             Mutated in place at rows
                             ``ssm_state_indices``.
    out                    — (B, 1, HV, V) preallocated output
                             buffer.  Mutated in place.
    ssm_state_indices      — (B,) int64 row indices into the slab.
    scale                  — Python float, ``1/sqrt(K)``.
    layer_idx              — per-layer dispatch key (matches the
                             ``gdn_attention_v2`` registry).
    use_qk_l2norm_in_kernel — passthrough to the FLA kernel.

  Output:
    Returns ``out`` (same tensor, after mutation).  Returning the
    mutated tensor matches ``arbi_serve::tkv_attention``'s
    ``mutates_args=("output", ...)`` pattern — the caller chains on
    the returned alias to keep the data-dep edge in the FX graph
    explicit.

No K↔V transpose at the boundary — unlike ``gdn_attention_v2``,
this op consumes the slab in its native ``(N, HV, V, K)`` layout
directly (the kernel's ``transpose_state=True`` math handles V-
outer addressing).

Run the vendored vLLM GDN packed-decode kernel (no snapshot).

Mutates ``recurrent_state_slab`` (at rows ``ssm_state_indices``)
and the pre-allocated ``out`` buffer in place; returns a clone of
``out`` (custom-op aliasing rules forbid returning the input).

This is the no-snapshot decode path. ``mutates_args`` is a fully
static 2-tuple so Dynamo's version-bump (ADInplaceOrView) indexes
only positions that always exist in the boxed call. The MTP
rollback-snapshot variant is the sibling op
:func:`gdn_packed_decode_v2_snapshot` — keeping the optional
snapshot out of this schema avoids the torch footgun where an
optional trailing arg that is also a declared mutated arg breaks
tracing when the caller omits it (``args[idx]`` out of range).

Run the GDN packed-decode kernel with the in-launch MTP snapshot.

Identical to :func:`gdn_packed_decode_v2` except the kernel also
writes the committed final state into ``snapshot_out`` (the MTP
rollback snapshot slice ``snap_rec_full[0]``) in the same launch —
replacing the host-side ``slab.index_select`` + ``snap.index_copy_``
round-trip (a per-GDN-layer fp32 slab-sized copy). This preserves
that behavior while keeping the schema clean: ``snapshot_out`` is
a required ``Tensor`` here, so every declared mutated arg always
has a positional slot in the boxed call (unlike an optional
trailing mutated arg, which breaks Dynamo's version-bump when
omitted). Mutates ``recurrent_state_slab``, ``out``, and
``snapshot_out`` in place; returns a clone of ``out``.

GDN prefill rolling-buffer conv custom ops.

The GDN prefill depthwise-conv loop wrapped as the opaque custom op
``gdn_prefill_conv_flat`` (flat ``(T_total, C)`` input + ``cu_seqlens``).
Also hosts the Dao-AILab ``causal_conv1d_fn`` opt-in gate and the varlen
Triton availability probe. Registration is fired by
:func:`arbi_serve._custom_ops.recurrent_gdn.register_gdn_ops`.

Mark ``fn`` constant for Dynamo without importing the compiler frontend.

``torch._dynamo.assume_constant_result`` sets ``_dynamo_marked_constant`` on
the function and returns it; Dynamo reads that attribute at trace time.
Reaching the decorator pulls ``torch._dynamo`` — the whole compiler
frontend, sympy included — into every process that imports this module,
architecture-independently and whether or not anything compiles. Setting the
marker directly is the same contract at import cost zero.
``tests/test_boot_import_cost.py`` pins it to torch's decorator.

Opt-in flag to route the GDN prefill conv through Dao-AILab's tuned
``causal_conv1d_fn`` (batched/uniform branch) instead of the in-tree
varlen Triton kernel — see
``arbi_serve/kernels/gdn_prefill_conv_causal_conv1d.py``.

Default OFF: ``ARBI_GDN_CONV_CAUSAL_CONV1D`` unset → the canonical
Triton path. Set to ``1``/``true``/``yes``/``on`` to A/B the
``causal_conv1d_fn`` path on GPU. Read per-call (not cached) so the
flag can be flipped between server runs without re-import.

Gather this launch's conv window into the layer's staging slot.

The conv state after ``t`` tokens is the last ``K`` raw inputs before
``t`` (``gdn_prefill_conv_flat`` defines ``final_buf`` as exactly that
window at the row's end), so it is a slice of ``x_conv`` and no second
conv runs. Which rows: a device index table the metadata builder wrote
for the step (row 0, ``K`` times, when nothing is armed), so the gather
is the same kernels on every launch and enters no trace or capture as a
per-step value. Copied out NOW: ``x_conv`` is a compiled graph's
intermediate and its storage is reusable the moment this op returns.

Register ``arbi_serve::gdn_prefill_conv_flat`` — flat-input variant
of the GDN prefill rolling-buffer conv.

Accepts ``x_conv`` in the flat ``(T_total, C)`` layout the caller
produces via the ``cat([Q, K, V], dim=-1)`` projection, plus
``cu_seqlens`` describing per-row token boundaries. The (B, ...) /
per-row reshape and dispatch happen inside the op, behind the
``custom_op`` boundary, so Dynamo never sees the symbolic
``unflatten(x_conv, 0, (B, -1))`` / ``view(B, T_per_row, C)`` call
that fails Dynamo's FakeTensor numel check at B>=2 under layer
compile.

This op exists to sidestep two failure modes of doing the reshape
in traced Python:

(1) Compile-time: under ``@support_torch_compile`` on the GDN
decoder layer with ``fullgraph=True``, ``T_total`` (s59) and ``B``
(s13) come in as independent SymInts and Dynamo's symbolic-shape
solver cannot discharge ``Eq(s59, s13 * (s59 // s13))`` (sympy
FloorDiv is unprovable without the modular relation). The
trace dies with ``unflatten: Provided sizes [s13, -1] don't
multiply up to size of dim 0 (s59)`` — the engine step fails and
every request at B>=2 returns 1-2 tokens then goes silent.

(2) Runtime: a uniform-T_per_row if-branch under
``is_compiling()`` also assumes ``T_total % B == 0`` at runtime.
Under chunked prefill at B>=2 that is not always true — when a
new request is admitted into a slate where existing rows are
mid-prefill, per-row ``want = min(num_remaining_prompt,
chunk_prefill, budget)`` differs and ``view(B, T_total // B, C)``
crashes with a numel mismatch.

Hiding the reshape and the uniform/non-uniform dispatch inside the
op makes Dynamo see ``(T_total, C) → (T_total, C)`` only; no
B/T_total relationship is materialized in the FX graph, and the
runtime path picks the right impl based on the actual cu_seqlens.

Signature:

  Inputs:
    x_conv     : (T_total, C) — flat token-major conv input.
    prior_conv : (B, C, K)    — prior conv state, gathered + masked.
    w          : (C, K)       — depthwise conv weights.
    cu_seqlens : (B+1,) int32 — per-row token boundaries
                                (``[0, T_0, T_0+T_1, ..., T_total]``).

  Outputs:
    conv_out   : (T_total, C) — silu-activated conv output (flat).
    final_buf  : (B, C, K)    — final rolling buffer state.

No mutation: caller-owned slab writes happen outside the op.
Numerics: bit-identical to the in-line path on the
uniform branch (same kernel: F.conv1d + silu); the non-uniform
branch matches :meth:`GDNBlock._conv_prefill_pertoken` token-for-
token.

Run the GDN prefill rolling-buffer conv on FLAT ``(T_total, C)``
input, with the per-row reshape and uniform/non-uniform dispatch
hidden inside the opaque op.

Returns ``(conv_out, final_buf)``; ``cu_seqlens`` gives per-row
token boundaries. No mutation — the caller scatters ``final_buf``
into the slab after the op.

GDN MTP verify + fused-decode custom ops.

Capture-safe single-launch ops for the MTP verify recurrence and the
ARBI_ACCEPT_INVARIANT decode path: ``gdn_verify_replay_save``
(masked-replay chunked verify), ``gdn_verify_chained`` (bit-exact
conv+recurrence per token), and ``gdn_decode_fused`` (single-token fused
decode). Registration is fired by
:func:`arbi_serve._custom_ops.recurrent_gdn.register_gdn_ops`.

Register ``arbi_serve::gdn_tree_packed_scan`` — the TREE verify scan.

Same reason the chain's replay-save launch needs an op: the decode /
verify cudagraph capture COMPILES the decoder layer, and the fused
recurrent kernel cannot be traced (its FLA helpers carry
``@torch.compiler.disable``, which raises rather than graph-breaking
under ``fullgraph=True``). Without this the verify bucket at the
tree's width fails to capture and the whole verify pass runs eager —
correct output, and a step cost that cannot be compared with a
captured chain's.

``recurrent_state_slab`` is READ here and not written: the scan commits
no final state at all (``commit_final_state=False``), so a tree's packed
rows — all naming the one slab index — cannot race each other's ``h0``
loads. The op mutates nothing.
Output: the packed per-token mixer output, ``(L*D, HV, V)``.

Register ``arbi_serve::gdn_verify_replay_save`` — the masked-replay
(``ARBI_GDN_MTP_INGRAPH_ROLLBACK → replay``) verify recurrence as a capture-safe
custom op.

The fused kernel's FLA helpers carry ``@torch.compiler.disable`` and
its ``b_h`` fp32 accumulator trips Inductor if traced; the decode/
verify cudagraph capture compiles the model forward, so the replay-mode
verify must present one opaque launch per layer. No per-position
ladder and no base frame — the raw per-token inputs land in the
``replay_*`` buffers (in-kernel) and no state is committed, so the slab
still holds the h0 the launch read. The accept path's masked replay then
recomputes the accepted-prefix state bit-for-bit from that slab h0 and
these tensors with the same compiled binary.

Inputs are the chunked verify tensors with the replay input buffers
(mutated) carrying the rollback state.
Output: ``core_attn_out`` — (T_total, HV, V).

Register ``arbi_serve::gdn_verify_chained`` — ARBI_ACCEPT_INVARIANT
MTP-verify recurrence wrapped as a capture-safe custom op.

Background. Under ``ARBI_ACCEPT_INVARIANT=1`` the verify mixer must be
bit-identical to T sequential c=1 decode steps — conv included. On sm_89
the default verify path runs a hand-rolled PyTorch unfold conv while decode
runs the Dao-AILab ``causal_conv1d_update`` Triton kernel, and the two
diverge enough to break bit-parity at post_mixer.
:meth:`GDNBlock._dispatch_verify_invariant_chained` fixes it by
re-running the exact decode conv + recurrent kernels per token.

Why a custom op. The chained recurrence calls
``fused_sigmoid_gating_delta_rule_update`` directly. Its ``b_h`` fp32
state accumulator trips an ``InductorError: b_h fp32↔fp64`` if Inductor
traces it (the same reason ``_forward_decode_fla`` gates its fused path
off under a piecewise-cudagraph window). Wrapping the whole per-token
conv+recurrence+snapshot loop in this opaque op keeps Inductor out of the
kernel, so ``ARBI_ACCEPT_INVARIANT=1`` cudagraph-captures on sm_89 with no
eager-GDN fallback. Same opaque-boundary pattern as the verify ops.

Signature (all post-projection, flat over T_total = B*(K+1) tokens):
  x_conv_flat          — (T_total, conv_dim_local) [Q‖K‖V] conv input.
  b_raw, a_raw         — (T_total, HV) raw gate projections.
  conv_state_slab      — whole conv slab (N, C, K). Mutated.
  recurrent_state_slab — whole recurrent slab (N, HV, V, K). Mutated.
  state_idx_long       — int64 (B,) slab-row indices.
  state_idx_int32      — int32 (B,) row indices (causal_conv1d).
  layer_idx            — per-layer dispatch key.
  b, t                 — uniform-K shape scalars (B, T=K+1).

Output: core_attn_out — (T_total, HV, V) per-token mixer output.

No per-position snapshot. Only the final
post-token-(T-1) state commits to the slabs; partial-accept rollback
recomputes the accepted prefix host-side from the 1× base frame + the
per-token raw inputs the verify forward captured (recompute mode, forced
under the invariant flag by ``gdn_mtp_rollback_mode``).

Register ``arbi_serve::gdn_decode_fused`` — the single-token fused
``fused_sigmoid_gating_delta_rule_update`` decode recurrence as a
capture-safe custom op (``ARBI_ACCEPT_INVARIANT``).

Under the flag, decode must run the same fused kernel as the chained
verify recurrence (bit-parity) even inside the compile / piecewise-
cudagraph window. Calling the kernel directly under Inductor trips
``InductorError: Loop-carried variable b_h ... fp32 ... re-assigned to
fp64`` (Inductor infers the kernel's beta/threshold/scale scalars as
fp64). This opaque wrapper keeps Inductor out of the kernel so decode
captures on sm_89 while staying byte-identical to eager decode and to
``gdn_verify_chained``. Flag OFF never routes here (the decode fused gate
excludes the piecewise window), so default decode is unchanged.

Signature (post-conv, one token per row, cu_seqlens=[0,1,..,B]):
  q, k, v          — (1, B, H/HV, Dk/Dv).
  b, a             — (B, HV) raw gate projections.
  recurrent_state_slab — whole slab (N, HV, V, K). Mutated.
  state_idx_long   — int64 (B,) slab-row indices.
  snapshot0        — (N, HV, V, K) MTP rollback slot-0. Mutated.
  cu_seqlens       — (B+1,) int32.
  layer_idx        — per-layer dispatch key.
Output: core_attn_out — (B, HV, V).

Run the short-conv causal-conv1d kernel for layer ``layer_idx``.

Resolves the per-layer :class:`ShortConvBlock` and dispatches only
the ``causal_conv1d_*`` kernel (``causal_conv1d_fn`` for prefill
``mode==0``, ``causal_conv1d_update`` for decode ``mode==1``); the
projections and gate split stay in the block. ``conv_state`` is
mutated in place. Returns the conv output (same shape as ``Bx``).

``cu_seqlens`` / ``state_indices`` / ``has_initial_state`` are
threaded as real tensor args (mirroring ``gdn_attention_v2``'s
``state_indices`` / ``cu_seqlens``) — not a ``self._call_meta``
side-channel. Under ``torch.compile`` / piecewise capture the
Python ``forward`` that would set a side-channel attribute is
traced once (at capture time, B-padded shapes) and its
non-tensor side effects are not re-run on replay; a stashed meta
therefore goes stale (e.g. a B=2 capture's ``cu=[0,512,1024]``
leaking into a live B=1 15-token prefill, launching the conv
kernel with a zero-length chunk → ``cudaErrorInvalidConfiguration``).
Passing these as op args makes the live per-step values flow
through the graph correctly.

Pure-tensor greedy MTP verify.

Computes the per-slot main-model argmax + first-mismatch reduction
against the drafter picks. Identical math to
:func:`arbi_serve.spec_decode.verify._verify_greedy`; the difference
is that this returns batched tensors instead of building per-row
:class:`VerifyResult` objects (the per-row commit happens after the
op call when the caller pulls the small tensors to host).

Greedy TREE verify: walk the accepted root-to-leaf path.

A chain can identify its accepted prefix by position: slot ``k``
accepts iff ``argmax[k] == draft[k]``, so ``committed`` IS the
per-slot argmax. A tree cannot. Its nodes are siblings as often as
successors, so the token a node proposes is checked against the
argmax at its PARENT's row, and the accepted set is a path through
the tree rather than a prefix of a list.

Node ``i`` is accepted iff its own proposal matched AND its parent
was accepted. At most one node per depth can be accepted — siblings
share a parent, the target's argmax there is one token, and a
drafter's top-W picks are distinct — so the accepted set is a path
and its size is just the count.

``depth`` iterations, all gathers: no host sync, so the async verify
path can keep deferring its host pull to the drain.

Returns ``(accepted_count (B,), committed (depth+1, B), path_rows
(depth+1, B))``. ``committed`` matches
:func:`_mtp_verify_greedy_kernel`'s protocol: ``committed[:n]`` is
the accepted path's tokens and ``committed[n]`` is the recovery /
bonus token — the target's own argmax at the last accepted row, so a
step always commits at least one token and can never stall.

``path_rows`` is WHERE that path went, indexed identically:
``path_rows[d]`` is the verify-block row the walk sits on after ``d``
accepts, so ``path_rows[0]`` is 0 (the committed token's row) and
``path_rows[n]`` is the deepest accepted node's row. A chain does not
need this — its accepted node at depth ``d`` IS row ``d + 1``, so
every consumer reconstructs the row from the count. A tree's is
scattered, and three consumers need it: the bonus-hidden gather (the
drafter must seed from the hidden at the ACCEPTED node's flat slot,
not at slot ``n``), the frontier repair that shares that index, and
the KV compaction that moves the accepted rows down. Returning it
from the walk that already computed it is the only place it is
knowable without re-deriving the accept.

Pure-tensor Leviathan-2023 / Chen-2023 rejection-sampling kernel.

All randomness is provided externally:

  * ``uniforms`` drives the ``u * q <= p`` accept test per
    (slot, row);
  * ``gumbel_noise`` drives the recovery + bonus draws via
    ``argmax(log p + g)`` (Gumbel-max trick).

The op is graph-capture-safe by construction (no ambient
``torch.Generator`` access) and bit-deterministic-on-input.

``draft_probs=None`` is the index-form point-mass proposal (the
default greedy-draft path): ``q[k, b, v] = 1`` iff ``v ==
draft_tokens[k, b]``. The math is bit-identical to passing the dense
one-hot: the accept test's ``u * q_at_x`` has ``q_at_x == 1.0``
exactly (multiplication by 1.0 is exact in IEEE-754), and the
residual ``(p - q)+`` equals ``p`` with the drafted token zeroed —
``p[v] - 0.0 == p[v]`` bit-for-bit off the mass, and
``clamp_min(p[x] - 1.0, 0) == 0.0`` at it (softmax probs are
``<= 1.0``). Only the dense materialization is skipped.

Returns ``(accepted_count (B,) int64, accepted_tokens (K+1, B) int64)``
matching :func:`_mtp_verify_greedy_kernel`'s contract:
  * ``accepted_count[b]`` is the first-mismatch slot index in
    ``[0, K]``;
  * ``accepted_tokens[k, b]`` for ``k in [0, K)`` is the recovery
    sample from the residual ``(main_probs[k] - draft_probs[k])+``;
  * ``accepted_tokens[K, b]`` is the bonus sample from
    ``main_probs[K]``.

Counter-based (Philox) Leviathan-2023 rejection-sampling kernel.

The accept test is identical math to
:func:`_mtp_sample_residual_kernel` (``u * q <= p`` over the supplied
``uniforms``), so accept decisions are bit-identical to the
materialized-noise path for the same uniforms. The recovery + bonus
draws replace the ``(K+1, B, V)`` Gumbel-noise tensor + full-vocab
argmax passes with one single-pass Triton kernel whose noise is
generated in-kernel from Philox keyed on ``(seed, row, vocab_index)``
(:func:`arbi_serve.sampler.gumbel_argmax_triton.residual_bonus_gumbel_from_seed`)
— an exact draw from the same residual ``(p - q)+ / Z`` and bonus
``p_target[K]`` distributions. ``seed`` is read inside the kernel, so
the op body performs no host readback and no full-vocab writes.

``draft_probs=None`` is the index-form point-mass proposal (see
:func:`_mtp_sample_residual_kernel` for the bit-identity argument):
the accept test reduces to ``u <= p_at_x`` and the Triton kernel
computes ``q`` on the fly as ``offs == draft_tokens[k, b]`` — the
same 0.0 / 1.0 values a dense one-hot load would produce, with zero
``(K, B, V)`` q traffic.

Same output contract as :func:`_mtp_sample_residual_kernel`.

Fused-softmax sibling of :func:`_mtp_sample_residual_kernel`.

``main_scaled`` is the temperature-scaled (and, if the slate was
filtered, top_k/top_p-masked-to-``-inf``) verify logits — the softmax
is applied here, inside the op, so the caller never materializes the
dense ``(K+1, B, V)`` ``p_target`` probs tensor with a separate
full-vocab ``softmax`` launch. ``softmax(main_scaled)`` is
byte-identical to the ``p_target`` the materialized path builds (same
``torch.softmax`` over the same scaled/masked logits), so the accept /
recovery / bonus outputs are bit-identical to
:func:`_mtp_sample_residual_kernel` for the same ``uniforms`` /
``gumbel_noise``. This is the CPU / materialized-noise reference for
the fused path (the CUDA path is
:func:`_mtp_sample_residual_philox_from_logits_kernel`).

Fused single-pass counter-based (Philox) verify tail — softmax in-kernel.

A single fused rejection-sample: it consumes
the temperature-scaled verify logits directly and never materializes
the dense ``(K+1, B, V)`` ``p_target`` probs nor runs a separate
full-vocab ``softmax`` launch. A single ``logsumexp`` over the vocab
gives the ``(K+1, B)`` row normalizer, which drives both:

  * the accept test — ``p_target[x] = exp(logit_x - lse)`` per
    (slot, row); for the point-mass proposal (``draft_probs=None``)
    ``u * q`` collapses to ``u`` (``q_at_x == 1`` exactly), so accept
    is ``u <= exp(logit_x - lse)``;
  * the recovery/bonus draws — the single-pass Triton kernel
    (:func:`~arbi_serve.sampler.gumbel_argmax_triton.residual_bonus_gumbel_from_seed`
    with ``in_log_space=True``) recomputes ``p = exp(logit - lse)``
    per lane and generates the Gumbel noise in-kernel from Philox
    keyed on ``(seed, row, vocab_index)`` — no materialized
    ``(K+1, B, V)`` noise, byte-identical across SPMD ranks.

Accept decisions are distributionally identical to the materialized
path; ``exp(logit_x - lse)`` differs from a materialized
``softmax(...)[x]`` by at most a float32 ULP (both are the exact
``p_target[x]``), so for a fixed seed an accept can only flip at an
exact-tie boundary. Same output contract as
:func:`_mtp_sample_residual_kernel`.

``lse=None`` computes the normalizer here (``torch.logsumexp`` — the
no-filter slate's path, byte-identical to the pre-``lse`` op). A
filtered slate passes the ``(K+1, B)`` normalizer the top_k/top_p
mask kernel already produced in its own launch
(:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
with ``lse_out``), so masking costs the fused path zero extra
launches. The supplied ``lse`` must be the logsumexp of exactly
``main_scaled`` (the masked row) — same drift class either way (fp32
reduction order only).

Greedy TREE verify: walk the accepted root-to-leaf path.

The tree analogue of :func:`mtp_verify_greedy`. Its ``committed``
is a real gather rather than the per-slot argmax — a tree's
accepted token is the argmax at its PARENT's row, so the chain's
``committed IS main_argmax`` identity does not hold. The third
output, ``path_rows``, is the verify-block row the walk stands on
after each accept: a chain's is ``d + 1`` by construction, a
tree's is not, and the drafter seed / frontier repair / KV
compaction all address by row rather than by count.

Stochastic MTP verify via Leviathan-2023 rejection sampling.

Pure tensor function over already-prepared probs (SamplingParams
applied by the caller). ``draft_probs=None`` encodes the
index-form point-mass proposal (one-hot on ``draft_tokens``),
bit-identical to the densified one-hot. Returns
``(accepted_count, accepted_tokens)``.

Stochastic MTP verify with in-kernel counter-based Gumbel noise.

CUDA-only sibling of ``mtp_sample_residual``: same accept math and
output contract, but the recovery/bonus randomness is Philox
generated inside a single-pass Triton kernel keyed on ``seed``
instead of a materialized ``(K+1, B, V)`` noise operand.
``draft_probs=None`` encodes the index-form point-mass proposal
(one-hot on ``draft_tokens``), bit-identical to the densified
one-hot — the kernel synthesizes ``q`` from the token index.

Fused-softmax MTP verify: softmax the scaled logits in-op.

Byte-identical outputs to ``mtp_sample_residual`` fed the
materialized ``softmax(main_scaled)`` — the fused path just never
builds that dense ``(K+1, B, V)`` probs tensor outside the op. The
materialized-noise reference / CPU + Philox-off path for
``ARBI_MTP_FUSED_REJECTION``. ``draft_probs=None`` encodes the
index-form point-mass proposal. Returns ``(accepted_count,
accepted_tokens)``.

Fused single-pass Philox MTP verify tail — softmax in-kernel.

CUDA-only. A single fused rejection-sample: a
single ``logsumexp`` normalizer drives the accept test and the
in-kernel-softmax residual/bonus Triton kernel over the scaled
logits — no dense ``(K+1, B, V)`` ``p_target`` probs, no separate
softmax launch, no materialized noise. ``draft_probs=None`` is the
index-form point-mass proposal. ``lse=None`` computes the
normalizer in-op (no-filter slate); a filtered slate passes the
normalizer the mask kernel's LSE epilogue produced. Same output
contract as ``mtp_sample_residual``.

The lossless multi-candidate (tree) verify accept, as a custom op.

The greedy tree accept lives next to the greedy chain accept in
:mod:`arbi_serve._custom_ops.spec_decode`; this is its stochastic
sibling. It is split out for the same reason the GDN verify ops are:
the file it would otherwise join is at its size cap.

The rule it implements is :func:`arbi_serve.spec_decode.tree_spec.
accept_stochastic` — one uniform per depth, inverse-CDF over the
candidate set at that depth, and a draw from the sibling residual at the
rejection point. That module owns the algebra and the argument for why
the rule is lossless; this owns the expression of it as gathers over a
static tree, so the async verify path can keep deferring its host pull
to the drain.

Lossless TREE accept: walk the accepted path, draw at the frontier.

Returns ``(accepted_count (B,), committed (depth+1, B), path_rows
(depth+1, B))`` — the same triple
:func:`arbi_serve._custom_ops.spec_decode._mtp_verify_greedy_tree_kernel`
returns, so every downstream consumer (the drafter seed, the frontier
repair, the KV compaction) reads a tree's scattered accept the same
way whatever the temperature is.

``committed[:n]`` is the accepted path's tokens and ``committed[n]``
is the frontier draw, so a step always commits at least one token.
``path_rows[d]`` is the verify-block row the walk stands on after
``d`` accepts.

THE WALK. At each depth the row stands on some block row; its
children are the candidates for that position. Their target masses
are accumulated in node order and the depth's single uniform picks
the first candidate whose cumulative mass reaches it — inverse-CDF
sampling restricted to the candidate set, which lands on candidate
``x`` with probability exactly ``p(x)``. A duplicate proposal adds no
new mass, so it is excluded from both the accumulation and the hit
set; counting it twice would over-emit that token.

THE FRONTIER DRAW. One draw per row, at the row the walk stopped on,
from that row's target distribution with its candidates zeroed. That
ONE expression covers both cases a chain splits: a row that rejected
stands on a row whose children are the candidates it just declined,
so the zeroing is the sibling residual and the draw is the recovery;
a row that accepted every depth stands on a LEAF, which has no
children, so nothing is zeroed and the draw is the bonus from the
unmodified distribution. The residual is renormalized only in the
first case, which is what makes a width-1 tree's draw bit-identical
to the chain's rather than merely equal in law.

All of it is index_selects and gathers over a static tree: ``depth``
iterations, no data-dependent Python, no host sync.

Lossless multi-candidate TREE accept — the sampled slate's walk.

The stochastic analogue of ``mtp_verify_greedy_tree`` and the
tree analogue of ``mtp_sample_residual``: the per-row
SamplingParams chain is applied to ``p_target`` OUTSIDE this
boundary so the op core stays a pure-tensor function, and the
randomness arrives as operands (the depth-wise accept uniforms,
and either a materialized Gumbel block or the on-device step seed
the Philox draw reads in-kernel) so the op is capture-safe and
bit-deterministic on its inputs.

Returns ``(accepted_count, committed, path_rows)`` — the greedy
tree op's contract, so the drain, the drafter seed and the KV
compaction are temperature-blind.

The copy-engine TP all-reduce ops.

See :mod:`arbi_serve._custom_ops` for the package-level docstring, and
:mod:`arbi_serve.distributed.copy_engine_reduce` for why a second all-reduce
transport exists at all.

The ops are opaque for one reason: the prefill collectives are issued from
inside a Dynamo-traced block, and the transport's peer handshake is host-side
control flow that cannot be traced. An opaque op is baked into the compiled
artifact as a call, its FAKE runs at trace time so the handshake never executes
while tracing, and its real body runs on every execution of that artifact.

Registered here rather than beside the coordinator so that importing
:mod:`arbi_serve._custom_ops` registers them like every other family — a caller
that reaches ``torch.ops.arbi_serve.tp_copy_engine_reduce_*`` without an engine
boot having run must still find the schema.

This module never imports the coordinator. An armed group hands itself in
through :func:`bind_group`, which keeps the dependency one-way and keeps the
group-scoped ``torch.distributed`` call inside ``parallel_state`` — the only
module allowed to make one.

The registered ``(start, wait)`` callables.

Handed out rather than named at each call site, because two gates pull the
other way and this is what satisfies both. A module that spells
``torch.ops.arbi_serve.*`` must import these registrations at module scope
(``test_every_op_caller_registers_the_ops``) — but ``parallel_state`` sits
on the package's root import chain, which must stay torch-free
(``test_package_import_stays_torch_free``), so it cannot carry that import.

Handing out the callables keeps the guarantee the first gate exists for
without the import: a caller cannot obtain a handle without importing this
package, and importing it is what registers the schemas. The dependency is
carried by the value rather than by an import line.

Is the calling thread inside a cudagraph capture?

A named module-level seam rather than an inline call, for two reasons.
It is the ONE place the question is asked on this path, so the guard and
the test that proves the guard fires read the same predicate; and it can
be substituted, which is what lets the refusal be exercised — and its
positive control run — on a host with no device in it.

A build with no CUDA at all is not capturing, and asking torch is an
error there rather than a False, so that case is answered first.

Refuse to run an ARMED exchange inside a cudagraph capture.

A cudagraph is recorded once and replayed. The replay re-issues the
recorded device work and re-runs none of the host control flow that
produced it — and the peer handshake and the consume-ack that make this
transport correct ARE that host control flow. A recorded exchange would
therefore replay its copies against host slots nobody refilled and flag
words nobody advanced: stale bytes, reduced confidently, with no symptom.

So capture is refused here, at the op, rather than answered by a quiet
fall back to NCCL. The fall back would record a graph that is arithmetically
right while the tier the operator armed is silently absent, which turns a
configuration fault into no observable event at all. It is refused HERE
rather than only at arm time because capture begins long after arming, and
a guard that runs only at arm time cannot fire for the case that matters.

Refuse to let a boot pay a cold compile in silence.

Every expensive boot artifact -- the exl3 / xgrammar / cumem cpp extensions,
tkv's split-K kernels, and the ~2-3 min decode autotune table -- is written to
a directory named by an environment variable the image pins onto ``/cache``.
``compose.yaml`` persists ``/cache`` as a fixed-name volume, so a
``docker compose up`` reuses all of it.

A bare ``docker run`` does not. The variables still point at ``/cache/...``,
the directories still exist (the container's own writable layer), every write
still succeeds -- and ``docker rm`` deletes the lot. The next boot recompiles
from scratch and re-sweeps the autotune, having found nothing to reuse. The
Dockerfile already anticipates this launch path for benchmarks and ad-hoc runs.

Nothing about that failure is visible: there is no error, only minutes of
"still running" JIT lines that look exactly like a legitimate first boot on a
new machine. The distinction is one the process CAN make -- an ephemeral
directory sits on the container's root overlay, a persisted one sits on a bind
mount or volume -- so it should make it, once, and say so.

``[(mount_point, fs_type), ...]`` from ``/proc/self/mountinfo``.

Returns ``[]`` off Linux or when the file cannot be read, which makes
every caller treat the question as unanswerable rather than answering it
wrongly -- an unreadable mount table is not evidence of a bad mount.

``[(env_var, path, what)]`` for cache dirs on the container's own layer.

A dir is ephemeral when the mount that owns it is the container root and
that root is an overlay: nothing outside the container holds it, so
``docker rm`` takes it. A bind mount or named volume anywhere along the
path gives that path its own (longer) mount point, which is exactly what
distinguishes the two cases.

Log ONE actionable line when a boot's caches cannot outlive it.

A warning, not a refusal: a throwaway container that genuinely wants to
recompile is legitimate, and refusing it would break the ad-hoc runs this
exists to help. What is not legitimate is paying the cost without being
told, which is the only thing this changes.

Version key for JIT C++-extension build-cache directories.

A JIT-built torch extension (``.so`` + ninja state) is ABI-bound to the
exact torch build, the Python ABI, and the CUDA toolkit that produced it.
A build dir shared across venvs therefore serves a stale ``.so`` to the
wrong interpreter — an import error at best, a silent symbol mismatch at
worst. Every arbi-serve-owned extension build dir nests under
:func:`cpp_ext_cache_key` so distinct torch/python/CUDA builds get
distinct directories, mirroring how torch keys its own DEFAULT cache root
(``~/.cache/torch_extensions/py{maj}{min}_cu{ver}``) — a keying torch
SKIPS whenever ``TORCH_EXTENSIONS_DIR`` is set, which is exactly the case
for arbi-serve's shared cache dirs. ``torch.__version__`` is added on top
of torch's py/cu pair because the libtorch C++ ABI moves between torch
releases even at a fixed Python + CUDA version.

Torch-free when torch isn't already imported: the CLI bootstrap needs the
key before any torch import (it sets ``TORCH_EXTENSIONS_DIR`` for the
whole process, and light subcommands never load torch), so the fallback
reads the tiny generated ``torch/version.py`` (distribution metadata as a
last resort) instead of importing the package.

Contents of the installed ``torch/version.py``, without importing torch.

That file is a small generated module of literals — ``__version__ =
'2.12.1+cu130'``, ``cuda: Optional[str] = '13.0'`` — so reading it costs
microseconds vs the multi-second ``import torch``. Its literals are BYTE-
IDENTICAL to the imported ``torch.__version__`` / ``torch.version.cuda``
(which the wheel's dist metadata is not: it drops the ``+cu130`` local
tag), so both key paths agree.

Torch-free (version, cuda) via ``torch/version.py`` literals.

Falls back to distribution metadata for the version (which may lack
the ``+cuXXX`` local tag) and ``cpu`` for the toolkit — a degraded but
still venv-distinguishing key.

(torch version, CUDA toolkit version) — cheap in both worlds.

When torch is already imported, read the authoritative attributes;
otherwise stay torch-free via the generated ``torch/version.py``.
Both paths report the same strings for the same install, so the CLI
bootstrap (torch-free) and the in-engine callers key identically.

Path component isolating extension builds per torch/python/CUDA.

``torch{torch_version}-cu{cuda_version}-py{major}.{minor}``, each part
sanitized to filesystem-safe characters.

GPU arch is deliberately NOT in the key: the cuMem shim is host-only
C++ (``with_cuda=False``, no device code), and for device extensions
torch's per-extension source/flags hash already rebuilds on a changed
``TORCH_CUDA_ARCH_LIST`` inside the keyed dir.

Fast ``fla`` import: skip its unused subpackages, and start it early.

``fla/__init__.py`` eagerly imports ``fla.layers`` and ``fla.models`` —
nn.Module layer and full-model implementations. arbi-serve uses neither
(it has its own in :mod:`arbi_serve.models`); it wants only the Triton
kernel ops under ``fla.ops`` / ``fla.modules``. Those two subpackages
cost ~3.0s and ~1530 modules of the ~7.0s ``import fla`` charges, and
every ``fla.<anything>`` import runs that ``__init__`` first — so
importing a narrower submodule does NOT avoid it.

Two levers, independently switchable:

``install_stubs()``
    Places empty placeholder modules at ``fla.layers`` / ``fla.models``
    for the duration of the ``fla`` import, so ``__init__``'s eager pull
    is a no-op, then REMOVES them. Removing matters: a later genuine
    ``import fla.layers`` then loads the real package (paying its cost)
    instead of silently getting an empty module. Torch-free — it only
    touches :data:`sys.modules` — so it is safe to call from
    :mod:`arbi_serve`'s package ``__init__``, which must stay torch-free.

``start_background_import()``
    Imports ``fla`` plus the kernel helpers on a daemon thread so its
    remaining ~4s overlaps the flat-dump weight DMA instead of sitting on
    the critical path. Costs torch, so it must be called from a path that
    has already committed to torch — never from package import.

Python's per-module import locks make the background thread safe against
a concurrent main-thread ``import fla``: the main thread blocks until the
background import completes, which is exactly the join we want.

Opt out with ``ARBI_FLA_FASTIMPORT=0`` (stubs) / ``ARBI_FLA_BG_IMPORT=0``
(background thread); both default ON.

``import fla`` with the unused subpackages stubbed out.

Returns True when the stubbed path ran. A no-op returning False when
``fla`` is already imported (the cost is already paid) or the flag is
off. Never raises on a missing ``fla`` — the CPU/test path has none,
and the GDN call sites already degrade on ImportError.

Disable FLA's identity-keyed tensor cache + replace prepare_chunk_indices
+ pin ``chunk_gated_delta_rule`` fwd workspace to a persistent pool.

Three-part shim that makes hybrid prefill capture safe AND keeps the
captured-graph mempool small across multi-bucket sweeps:

1. Sets ``fla.utils.FLA_DISABLE_TENSOR_CACHE = True`` so every
   ``@tensor_cache``-decorated FLA helper (``prepare_lens``,
   ``prepare_chunk_offsets``, etc.) re-runs its underlying impl on
   every call.  Without this, identity-keyed single-entry caching
   pins stale ``data_ptr``s across multi-bucket capture sweeps and
   replays trip ``cudaErrorIllegalAddress`` in
   ``chunk_local_cumsum_scalar_kernel`` /
   ``chunk_gated_delta_rule_fwd_kernel_h``.

2. Replaces ``fla.ops.utils.index.prepare_chunk_indices`` with an
   identity-keyed persistent-cache wrapper that ALSO never evicts.
   The flag from (1) is read by the ORIGINAL FLA wrapper closure;
   our replacement bypasses that closure entirely, so we control
   the discipline directly.  Required because under the env-var-only
   path, the CAPTURE forward calls ``prepare_chunk_indices`` for
   the FIRST time inside ``torch.cuda.graph(...)`` — the
   ``.tolist()`` host sync inside the underlying impl trips
   ``cudaErrorStreamCaptureInvalidated`` and aborts every bucket
   except N=1 (where the GDN block routes through the decode
   kernel and never calls ``prepare_chunk_indices``).

Part (3) — workspace pinning — is independent of (1) and (2): it
addresses a different problem (captured-pool VRAM footprint, not
capture-time correctness).  Across a prefill-bucket capture sweep,
the FLA fwd's intermediate workspaces add to the captured-graph
mempool per bucket, since each bucket captures its own workspace
allocations independently.  Routing those allocations to a
separate ``NamedMemPool`` (the engine's ``gdn_workspace_pool``,
sized to the largest bucket) bounds the workspace cost at one
pool's worth.  Pool routing uses ``_cuda_beginAllocateCurrentThreadToPool``
directly (NOT ``torch.cuda.use_mem_pool``, which would call
``_cuda_releasePool`` on exit and destroy the pool the first time
the wrapper exits).

Why both parts.  Part (1) alone is insufficient: the host sync still
runs on the CAPTURE forward's first call.  Part (2) alone is
insufficient: the OTHER ``@tensor_cache``-decorated helpers
(``prepare_chunk_offsets``, ``prepare_lens``) still cache by
identity, so their fresh allocations pile up and stale ``data_ptr``s
break replay.  Together they let captures complete cleanly AND keep
replay-time ``data_ptr``s valid:

  * Pre-warm: ``_capture_one_bucket``'s warmup forward (OUTSIDE
    ``torch.cuda.graph``) populates our identity-keyed cache for
    ``prepare_chunk_indices`` AND the FLA-flag-disabled helpers
    allocate freshly (data_ptrs land in the regular allocator,
    cached forward storage stays alive via the captured-graph's
    internal references).
  * Capture: same persistent-cache returns the warmup-allocated
    ``chunk_indices``; FLA-flag-disabled helpers re-allocate fresh
    inside the ``torch.cuda.graph(pool=mem_pool)`` region — those
    allocations land in the cudagraph mempool and are correctly
    re-allocated to the same data_ptr on replay.

Memory.  Our cache holds one ``chunk_indices`` tensor per captured
prefill bucket (~few KB).  Negligible.

Ownership.  Both caches hold DEVICE tensors allocated in whichever
member's pool was routing when they were filled, and the wrappers are
installed once for the PROCESS (both installers are idempotent by tag /
name), so a second stable-VA member reuses the first member's wrapper.
That makes the cache contents member state living behind a process-global
name — the family
:mod:`arbi_serve.engine.member_scratch_retire` exists for.  Both are
therefore module-level mappings (:data:`CHUNK_INDEX_CAPTURE_CACHE`,
:data:`CHUNK_INDEX_STEP_MEMO`) rather than closure cells, so the park can
claim them onto the parking member's record, the build seam can retire
what a park did not, and the wake can hand a member back its own entries
at the VAs its captured graphs baked.

Implementation.

  * :func:`patch_fla_chunk_indices_cache` walks ``sys.modules`` and
    rebinds every ``prepare_chunk_indices`` reference (both the
    source module ``fla.ops.utils.index`` and every downstream
    ``from fla.ops.utils.index import prepare_chunk_indices``
    re-export) to our wrapper.
  * Idempotent: re-application is a no-op (the wrapper carries a
    ``_PATCH_TAG`` attribute).
  * Called from :func:`precapture_layer_graphs` BEFORE any per-
    bucket capture fires.

Numerical correctness.  Identical math to upstream — we call the
underlying ``__wrapped__`` impl on every cache miss; only the cache
discipline changes.

Point every module-level name bound to ``old`` at ``new``, reversibly.

ONE walk, used by both patches. Two copies of it is what let the capture
wrapper and the step memo drift apart on which references they moved.

Undo every rebind this module made, newest first. Returns how many.

THE INVERSE OF THE TWO PATCHES, and the reason it exists is that they are
process-global: :func:`patch_fla_chunk_indices_cache` is reached from three
capture-admin sites, so any in-process boot installs the capture wrapper
for the life of the interpreter. In a server that is correct — there is one
boot. In a test session it is not: the wrapper outlives the test that
installed it, :func:`install_chunk_indices_step_memo` then declines on the
``_PATCH_TAG`` guard, and a later test's ``prime_chunk_indices_step_memo``
finds no ``store`` to publish through. That is an ordering interaction, not
a flake — the second test's result depends on whether the first one ran.

Also drops both caches: their values are device tensors from the pool of
the member that filled them, so carrying them past a restore would hand a
later caller a table allocated in a pool that no longer exists.

Identity-keyed persistent cache.

Key = ``(id(cu_seqlens), int(chunk_size), id(cu_seqlens_cpu))``.
Value = ``(cu_seqlens, cu_seqlens_cpu, chunk_indices)`` — strong
refs so neither input can be GC'd (which would let Python recycle
its ``id()``) and the chunk_indices storage is pinned for as long
as the member that allocated it can still replay a graph that baked
the address.

Identity-keyed lookup avoids any host sync on the cache hit path
— essential because cache lookup runs INSIDE the
``torch.cuda.graph(...)`` capture region on the second call per
bucket (the warmup forward populates; the capture forward looks
up).

The cache itself is :data:`CHUNK_INDEX_CAPTURE_CACHE` — module-level so a
member park can claim it; see that attribute.

Drop every entry FLA's ``tensor_cache`` memoised before it was disabled.

``fla.utils.tensor_cache`` keeps a module-level ``deque`` of
``(args, kwargs, result)`` triples per decorated helper (``prepare_lens``,
``prepare_chunk_indices``, ``prepare_chunk_offsets``, …).  Disabling the
cache stops new entries but does not release the ones the boot's
pre-patch forwards already made — and those forwards run inside
``scratch.forward_arena``, so a handful of small index tensors ends up
PERSISTENT in a scratch pool whose free physical the engine gives back at
its idle quiesce point.  Nine of them are enough to forfeit that whole
release (:func:`~arbi_serve.engine.inprocess_capture.
release_idle_forward_arena` refuses on a non-empty arena, whatever its
size).

Dropping them costs one recomputation, which the now-disabled cache would
force anyway.  Runs ONCE, from the first patch application, which the
caller's contract places before any per-bucket capture fires — so no
captured graph can have baked one of these addresses.  Returns the number
of caches cleared.

Turn FLA's identity-keyed ``tensor_cache`` off and drop what it holds.

Separate from :func:`patch_fla_chunk_indices_cache` because it must run for
EVERY model, not only the GDN-layer models that need the capture-time
patches: the cache is populated by any FLA helper the stack reaches, and
its entries are what strand the forward arena. No-op when FLA is not
installed. Returns the number of caches cleared.

Single-entry, per-step memo for ``prepare_chunk_indices``.

Restores the compensation that :func:`disable_fla_tensor_cache` removes.
FLA's own ``@tensor_cache`` is a single-entry identity cache; every GDN
layer in a step calls ``prepare_chunk_indices`` with the SAME
``cu_seqlens`` object (``GdnMetadataBuilder._allocate_meta`` aliases
``batch.cu_seqlens_q`` onto the meta once per step), so upstream served
47 of 48 calls from that cache. Disabling it — which capture needs, and
which ``engine/build.py`` therefore does unconditionally — turns the
impl's blocking device read into a per-layer stall. That read is
``_segmented_arange``'s ``torch.repeat_interleave``: with device
``counts`` it must learn the data-dependent output length
``counts.sum()`` before it can allocate, so the host waits. FLA's own
docstring names the remedy ("pass host-side counts to avoid it"), which
is what the ``cu_seqlens_cpu`` route beside this memo does.

SINGLE ENTRY, not a growing dict, and the entry is held by WEAK
reference. Both properties are load-bearing:

  * The eager path mints a fresh ``pb.cu_seqlens_q[: B + 1]`` VIEW every
    step, so an identity-keyed dict that holds strong refs (as
    :func:`_build_persistent_wrapper` deliberately does, for capture)
    would grow without bound across a server's life and never evict.
  * Holding the key tensor weakly means a freed ``cu_seqlens`` cannot
    have its ``id()`` recycled onto a live entry: the entry dies with it.
    Identity is then sound to compare, which is what lets the hit path
    run with no host sync at all.

Contents are never re-read: the cached value is keyed on the tensor
OBJECT, and a step that writes new offsets into the persistent buffer
hands out a new view object for them.

THE HOST TWIN IS NOT PART OF THE KEY. ``cu_seqlens`` and its twin
describe the same boundaries — :func:`~arbi_serve.models.
_gdn_fla_forward._resolve_cu_seqlens` returns a twin only when it
provably describes the very tensor beside it — and the two routes build
a table that is equal element for element, so a hit computed with the
twin answers a call made without one. Keying on the twin instead splits
one step's table in two: the layers inside the opaque custom op have no
twin to pass, so a table resolved with one could never serve them, and
every step paid the device route at least once. Identity of the device
tensor is the whole key because it is the whole question.

FLA's chunk-index table, built on the host and staged through pinned memory.

Row ``r`` of the table is ``(sequence, chunk-within-sequence)`` for the
``r``-th chunk of the flattened batch, sequences in order — the same table
``fla.ops.utils.index.prepare_chunk_indices`` returns, and pinned by a test
against it.

Built here rather than asked of FLA because FLA's host route is only half
host: it derives the table from host counts and then uploads it with
``.to(cu_seqlens)`` on PAGEABLE memory, which the CUDA runtime serializes
as a synchronous copy. Staging through the caching host allocator's pinned
memory makes the same upload asynchronous, which is the whole point of
resolving the table before the forward rather than inside it.

``like`` supplies the dtype and device, exactly as FLA's ``.to(cu_seqlens)``
does.

Resolve this step's chunk-index table from the HOST twin, once.

The GDN layers reach ``prepare_chunk_indices`` from inside the opaque
``arbi_serve::gdn_attention_v2`` custom op, which carries no host twin
and cannot be handed one: a CPU tensor consumed inside the captured
region is not what the vendored entry's ``chunk_indices`` parameter
exists to avoid. So the twin has to be spent BEFORE the forward, on the
host, where the step's metadata is built — and then the memo hands every
layer the same table.

Without this the first layer of every step resolves the table the only
way it can from inside the op: FLA's device route, whose
``_segmented_arange`` sizes a ``repeat_interleave`` by ``counts.sum()``
and so blocks the host on a device read.

Returns whether a table was resolved, so the caller can account for it.

Drop the per-step chunk-indices entry. Returns entries released.

Reached from the forward-arena release
(:mod:`arbi_serve.engine.arena_step_caches`), which destroys the pool the
entry's tensor was allocated in. Sound there because the memo is keyed on
the IDENTITY of the step's ``cu_seqlens`` view and the eager path mints a
fresh view every step: the next call misses whether or not this runs, so
the drop costs nothing the disabled upstream cache was not already costing.

Scoped to the STEP memo by name. The capture path's persistent cache
(:func:`_build_persistent_wrapper`) deliberately never evicts — its entries
are addresses replays have baked — so this must find nothing to clear when
that one is installed, and it does: that wrapper carries no ``clear_memo``.

Install the per-step ``prepare_chunk_indices`` memo. Idempotent.

Called UNCONDITIONALLY from the boot, at the same seam as
:func:`disable_fla_tensor_cache` and for the same reason it is
unconditional: the disable is what creates the need, and the disable is
not keyed to GDN layers.

Deliberately does NOT carry ``_PATCH_TAG``, and exposes the RAW impl as
``__wrapped__``. A capture boot's :func:`patch_fla_chunk_indices_cache`
therefore still finds an untagged callable with the raw impl under it and
replaces this wrapper outright — the capture path behaves exactly as it
did before this memo existed.

Returns True when a wrapper was installed.

Register the workspace pool the FLA chunk-fwd allocations route through.

Called from the engine's piecewise-capture boot path BEFORE the
capture sweep fires.  Until set, :func:`patch_fla_chunk_workspace`
leaves ``chunk_gated_delta_rule`` unwrapped.  Passing ``None``
clears the registration.

Return the currently-registered FLA chunk-fwd workspace pool (or
``None`` when unset — the default-allocator / never-pin routing).

Used by the boot activation profiler to fail loud if the workspace is
pinned while the profile forward runs: a pinned pool would divert the
eager FLA chunk-fwd workspace away from the profiled ``scratch.forward_arena``
arena and under-reserve the activation peak.

Wrap ``chunk_gated_delta_rule`` to route its internal
intermediate-tensor allocations through ``_GDN_WORKSPACE_POOL``.

The wrapper runs the FLA call inside the pool's own
:meth:`arbi_serve.runtime.named_pool.NamedMemPool.use` context — NOT a raw
``_cuda_beginAllocateCurrentThreadToPool`` — because ``use()`` is the ONLY
safe routing primitive on the serving path.  It handles, in one place, the
three things a raw begin/end does NOT and which each independently corrupt
the allocator when the eager GDN scan runs between the captured per-layer
graphs:

  * **Nested same-pool entry** — a second begin for a pool already being
    recorded raises ``RuntimeError: beginAllocateToPool: already recording
    to mempool_id``.  ``use()`` refcounts (``_active_depth``) and drives
    begin/end only on the outermost pair.
  * **``expandable_segments:True``** — the engine's global default, under
    which ``torch.cuda.MemPool`` begin/end refuses to run.
    ``use()`` scopes ``expandable_segments`` OFF for the begin→alloc→end
    window on cuMem-backed pools.
  * **cuMem tag stamping** — ``use()`` stamps each allocation with the
    member's pool tag so a residency park's ``sleep_all(namespace=…)``
    can UNMAP the workspace.  A raw begin leaves the bytes untagged, so
    even a routed workspace would survive park — defeating the whole point.

``use()`` also no-ops during cudagraph capture (allocations flow into the
captured graph's own pool).  When ``_GDN_WORKSPACE_POOL`` is ``None`` (CPU /
test / arena-backed OFF), falls through to the unwrapped call.

Wrap ``chunk_gated_delta_rule`` so its internal allocations route
through ``_GDN_WORKSPACE_POOL``.

Idempotent.  No-op when FLA isn't installed.  Walks ``sys.modules``
and rebinds every ``from fla.ops.gated_delta_rule import
chunk_gated_delta_rule`` re-export to the wrapper, mirroring the
sys-modules walk in :func:`patch_fla_chunk_indices_cache`.

Publish a table someone else resolved, under this memo's key.

The caller that has the host twin is the metadata build, and it
cannot reach FLA's own resolution without paying for it (FLA
uploads its host-built table with a PAGEABLE ``.to(cu_seqlens)``,
which the runtime serializes as a synchronous copy). So it builds
the table its own way and publishes it here; the layers then hit
the memo exactly as they would have.

Run ``underlying`` with its allocations routed through the GDN
workspace pool via ``NamedMemPool.use()``, or fall through unwrapped
when no pool is set or the routing context cannot be entered.

Defensive: if entering the pool context raises (an unexpected
allocator-state edge on the serving path), the FLA call runs UNROUTED
rather than IMA-ing a live request — the workspace then lands in the
default allocator and the park-time co-residency assert fails LOUD at
the next swap (a loud boundary error, never a serving-hot-path crash).
We only guard the context ENTER; ``underlying``'s own exceptions
propagate, and it is run exactly once on every path.

Run a callback the moment a module is imported, not before.

A patch or a hook that targets a third-party module has to be installed before
anything can use it. Importing the module to install it makes that dependency
unconditional: every boot pays for it even when nothing on the path would have
loaded it. Arming on the import itself keeps the ordering guarantee — the module
cannot be used before it is imported — at zero cost when it never is.

Torch-free; a single ``sys.meta_path`` finder serves every registration.

Call ``callback(module)`` once ``name`` is imported.

Runs immediately when the module is already in ``sys.modules``. Callback
exceptions are swallowed: a hook that fails must not break the import that
triggered it.

Self-healing replacement for torch's cpp_extension JIT build lock.

``torch.utils.file_baton.FileBaton`` is a presence-only lock: ``try_acquire``
creates an EMPTY lock file, ``wait`` polls ``os.path.exists`` FOREVER, and
``release`` removes it. If the holding process dies mid-build (crash, OOM,
SIGKILL), the lock file is never removed — so every later JIT compile of that
extension blocks in ``wait`` indefinitely. We hit exactly this: killed GPU runs
left ``~/.cache/torch_extensions/*/lock`` behind and the next server's kernel
warm wedged on the dead baton.

This patches ``FileBaton`` to be PID-aware and self-healing:

  - ``try_acquire`` stamps the holder PID into the lock file.
  - A contender that finds the lock checks whether that PID is still alive. If
    the holder is DEAD it steals the stale lock (removes it and acquires), so a
    crash can't wedge future builds. A live holder is never stolen (no
    concurrent-build corruption). A PID-less lock (pre-patch, or written by a
    non-arbi process) is treated as stale only after ``_STALE_SECONDS``, well
    beyond any real compile.
  - ``wait`` breaks the lock and raises :class:`JitLockAbandoned` when the
    holder dies mid-wait. ``torch.utils.cpp_extension._jit_compile`` compiles
    ONLY on the ``try_acquire() -> True`` branch and imports the ``.so`` as
    soon as ``wait`` returns, so returning here after breaking an abandoned
    lock leaves nobody building and the caller dies with ``ImportError:
    <name>.so: cannot open shared object file``. The removed lock lets the
    next ``try_acquire`` become the builder.
  - ``wait`` is BOUNDED and announces itself. The steal path above only
    covers a holder we can prove is gone; a lock whose holder looks live but
    never finishes — or whose liveness signal is unreadable (``/proc``
    masked, a different PID namespace) — is not stealable, and torch's loop
    then sleeps with no deadline and no output. That is indistinguishable
    from a slow build while it burns a boot. Past ``ARBI_JIT_LOCK_TIMEOUT_S``
    it raises :class:`JitLockTimeout` naming the lock and the remedy.
    turbo-attn's ``TKV_JIT_LOCK_TIMEOUT_S`` governs the same wait once its
    own bounded baton installs (at its first kernel build); both are bounded,
    so the boot cannot hang under either.

The bound matters for extensions this module is the ONLY guard for.
``_prebake_loader`` wraps ``cpp_extension.load``, but ``load_inline`` reaches
``_jit_compile`` directly and never passes through it — so the ``flock``
serialization and ``clear_stale_extension_lock`` in
:mod:`arbi_serve._cpp_ext_locks` do not apply to any ``load_inline`` caller.
The cuMem shim (:mod:`arbi_serve.runtime._cumem_shim`) is one, and it builds
into its own ``build_directory`` outside ``TORCH_EXTENSIONS_DIR``, so no
sweep reaches its lock either.

PID REUSE across container restarts: a bare PID-liveness check is NOT enough.
The container PID namespace resets to low PIDs on every boot, so an abandoned
lock stamped (say) ``pid:93`` from a prior killed run routinely collides with
the CURRENT rank-0 process — which also lands on PID 93. ``_pid_alive(93)``
then returns True, the lock looks live, and EVERY rank (including the one whose
number was reused) spins in ``wait`` forever on a lock no live builder holds —
a hard TP boot deadlock. The fix:
treat a stamped PID as the holder only if that PID's process STARTED BEFORE the
lock was written (``_holder_started_before_lock``). A lock older than the
process that now bears its number was written by a dead predecessor → steal it.

Installed once from ``arbi_serve.__init__`` (torch-present path), before any
submodule can trigger a ``load_inline``. Idempotent.

Monotonic-ish process start time of ``pid`` (kernel ticks since boot).

Reads field 22 (``starttime``) of ``/proc/<pid>/stat``. Returns None when
the proc is gone or ``/proc`` is unreadable. The unit (clock ticks since
system boot) is comparable across processes in the SAME PID namespace,
which is all we need to order a process against a lock-file mtime.

``stat`` field 22 is positionally safe even though ``comm`` (field 2) may
contain spaces/parens: we split on the LAST ``)`` and index from there.

True if ``pid``'s process predates the lock file (so it could own it).

Converts the lock's wall-clock mtime and the process's boot-relative
starttime to the same boot-relative seconds axis via the system boot epoch
(``/proc/stat`` ``btime`` + ticks). When either signal is unreadable we
answer True (conservative: keep treating a live PID as the holder, falling
back to the pre-existing liveness-only behaviour) so we never steal a lock
a real concurrent builder holds.

Identity of the CURRENT lock file at ``path``:
``(st_ino, st_mtime_ns, st_size)``.

Used to make the steal path generation-safe: the staleness verdict is
computed on one lock file, but by the time ``os.remove`` runs a NEW
builder may have re-created the lock at the same path — removing that
fresh lock hands the "lock" to two processes at once, which is exactly
the concurrent-build/partial-``.so``-dlopen race ("invalid ELF header").
A steal must therefore name the generation it judged stale and refuse
to remove any other.
Returns ``None`` when the lock is already gone.

Nanosecond mtimes come from the kernel's coarse file-time clock (ms-ish
granularity), and inode numbers are eagerly reused — so neither alone
distinguishes a lock re-created immediately at the same path. The
triple (+ size: the PID stamp differs across holders) makes a same-
granule collision require inode reuse AND equal-length stamp AND
same-tick creation; the residual window is closed for arbi processes
by the flock serialization (see :func:`_steal_lock_generation`).

Remove the lock at ``path`` iff it is still the ``generation`` that
was judged stale. Returns True when the path is clear (removed by us or
already gone), False when a DIFFERENT (fresh) lock now sits there.

Note the residual window between the re-stat and the ``os.remove`` is
nanoseconds (vs. the minutes-long build window of the naive remove) and
is fully closed for arbi processes by the ``flock`` serialization in
:mod:`arbi_serve._cpp_ext_locks` — every arbi ``load()`` holds that lock,
so no live arbi builder can re-create the baton while another arbi
process is inside this code.

Remove the lock at ``path`` iff it is stale and still the generation
that was judged stale. True ONLY when this call removed it.

Distinct from :func:`_steal_lock_generation`, whose True also covers "the
lock was already gone" — the right answer for ``try_acquire`` (retry the
create) and the wrong one for ``wait``, where a holder releasing normally
between the poll and the stat is a COMPLETED build, not an abandoned one.

Whether the lock at ``path`` can be safely stolen.

Stale iff the stamped holder PID is provably dead OR provably a PID-reuse
impostor (its process started AFTER the lock was written, so it cannot be
the original builder — see :func:`_holder_started_before_lock`), or — for a
PID-less lock — it is older than :data:`_STALE_SECONDS`. A live stamped PID
that predates the lock is NEVER reported stale, so a genuine in-progress
build is never stolen.

Make ``libcudart`` symbols globally visible.

TileLang (pulled in transitively by ``fla`` / flash-linear-attention,
used by GDN/Mamba arches, and by the decode-capture path) resolves the
CUDA runtime via ``dlsym(RTLD_DEFAULT, ...)`` — it only finds the
symbols if some library loaded libcudart with ``RTLD_GLOBAL``. When the
symbols are not global it HARD-ABORTS the interpreter ("TileLang Error:
libcudart symbols not found globally"). The exllamav3_ext extension is
``dlopen``\ed ``RTLD_LOCAL``, so when an EXL3 checkpoint and a GDN/Mamba
arch coexist in one process the runtime ends up private and TileLang
aborts.

Loading libcudart ourselves with ``RTLD_GLOBAL`` BEFORE ``fla`` /
TileLang import makes the symbols global so both the EXL3 kernel and
TileLang resolve against the same runtime regardless of load order.
Best-effort and idempotent: a static-linked or oddly-named libcudart
just no-ops and the pre-existing behaviour stands.

Emit a log line only when it says something the last one did not.

For a diagnostic that is RE-DERIVED several times in one boot and comes
out the same every time — a memory floor recomputed at each grow
attempt, a provenance breakdown re-stated per phase — the first emission
is the whole content and the repeats are noise. Repeating a dense line
verbatim also costs a reader more than a terse one: it hides the lines
between the copies.

:func:`log_if_changed` keeps the full text on the emission that carries
new content and demotes an exact repeat to DEBUG. Nothing is shortened
and nothing is dropped — a re-derivation whose result CHANGED still
prints in full at INFO, which is the case an operator must not miss.

Keys are caller-chosen and process-global; use one per call site.

Log ``msg % args`` at ``level`` if it differs from ``key``'s last.

An exact repeat goes to ``repeat_level`` instead. Returns ``True``
when the message was new (and so reached ``level``).

The message is rendered eagerly — comparison needs the text — so this
belongs on boot / control paths, not on a per-step hot path.

Forget the last message for ``key`` (or all keys when ``None``).

A fresh process starts empty; this exists for tests and for a
re-boot-in-process (model reload) that should restate its diagnostics.

``python -m arbi_serve dump_openapi`` — write the OpenAPI spec to stdout.

Used by ``scripts/generate-client.sh`` to drive
``openapi-python-client`` without having to boot a server. The OpenAPI
spec is a pure function of the FastAPI route table + Pydantic models —
no engine, no GPU, no calibration files needed.

Why a separate FastAPI instance instead of ``build_app(cfg)``? The
production ``build_app`` requires a fully-validated :class:`ServerConfig`
(including a model path) and triggers the engine build inside its
lifespan. For schema-only export we just need the routes registered;
re-using the routers is enough.

Stays import-light by design: we touch only ``fastapi`` and the two
router modules, no torch / no uvicorn / no recipes.

Group endpoints into client API modules by URL prefix.

The runtime routers don't carry FastAPI tags (we wire them lean
for serving), but ``openapi-python-client`` uses tags as the
module-grouping key for the generated ``api/`` tree. Without
tags every endpoint lands in ``api/default/``, producing a
single bloated module. We infer a tag from the path so the
generated client mirrors the URL hierarchy (admin / chat /
completions / models / loras / health).

Wrap torch's JIT extension entry points — global stale-lock guard +
optional prebaked-.so short-circuit.

Three jobs across the two entry points every JIT extension flows through:
``cpp_extension.load`` (exl3, AWQ Marlin, xgrammar, …) and
``cpp_extension.load_inline`` (every tkv CUDA kernel), which reaches
``_jit_compile`` on its own and so needs the lock guard wrapped separately:

  1. **Stale-lock guard (always on).** Before every build, clear any stale
     ``FileBaton`` lock for the extension (see
     :mod:`arbi_serve._cpp_ext_locks`). torch's lock has no release on
     process death, so a SIGKILLed build wedges every later ``load()``
     forever. Patching at this single chokepoint covers every JIT
     extension, on every box (not just the slim image).

  2. **Prebake short-circuit (slim image only).** The slim Docker image
     strips nvcc/g++ but pre-bakes select extensions into
     ``/opt/cache-baked/``; torch re-emits build.ninja each process and
     would try (and fail) to recompile. When a pre-baked
     ``<dir>/<name>/<name>.so`` is present and the caller used
     ``is_python_module=False``, skip the JIT path and load via
     :func:`torch.ops.load_library`.

  3. **Source-digest guard (prebake path).** A pre-baked ``.so`` is keyed
     by extension NAME alone, so any change to the installed sources — a
     version bump, a git-ref pin, a bind-mounted checkout — resolves to
     the same artifact. The compile path records a digest of the sources
     it built from; the short-circuit recomputes it and refuses a
     PROVEN mismatch (recorded digest != installed sources).

     A bake that recorded NO digest is unverifiable, not wrong — absence
     of a record is not evidence of a mismatch, and refusing it would
     reject every image whose bake predates digest recording. Those load
     with one warning per extension.
     ``ARBI_SERVE_PREBAKE_REQUIRE_DIGEST=1`` turns unverifiable back into
     fatal for deployments that want to gate on a labelled bake.

Content digest over an extension's sources and their sibling headers.

Hashes the ``sources`` torch was handed plus every source/header file
under their common ancestor, so a change confined to a ``.cuh`` — where
templated CUDA kernels actually live — moves the digest. Paths are
recorded relative to that ancestor, so a digest travels between the
builder and runtime stages. Returns ``None`` when there is nothing to
hash, which leaves the caller's behaviour unchanged.

Load a prebaked ``.so`` that carries no recorded source digest.

Unverifiable is not the same as mismatched: a bake predating digest
recording writes no sidecar, so refusing it rejects a binary that is
almost certainly correct and gives the operator no way forward. Warn
once per extension and proceed; ``ARBI_SERVE_PREBAKE_REQUIRE_DIGEST=1``
restores the hard refusal where a labelled bake is a deployment gate.

Decide what a prebaked ``.so`` that fails its source-digest check gets.

Returns True when the bake is proven stale AND this runtime can build the
extension from the installed sources, in which case the caller falls
through to the ordinary JIT build (serialized, cached under
``TORCH_EXTENSIONS_DIR`` like any first build) instead of serving machine
code the sources no longer describe. The image's bake is a boot-time
saving, not the source of truth: under a bind-mounted or re-pinned tree
the installed sources are what the deployment claims to serve, and a
stale bake that silently disabled an optional leg (the exl3 int8 GEMM
refuses by name when its kernel does not load) would report a
performance number for code the tree does not contain.

A runtime without nvcc (the slim image) cannot build anything, so there
the digest failure keeps refusing loudly, as before. Whether the
extension is a Python module (torch's default, and what exllamav3 and
the exl3 int8 GEMM both are) changes only how torch imports the result,
not whether it can build it.

Serialize an inline-source JIT build the same way ``load`` is.

``load_inline`` reaches ``_jit_compile`` directly and never calls
``load``, so the flock above did not cover it — and every tkv CUDA
kernel (split-K decode, MTP verify, the native int8 prefill kernel)
is a ``load_inline`` caller building into a ``TORCH_EXTENSIONS_DIR``
that co-tenant containers and TP ranks share. Unserialized, two
builders link the same ``.so`` concurrently and a third dlopens the
half-written file; torch's own FileBaton cannot fix that, because
its staleness recovery can delete a live builder's lock. Same guard,
same reasons, second entry point: the seam is the extension name,
not the API that reached it.

No prebake short-circuit here. ``load_inline`` compiles from source
TEXT, so there is no ``sources`` list to digest and nothing to
short-circuit to; a baked module reaches it by being seeded into the
build dir (``cli.bootstrap._seed_prebaked_torch_extensions``), which
carries its own source-digest guard.

Process-wide shims the RUNTIME needs, installed on demand.

Importing anything under ``arbi_serve`` — the admin console included — must not
pay for torch, fla or the GDN custom ops; the console is an HTTP proxy that
needs none of it.

Ordering is load-bearing and is guaranteed by DEPENDENCY rather than by
package-root eagerness: :mod:`arbi_serve._custom_ops` calls
:func:`install_runtime_shims` before it registers anything, so any path that
reaches the kernels installs the prerequisites first. Idempotent — later calls
return immediately.

Force transformers' flash-attn probes to False, as transformers loads.

flash-attn is never installed on this stack; the attention path is
Turbo prefill. transformers' availability probes
(``is_flash_attn_{2,3,4}_available`` & friends) index
``PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]`` UNGUARDED, and that lookup is
reached whenever ``find_spec("flash_attn")`` is truthy — which the EXL3 shim
makes happen by installing a temporary ``flash_attn`` sys.modules stub
(exllamav3 imports flash_attn at module scope). The probe then believes
flash-attn is present and the unguarded lookup raises
``KeyError: 'flash_attn'``, crashing the import of
``transformers.activations`` — reached transitively by
``fla.layers.log_linear_mamba2`` (GDN) and by exllamav3's model stack at
EXL3 load.

Armed on the import rather than by importing transformers here: transformers
is a second-scale import and the server does not use it as a model runtime,
so a boot that never reaches it must never pay for it.

Canonical byte units and ceil division.

Single source for the binary byte units and the ceil-division idiom so
the same literals are not re-spelled across the codebase. Pure stdlib
arithmetic; imports nothing from arbi_serve (safe to import anywhere).

Column count for a captured/persistent paged-KV block-table row.

A row addresses the pages one sequence can reference — the larger of
its ``max_context`` window and the pool's page ceiling (the max pages
the growable pool can map). Sizing at the ceiling keeps ``max_context``
settable post-grow without re-capture; the extra columns are int32
zeros the kernel never reads. Falls back to the ``max_context`` window
when the ceiling is unknown (``0`` / pre-profile / non-paged-KV).

Adapter Protocol — generic per-linear adapter dispatch.

Per-linear adapter dispatcher. The model's parallel linears are
wrapped at construction by a generic adapter dispatcher; per-request
:class:`AdapterState` selects which adapter (if any) is active for
that request.

The :class:`LinearBase` hook is a thin trampoline through
:func:`apply_adapter`.

Per-linear adapter dispatch.

The contract is intentionally minimal: given the linear's already-
computed output ``y`` and the input ``x`` plus the linear's
``layer_key`` (for routing the right per-layer adapter weights),
return the adjusted output. LoRA implements this as
``y + (x @ A) @ B * scale``.

The wrap happens at load time. The forward path branches on the
request's :class:`AdapterState`; if no adapter is active for the
request, the dispatcher short-circuits to ``y``.

Implementations MAY return ``y`` itself (in-place update) or a
fresh tensor; callers must use the returned value. The default
when no per-layer weights exist for ``layer_key`` is the identity
(return ``y``).

Per-request adapter selection + the dispatch entry point.

Lives on :class:`arbi_serve.engine.request.Request.adapter_state`
and on :class:`arbi_serve.engine.batch.ScheduledBatch` per-step
(the engine collapses per-request states into one per-step state
that the model sees).

The Protocol unifies two roles:

  1. Identity / activation status (``adapter_id``, ``is_active``).
  2. The :meth:`apply_to_output` entry point (same shape as
     :class:`Adapter`).

Conflating the two is deliberate: :class:`LoraBatchState` already
holds both identity (``adapters`` list) and per-target weights.
Splitting them would force every consumer to carry a
``(state, adapter)`` tuple where one is enough.

Generic per-linear adapter dispatch trampoline.

The :class:`arbi_serve.models.linear.LinearBase` family calls this
once per forward, after the base matmul. It encodes the four
short-circuits that every linear should observe:

  1. ``adapter_state`` is ``None`` → no adapter on the batch.
  2. ``adapter_state.is_active`` is ``False`` → adapter present
     but every active adapter is a no-op for this batch.
  3. ``layer_key`` is ``None`` → linear is adapter-opaque (e.g. a
     router weight in MoE that no LoRA is allowed to touch).
  4. The adapter's :meth:`apply_to_output` returns ``y`` itself
     when no per-layer weight exists for the linear's key.

All four collapse to "return ``y`` unchanged". The dispatcher is
deliberately small so it inlines on the fast path.

Misuse: if a linear is wrapped without a layer key but the
request carries an active adapter targeting that layer, the
contribution silently drops (we have no key to route by). Models
that opt linears into the adapter mechanism must call
:meth:`set_lora_target_name` at construction.

Multi-LoRA serving package.

LoRA adapters are first-class per-request dials, same shape as
``?attention_backend=`` and ``?calibration=``. Orthogonal to TKV: a
LoRA modifies projection weights, TKV compresses the KV cache; a
single request can carry both
(``?lora=X&attention_backend=tkv-k4v4``).

Public surface:

- :class:`LoraAdapter`     — one loaded adapter (rank, α, per-target weights).
- :class:`LoraInfo`        — admin-API view of an adapter (no tensors).
- :class:`LoraStore`       — async load/unload + GPU residency, in-flight refs.
- :class:`LoraBatchState`  — per-step row → adapter assignment; satisfies
                             :class:`arbi_serve.adapters.adapter.AdapterState`
                             via ``apply_to_output`` (the adapter
                             dispatcher entry).
- :func:`apply_lora`       — BGMV-style segment-GEMM correction (Triton + CPU).
- :func:`load_peft_adapter`— PEFT ``adapter_model.safetensors`` loader.

Punica/BGMV-style segment-GEMM for mixed-LoRA batches.

Given:

  - ``x``: ``(N_tokens, in_features)``  the linear's input
  - ``lora_indices``: ``(N_tokens,)`` int — one adapter id per token;
    ``0`` means "skip"
  - ``packed_a``: ``(K+1, max_rank, in_features)``  — adapter ``k`` ≥ 1
    holds its A^T (vLLM convention); slot 0 is zero
  - ``packed_b``: ``(K+1, out_features, max_rank)``  — adapter ``k`` ≥ 1
    holds B^T (vLLM convention); slot 0 is zero
  - ``scalings``: ``(K+1,)`` float — α/r per adapter (slot 0 = 0)

Compute, for each token ``t`` and ``k = lora_indices[t]``::

    y[t] += scalings[k] * (x[t] @ packed_a[k] ^ T) @ packed_b[k] ^ T

In one launch, all tokens proceed through the same kernel. Tokens
whose ``k = 0`` cost only a load + zero-multiply (no correction
applied) thanks to slot 0 being zero, which keeps the kernel
branch-free.

We ship two impls:

  - ``apply_lora_triton`` (default on CUDA) — Triton kernel launched
    per ``(N_tokens, BLOCK_OUT)`` grid; segment GEMV per token.
  - ``apply_lora_pytorch`` (CPU + correctness fallback) — vectorized
    via ``torch.einsum`` over slot dim. Used by tests on a CI box
    without a GPU and also as the apples-to-apples reference.

Both are bit-equivalent at fp32 (within rounding); the public entry
:func:`apply_lora` selects automatically by tensor device.

LoRA is one impl of the :class:`Adapter` Protocol.
The :func:`apply_lora` function below is the kernel-dispatch entry
that :meth:`LoraBatchState.apply_to_output` (the Protocol method)
calls into. The dispatch boundary is intentionally narrow — only
``y, x, layer_key`` cross the Protocol seam; per-adapter packed
tensors stay on the state object.

PyTorch fallback / correctness reference for mixed-LoRA batches.

Steps:
  1. ``Aw[t] = x[t] @ packed_a[k]^T``     → ``(N, max_rank)``
  2. ``By[t] = Aw[t] @ packed_b[k]^T``    → ``(N, out_features)``
  3. ``y[t] += scalings[k] * By[t]``

Each per-token slot lookup uses ``index_select`` along slot dim 0.

Public entry. Selects Triton (CUDA) or PyTorch (CPU) automatically.

Asserts:

  - shapes match the docstring above;
  - all tensors share device;
  - ``lora_indices`` is int (any int dtype OK; cast to int64
    internally for ``index_select``).

Bit-equivalence: the PyTorch reference is the spec. Triton
is fp32 accumulation throughout (matching), with bf16/fp16 only
at load/store boundaries — same numerics envelope as the rest of
the model's TP=1 forward.

Per-program: one (BLOCK_N tokens × BLOCK_O outputs) tile.

Each program reads its tokens' lora indices, gathers the right
A and B slices, computes ``y += scaling * (x A^T) B^T``.
Tokens where ``lora_idx == 0`` add zero (slot-0 weights are
zero-filled), so the kernel is branch-free.

Triton-launched mixed-LoRA segment GEMM.

Falls back to the PyTorch impl if the in-features dimension is
below a threshold where the launch overhead dominates — the
Triton path is only the right call for prefill-shaped batches.

Per-step :class:`LoraBatchState` builder.

Inputs:

  - ``lora_assignments``: ``list[Optional[str]]`` — one entry per row in
    the batch (None = no LoRA on that row).
  - ``cu_seqlens_q``: ``(B+1,)`` int — token-boundary csum from the
    :class:`ScheduledBatch`. Used to expand per-row → per-token.
  - ``store``: the engine's :class:`LoraStore` (resolves names →
    :class:`LoraAdapter`).
  - ``device``: where to place the per-step tensors.
  - ``num_layers``, ``model_target_keys``: precomputed once at engine
    build time. Each ``model_target_keys`` entry is the
    ``"layers.{i}.{module_name}"`` key the parallel linear was bound
    to via ``set_lora_target_name``.
  - ``capture_pool``: optional :class:`LoraCapturePool`. When set, the
    builder writes the per-step tensors into the pool's persistent
    buffers via :meth:`LoraCapturePool.populate` and returns a view
    state — required for cudagraph capture/replay parity (the
    captured kernels reference the persistent ``data_ptr``s). When
    ``None`` (CPU paths, unit tests), the builder allocates fresh
    per-step tensors.

Output: a fully-populated :class:`LoraBatchState`. Tokens for rows
with no LoRA carry ``lora_token_indices = 0`` (the no-op slot).

Build the per-step state. Returns ``None`` if no row has a LoRA.

The returned state's ``per_target_packed`` covers exactly the
intersection of (a) target keys on the model's linears, (b) target
modules present in at least one of the active adapters. Linears
whose key isn't in the dict short-circuit to base in the adapter
dispatcher (``LoraBatchState.apply_to_output`` returns ``y``
unchanged).

The PEFT loader places ``weight_a`` / ``weight_b`` at the full
dim, so this path supports TP=1 only. TP>1 would require slicing
``weight_b`` by output-dim shard for column-parallel targets and
``weight_a`` by input-dim shard for row-parallel targets — same
shard math as the base linear's loader; we raise instead.

When ``capture_pool`` is supplied, the build proceeds in two
phases: (1) build a transient :class:`LoraBatchState` with fresh
per-step tensors, then (2) ``copy_()`` it into the pool's
persistent buffers and return a *view* state slicing those
buffers. The view's tensors share ``data_ptr``s with the pool —
every step's view references the same addresses, which is what
cudagraph replay requires.

Persistent worst-case LoRA buffers for cudagraph capture.

The default :func:`arbi_serve.adapters.lora.builder.build_lora_batch_state`
allocates fresh ``torch.zeros`` per step (``lora_token_indices``,
per-target ``packed_a`` / ``packed_b``, ``scalings``) — every call
returns tensors with new ``data_ptr``s. A captured cudagraph that read
those data_ptrs at capture time would dereference stale memory at
replay.

Worst-case-sized buffers are allocated once at engine init, then
per-step content is ``copy_()``'d into them. The captured graph
references the persistent ``data_ptr``s; replay re-runs against
whatever the engine just copied in.

Pool key extension. The captured graph also bakes ``num_active_loras``
(the rank-0 dim of ``packed_a`` / ``packed_b``) into kernel launch
shape. We quantize to a small bucket set (e.g. ``{0, 1, 2, 4, 8}``) so
the pool key extends cleanly without exploding the captured-graph
count. The bucket is derived via ``bisect_left``: a step with 3 active
adapters maps to the ``num_active_loras=4`` bucket and replays the
graph captured at that bucket, using zero-padded adapter slots for the
unused tail slots.

``CapturedGraphPool`` keys its captured graphs on
``(batch_size, seq_len, lora_bucket, is_prefill)``; the
``lora_bucket`` slot is what :func:`quantize_active_loras` produces.
The :class:`LoraCapturePool` here is the LoRA-specific buffer manager
(adapter slots, not captured kernels) that those keys index into.

Round ``num_active`` UP to the smallest bucket in
:data:`ACTIVE_LORA_BUCKETS` that fits.

A step with no LoRA in flight (``num_active == 0``) maps to the
``0`` bucket — the captured graph there is the no-LoRA baseline
where every linear's adapter dispatcher short-circuits. Steps
with ``num_active`` exceeding the largest bucket raise; operator
must extend :data:`ACTIVE_LORA_BUCKETS` and re-capture.

Worst-case-sized persistent buffers + per-step :func:`populate`.

Sizing. ``packed_a`` is ``(max_loras + 1, max_rank, in_features)``
and ``packed_b`` is ``(max_loras + 1, out_features, max_rank)``,
both per target key. Slot 0 is the canonical "no LoRA" sentinel
(same convention as :class:`LoraBatchState`). ``+1`` covers it.

Per-step :func:`populate` zero-fills the staging slots
``[1 : num_active + 1]`` then ``copy_()``s the active adapters'
weights into them, and zero-fills the tail beyond. ``scalings``
and ``lora_token_indices`` get the same treatment.

The returned :class:`LoraBatchState` is a *view* — its tensors are
slices of the persistent buffers, so the BGMV kernel reads stable
``data_ptr``s + the ``shape`` axis matches the active-LoRA bucket
the graph was captured at.

Zero adapter slots ``[1 : num_active + 1]`` across every
per-target buffer + ``scalings``.

Unused tail slots (``num_active + 1 :``) keep whatever they
held last; the kernel never reads them because
``lora_token_indices`` is bounded by ``num_active``.

Copy ``state``'s content into the persistent buffers and
return a view :class:`LoraBatchState` referencing slices of
those persistent buffers.

``num_active_bucket`` is the quantized active-LoRA bucket the
captured graph was captured at (see
:func:`quantize_active_loras`). The returned view's
``packed_a`` / ``packed_b`` slot dim equals ``bucket + 1`` so
the BGMV kernel launch shape matches the captured graph
exactly. Real adapters occupy slots ``[1 : state.num_active +
1]``; the tail ``[state.num_active + 1 : bucket + 1]`` is
zero-filled (no contribution).

Return a view :class:`LoraBatchState` over the persistent
buffers — pre-zeroed except slot 1 holds a dummy unit-scaling
adapter to ensure the BGMV kernel branch fires for the captured
kernel.

At capture time the graph must launch the BGMV kernel against
the persistent buffers so replay's ``copy_()``-then-replay
contract produces the right values. We seed ``scalings[1:]``
with a tiny non-zero so the kernel traces the full path
(slot-0 zero would let the compiler fold the contribution
away — instrumented production kernels don't, but the seed
keeps the contract crisp). At replay time the engine
``copy_()``s real values in.

``num_active_bucket == 0`` returns a state with ``num_active
== 0`` — captured graphs at the 0-bucket are the canonical no-
LoRA shape and the BGMV kernel never fires (the per-target
dict is empty so the adapter dispatcher
``LoraBatchState.apply_to_output`` short-circuits).

LoRA adapter dataclasses.

A LoRA replaces a base linear ``y = x W^T`` with::

    y = x W^T + (α / r) * (x A^T) B^T

where ``A: (r, in_features)``, ``B: (out_features, r)``. PEFT names
these ``lora_A.weight`` and ``lora_B.weight``; vLLM stores them as
``weight_a`` / ``weight_b`` after a transpose. We follow the vLLM
convention because every BGMV-style kernel (Punica, vLLM's lora_kernel)
expects that orientation.

We keep the per-target-module slot keyed by the **same name the model
uses internally** (e.g. ``"q_proj"``, ``"k_proj"``, ``"v_proj"``,
``"o_proj"``, ``"gate_proj"``, ``"up_proj"``, ``"down_proj"``,
``"gate_up_proj"`` for the merged gate/up linear). The PEFT loader is
responsible for normalizing PEFT's ``"q_proj"`` name to the same key
the model declared on the base linear via :meth:`set_lora_target_name`.

One target module's (A, B) pair on a single layer.

Stored in fp16/bf16 (matches the base linear dtype). Tensor shapes
follow vLLM's BGMV-input convention:

  - ``weight_a``: ``(rank, in_features)``  — applied as ``x @ A^T``
  - ``weight_b``: ``(out_features, rank)`` — applied as ``(x A^T) @ B^T``

For ``MergedColumnParallelLinear`` (gate+up fused), ``weight_b``
spans both output halves stacked on dim 0 — same memory layout as
the merged base linear.

One loaded LoRA adapter.

``target_modules`` is keyed by ``"layers.{idx}.{module_name}"`` —
e.g. ``"layers.5.q_proj"``. The model's parallel-linear instances
look themselves up here when ``LoraBatchState`` is set on forward.
Models that fuse gate+up under one base linear key
(``MergedColumnParallelLinear``) expect a fused ``weight_b`` with
both shards stacked on dim 0; the PEFT loader concatenates.

Load PEFT-format LoRA adapters into :class:`LoraAdapter`.

PEFT layout (``adapter_config.json`` + ``adapter_model.safetensors``):

  - ``adapter_config.json``:
      ``r`` (rank), ``lora_alpha``, ``target_modules`` (list of names),
      ``lora_dropout`` (ignored at inference), ``base_model_name_or_path``.

  - ``adapter_model.safetensors``: tensors keyed
      ``base_model.model.model.layers.{i}.self_attn.q_proj.lora_A.weight``
      — the ``lora_A`` shape is ``(rank, in_features)``;
      ``lora_B`` shape is ``(out_features, rank)``. Some PEFT exports
      use ``base_model.model.<name>.lora_A.weight`` patterns; we accept
      both via a regex-style key matcher.

Validation. We reject (raise :class:`PeftValidationError`) at load
time if:

  - rank is not positive, or alpha is not positive;
  - any target module is not in the model's known target list;
  - any layer's ``in_features`` / ``out_features`` doesn't match the
    base linear's expected shape;
  - the safetensors file is missing for a declared target module;
  - the rank inferred from a tensor disagrees with the config.

PEFT itself is never imported at runtime; this module reads
``adapter_config.json`` and ``adapter_model.safetensors`` directly.

Slice a full-dim ``(weight_a, weight_b)`` pair to one TP rank.

Mirrors the base linear's shard (``models.linear``): a
**column-parallel** target (q/k/v/gate/up — output-sharded base)
shards the LoRA up-projection ``weight_b`` on the output dim (dim 0);
a **row-parallel** target (o/down — input-sharded base) shards the
down-projection ``weight_a`` on the input dim (dim 1). The
complementary matrix stays full.

Correctness rests on the apply-before-collective contract: LoRA is
applied to the local matmul output BEFORE the column-parallel
all-gather / row-parallel all-reduce (see ``ColumnParallelLinear`` /
``RowParallelLinear``), so each rank's sliced LoRA contribution
composes (column) or accumulates (row) with its peers exactly like
the base weight shard — no extra reduction is introduced.

``tp_size <= 1`` returns the inputs unchanged.

Load a PEFT adapter directory into a :class:`LoraAdapter`.

Arguments:
  - ``path``: directory holding ``adapter_config.json`` +
    ``adapter_model.safetensors``.
  - ``adapter_name``: how callers refer to it (the registry key).
  - ``device``, ``dtype``: where + how to materialize.
  - ``base_layout``: optional dict
    ``{"q_proj": (in_features, out_features), …}``. When set,
    load fails fast on shape mismatches.
  - ``num_layers``: optional — when set, every layer in
    ``[0, num_layers)`` must have at least one target weight or
    the load is rejected (catches partial-export adapters).
  - ``validator``: optional ``f(config, base_layout) -> None``
    for engine-level cross-checks (raises if mismatched).
  - ``max_rank``: optional engine-side ceiling. Adapters whose
    config rank exceeds this are rejected. Sized to match the
    engine's :class:`LoraCapturePool` so a captured cudagraph's
    rank-axis kernel-launch shape always fits the loaded
    adapter.
  - ``tp_rank`` / ``tp_size`` / ``tp_shard_plan``: tensor-parallel
    sharding. The adapter file is always full-dim; shape validation
    runs against the (full) ``base_layout`` first, then each target
    is sliced to this rank via :func:`shard_lora_target` using
    ``tp_shard_plan[module] -> "column" | "row"`` (built from the
    model's linears). ``tp_size <= 1`` (the default) is a no-op, so
    the single-rank path is unchanged.

The runtime never imports the ``peft`` package — we just respect
its on-disk format.

Per-step LoRA batch state.

The engine builds one :class:`LoraBatchState` per forward step from
``ScheduledBatch.lora_assignments`` (one row per request). Linear
layers receive it via the keyword arg ``lora_state=``; ``None`` means
"no LoRA at this step" — every parallel linear short-circuits to the
base ``F.linear`` call.

Per-token row indices. Each request's tokens span a contiguous range
``[cu_seqlens_q[i], cu_seqlens_q[i+1])`` of the flat token tensor.
We expand the per-row LoRA assignment into a per-token ``lora_indices``
tensor at state-build time; the BGMV kernel reads one ``lora_id`` per
token. ``-1`` is the "no adapter" sentinel.

LoRA is one impl of the :class:`Adapter` Protocol. The state object
both carries the per-step BGMV-kernel inputs (the linear's
``lora_state=`` kwarg contract) AND satisfies
:class:`arbi_serve.adapters.adapter.AdapterState` so the generic
adapter dispatcher in :mod:`arbi_serve.models.linear` can route through
:meth:`apply_to_output`. The ``lora_state=`` kwarg is accepted directly
by linear forwards as an alternative to the generic dispatch path.

Inputs the BGMV kernel needs to apply mixed-LoRA in one call.

Fields:

  - ``adapters``: index → adapter. ``adapters[0]`` is the canonical
    "base only" sentinel (an empty adapter); real adapters start at
    index 1. The convention keeps ``-1`` available as the "this row
    has no adapter" runtime signal in user-facing tensors, but
    internally we shift to a 0-base scheme (``0`` == no LoRA) so
    ``index_select`` works without masking. The kernel receives
    ``lora_indices`` shifted (0 == no LoRA, k == adapters[k]).

  - ``lora_token_indices``: ``(N_tokens,)`` int32. Per-token
    adapter index. ``0`` means "skip" (apply no correction).

  - ``per_target_weights``: dict keyed by the target-module name
    (the same key the model passes to ``forward(..., lora_state)``,
    e.g. ``"layers.5.q_proj"``). Value is a precomputed
    ``(num_active_adapters + 1, rank, in_features)`` /
    ``(num_active_adapters + 1, out_features, rank)`` pair. Slot 0
    is zero-filled (the no-LoRA path, never read by a token whose
    ``lora_token_indices`` is 0). Slot ``k`` holds adapter ``k``'s
    weights for that target.

  - ``scalings``: ``(num_active_adapters + 1,)`` float — α/r per
    adapter (slot 0 == 0).

Ranks may differ across adapters; we right-pad each adapter's
``weight_a`` / ``weight_b`` to the max rank in the batch
(``max_rank``). Padding rows / columns are zero so the
contribution is bit-exact zero. This single-rank convention keeps
the kernel a single launch — the alternative (per-adapter ranks)
is significantly more complex for marginal memory savings on a
handful of in-flight adapters.

``False`` short-circuits the adapter dispatcher.

We're "active" iff at least one row carried a LoRA assignment,
which is exactly when ``num_active > 0``. The
:func:`build_lora_batch_state` builder already returns ``None``
for a fully-base batch, so an instantiated state with
``num_active == 0`` is uncommon but well-defined.

Adapter Protocol entry point. Adds the BGMV correction in place.

Applies the BGMV correction for the given linear. If this state
has no per-target weights for
``layer_key``, the correction is the identity (return ``y``).
Otherwise dispatch to the BGMV kernel; ``y`` is updated
in-place and returned for call-chain ergonomics.

``y`` and ``x`` may be 3-D (when the model passes a
``[B, T, F]`` tensor through a merged linear); we reshape to
``[N, F]`` for the kernel and rely on the in-place update to
keep the upstream view consistent.

Server-side LoRA registry: load / unload / list / refcount.

The public surface is async; CPU-bound safetensors parsing runs via
``asyncio.to_thread`` so the engine event loop is never blocked.

  - ``load(name, path)`` reads ``adapter_config.json`` +
    ``adapter_model.safetensors``, validates against the base model's
    expected shapes, uploads to GPU, returns a populated
    :class:`LoraAdapter`.
  - ``unload(name)`` raises :class:`InFlightError` (HTTP 409) when the
    adapter has live references.
  - ``acquire(name)`` / ``release(name)`` form the in-flight refcount;
    the engine bumps it when a request enters scheduling and drops it
    when the request finishes.
  - ``list_loaded()`` returns :class:`LoraInfo` snapshots for
    ``GET /v1/loras``.

Every loaded adapter stays resident in VRAM; the store reports total
bytes via :meth:`gpu_bytes_total` for the admin endpoint.

Estimate adapter GPU bytes from ``adapter_config.json`` + base layout.

Returns ``0`` when the config can't be read or ``base_layout`` /
``num_layers`` are missing — admission then degrades to "no
estimate, allow through".

Per-target bytes per layer:
    ``rank * in_features + out_features * rank``  weights × dtype.

Summed across every config-declared ``target_modules`` entry that
appears in ``base_layout`` and across all ``num_layers`` layers.

Async-safe registry of loaded LoRA adapters.

Threading model. The store runs on the main asyncio loop. The only
blocking operation is :meth:`load`, which offloads the safetensors
parse + GPU copy to a worker thread. Unload, list, get, acquire,
release are O(1) and never block.

Refcounting. ``acquire(name)`` is called by the engine when a
request that targets that LoRA is scheduled into a step. Release
happens when the request finishes (or is cancelled / errored).
The refcount itself lives under an asyncio lock; reads via
``in_flight(name)`` are race-free at single-event-loop granularity.

Lifecycle vs model reload. The store is built per engine (build.py
``lora_target_key``
— it does NOT alias the model's linears, so a forward looks adapters
up by key each step. A model SWAP/RELOAD that rebuilds the engine
rebuilds the store too (adapters must be re-loaded). A reload that
keeps the store but changes a target's dims (or its TP shard, see
``peft_loader.shard_lora_target``) would leave stale-shaped adapters;
the load-time ``base_layout`` check catches a mismatch only on the
next load, so callers that hot-swap the model must drop + reload
adapters. There is no in-place reshape.

``validator`` is an optional callable
``f(adapter_config, base_layout) -> None`` that the engine
passes to assert that the LoRA's ``target_modules`` and
``in_features`` line up with the base model. None disables
the cross-check (used by unit tests that load adapters
without a base model).

``mem_pool`` is the engine-supplied named ``MemPool`` for
adapter weights — when set, every ``load`` routes its GPU
allocations through that pool. Adapter weights are static
for an adapter's lifetime, so isolating them from per-step
transient allocs prevents the unload-then-load churn from
fragmenting the default heap.

``max_rank`` is the engine's :class:`LoraCapturePool`
rank-axis ceiling. ``load`` refuses adapters whose
configured ``r`` exceeds this; the captured cudagraph at
replay time would otherwise read past its persistent
``packed_a`` / ``packed_b`` buffers. ``None`` disables the
check (used by unit tests).

Raise :class:`LoraAdmissionError` when loading would push
the ``lora_pool`` past its configured ``size_target_bytes``.

Returns silently when ``mem_pool`` is None (no pool wired —
legacy code path), when ``size_target_bytes <= 0`` (no cap
configured), or when we can't estimate expected bytes
(missing ``base_layout`` / ``num_layers`` / unreadable
``adapter_config.json``).

Load + register a LoRA adapter from a PEFT directory.

Raises:
  - :class:`LoraAlreadyLoadedError` if ``name`` is taken;
  - :class:`LoraAdmissionError` if the adapter would push the
    engine's ``lora_pool`` past its ``size_target_bytes`` cap;
  - :class:`PeftValidationError` if the adapter shape doesn't
    match ``base_layout`` (or its config is malformed).

The load itself runs in a worker thread (CPU-bound: file I/O +
safetensors parsing). The GPU copy + tensor materialization
happens at the end on the main thread to keep the CUDA stream
with the engine.

Bump the in-flight refcount for ``name``.

Engine calls this when a request enters scheduling. The store
does NOT block on a missing name — that's a programming bug
(the engine should have validated at submit time). We raise
loudly so it shows up in tests.

Standalone admin-console service — NOT part of any served engine process.

Serves the admin UI (``arbi_serve/server/admin_ui/index.html``, the same
file a lone arbi-serve instance would otherwise host) and proxies every
call the page makes to an arbi-serve instance, local or remote. One console
instance can therefore manage a whole fleet — it never runs inside the
process being managed, and it never assumes the target is on the same host.

WHICH instance a call goes to is chosen per request, from a server-side
list (:mod:`arbi_serve.admin_console.registry`): the page names an ENTRY,
never an address. Entry zero is the engine this console shares a network
namespace with, derived from the deployment and verified before anything
is served (:mod:`arbi_serve.admin_console.local_engine`); the rest are
added at runtime, each with its own credential, which the console sends
server-side and never renders back to the page.

Two clients, one per protocol, deliberately not merged into a single
hand-rolled HTTP layer:

  * :mod:`openai` (``AsyncOpenAI``) for the OpenAI-compatible completions/
    chat surface — the purpose-built, maintained client for that exact
    wire protocol (the one every OpenAI-compatible server, vLLM included,
    tells its users to use).
  * :mod:`arbi_serve_client` (this repo's own generated client,
    ``client/arbi_serve_client/`` — see ``scripts/generate-client.sh``)
    for the admin surface, used directly via its generated request
    functions. No hand-written SDK layer on top: the admin routes don't
    declare a typed ``response_model`` (they return a plain dict), so the
    generated ``.parsed`` convenience always comes back ``None`` on
    success — ``targets._call_admin`` reads ``response.content`` off the
    generated ``*_detailed`` call's ``Response`` object instead, which is
    still just consuming the generated client's httpx transport, not a
    parallel client.

Run standalone: ``python -m arbi_serve.admin_console``.

``python -m arbi_serve.admin_console`` — run the standalone admin console.

Deliberately a separate process/port from any arbi-serve instance it
manages (see the package docstring) — this is the control plane, not a
managed node.

Return ``(keyfile, certfile)``, generating a self-signed pair on
first call and reusing it on every later one — regenerating on every
``docker restart`` would burn every browser's "I trust this cert"
decision each boot, right back to the raw click-through warning this
exists to make a one-time cost.

Shells out to the system ``openssl`` (present on essentially every
Linux base image already; the alternative, depending on the
``cryptography`` PyPI package, would add a real dependency to a
console whose whole point is being the lightweight control-plane
tool) rather than reimplementing X.509 generation.

Why this needs to exist at all: browsers only expose
``navigator.mediaDevices``/``getUserMedia`` (what the mic-capture Audio
tab needs) in a "secure context" — https://, or the special-cased
localhost/127.0.0.1. This console is normally opened against a remote
GPU box's bare IP over plain http://, which is neither, so the mic API
silently does not exist. A self-signed cert is not "less secure than
real TLS would be" for THIS purpose — the browser's secure-context
check only cares that the transport is encrypted, not that the CA is
publicly trusted — so it fully satisfies the requirement, at the cost
of one manual "proceed anyway" click per browser on first visit.

Native MTP vs DFlash2 as one switch, and the proof that it took effect.

The two speculative-decode configurations this server ships are selected
by three params that must move together — the drafter checkpoint, the
MTP master switch and the speculation depth K. Setting one without the
others is not a half-switch, it is a different configuration: clearing
``dflash_draft_path`` while leaving K at the DFlash depth boots the
bundled head at a depth it was never calibrated for. They are therefore
applied as one payload and named as one mode.

Why the receipt is a COUNTER FAMILY and not a flag read-back
    Every param here is capture-affecting: the engine builds and captures
    a new pool member and swaps to it, and until that swap completes the
    registry already reports the new value. A read-back is therefore
    evidence that the request was accepted, never that the drafter
    changed — the exact substitution that let two prefill flags report as
    applied while a thousand pinned calls ran with the pin off. What
    distinguishes the two modes is which code fires: the DFlash driver
    and its DSpark selector own a counter family no bundled-head chain
    can touch, so the family moving IS the drafter, and the family
    staying still while the drafter demonstrably ran is the bundled head.

Two independent instruments, because one can be silently wrong
    ``did_speculate`` reads committed tokens per step off the served
    usage: a decode with no drafter commits exactly one token per step,
    so anything above one proves speculation happened at all. The counter
    family then says WHICH drafter. A mode is confirmed only when both
    agree; when they disagree the switch is reported unconfirmed with the
    disagreement stated, because a plausible number from an unproven
    configuration is worse than no number.

A zero outside a counter's firing window is not a zero
    A GRAPH_RECORD counter fires at capture and cannot fire at replay, so
    reading one during serving reads zero by construction. Those counters
    are excluded from the serving-window verdict and reported separately
    rather than counted as silence — an uninterpretable zero must never
    be allowed to read as "the path did not run".

One speculative-decode configuration, as the operator picks it.

``overrides`` is the whole payload: every param that has to move for
the mode to mean what its name says. ``expects_dflash`` is what the
counter family must show for the switch to be confirmed.

Mean committed tokens per decode tick, and WHICH accounting produced it.

Stated as its own quantity because it is not the accept rate, and the
two have already been read for each other: accept rate is
``accepted/proposed``, this is tokens per tick, and its floor is 1.0
because every tick commits its verified token whether or not a draft
survived.

Two accountings exist and they do not agree, so the source travels
with the number. The EFFECTIVE one — ``mtp_row_tokens /
mtp_row_ticks`` — counts every decode row-tick including the cold
ones, and is the figure that determines TPOT. It is bumped on the
SPMD commit path, so a TP1 boot leaves it flat and the per-request
ratio stands in. That fallback is WARM-ONLY and reads HIGH: a cold
row drafts nothing and commits one token, consuming a whole tick
while adding to neither counter. Presenting it unlabelled as the
effective number is the conflation this return shape prevents.

Whether ANY drafter contributed, independent of the counters.

A decode with no speculation commits exactly one token per step, so a
figure above the floor is proof a drafter ran without reading a
single counter. This is the second instrument the verdict needs: a
counter family can be miswired, and agreement between two unrelated
readings is what makes the receipt worth having.

Counters that moved, split by whether their window makes them readable.

``rows`` carries the per-counter contract metadata the flag-truth
endpoint reports alongside the raw counts; the declared window is
read from it. A counter whose window is GRAPH_RECORD is sorted out of
the serving verdict entirely — its zero is guaranteed by construction
and can never be disconfirmed by more traffic.

Did the switch to mode ``key`` actually take effect?

Confirmation needs BOTH instruments to agree: the tokens-per-step
figure has to show a drafter ran at all, and the DFlash counter
family has to be moving for DFlash2 and still for native MTP. Any
other combination returns ``confirmed: False`` with the specific
disagreement, because the whole point of the receipt is that an
unproven configuration must not be allowed to present a plausible
number.

The config-override payload that puts the target in this mode.

``draft_path`` is required by, and only by, the DFlash mode: the
checkpoint is deployment state, not a constant, so it is read off
the target or supplied by the operator rather than baked here.

The fidelity half of the console A/B: KL divergence and argmax flips.

A speed-only A/B lets a numerics regression ship looking like a win, so a
console result is four numbers and not two: prefill rate, decode rate, KL
against the baseline, and the count of positions where the arm's top-1
token differs from the baseline's. The first two describe the path we
serve; these two describe whether it still computes the same thing.

Why the fidelity half uses a DIFFERENT endpoint than the speed half
    Speed is measured over ``/v1/chat/completions`` because that is the
    path an operator actually runs, sampler and all. KL cannot be
    measured there: each arm would follow its own continuation, and
    differencing two different texts is not a divergence. Fidelity is
    therefore measured over ``/v1/completions`` with ``echo`` and
    ``prompt_logprobs``, which scores a FIXED token sequence
    teacher-forced — the same weights, the same boot, the same flags,
    one request per prompt, positions aligned by construction.

Scored over a PINNED window, never over every aligned position
    Position 0 has nothing predicting it and the first few positions
    carry the fresh-diagonal step rather than the steady plateau;
    averaging them into the same statistic answers a different question
    at every prompt length. The window is stated with the number, and a
    row too short to fill it is dropped rather than scored short.

The bound is a lower bound, and its slack is reported with it
    :func:`kl_truncated` mirrors turbo-attn's
    ``benchmarks/eval/divergence/metrics.py`` so a console figure is
    comparable with every receipt already recorded there. Both sides
    report only their own top-K, so exact KL is unavailable: a reference
    id missing from the arm's top-K gets the tightest available upper
    bound on its arm probability, which MINIMISES the log ratio. The
    result therefore understates drift, and ``missing_mass`` — the
    reference mass that had to be imputed — travels with every KL figure
    so a reader can see when the bound stopped being tight.

A null arm that disagrees with itself invalidates the run
    The served sampler is seeded and the scoring path is teacher-forced,
    so re-scoring the baseline against itself must produce exactly zero
    flips and zero KL. When it does not, the harness is measuring
    something other than the flag, and :func:`null_control` says so as a
    refusal — every arm behind it is uninterpretable, not merely noisy.

``choices[].prompt_logprobs`` as ``{ids, probs}`` rows, one per position.

The wire format is a map from token id (as a string key) to an entry
carrying ``logprob``. Probabilities are exponentiated here once so
every consumer scores the same numbers, and a position the server
reported as ``null`` stays ``None`` rather than becoming an empty
distribution — absence of data must never read as agreement.

The argmax token id at each position, or ``None`` where unscored.

The established spelling for the argmax column, so a console run and
a divergence receipt can be compared without a translation step.

Truncated ``KL(P_ref || P_arm)`` at ONE position, over the reference's top-K.

A LOWER bound, and deliberately not a flattering one. The sum runs
over the reference's top-K ids only, so reference mass outside them
is missed — returned as ``support_mass``. A reference id absent from
the arm's top-K has an unknown arm probability bounded above both by
the arm's smallest reported probability and by its unreported
residual mass; substituting that upper bound minimises the log ratio,
which is what makes the result a lower bound. ``missing_mass`` is the
reference mass that needed imputing and MUST be reported with the KL:
on a peaked distribution it is ~0 and the bound is tight, and when it
is not ~0 the number is not trustworthy.

``kl`` is NaN when the reference side is empty, so an unscored
position cannot read as agreement.

The pinned window of ``rows``, or ``None`` when it does not fill.

A short row is refused rather than scored over whatever it has: an
average taken over 12 positions and one taken over 95 are different
statistics, and silently returning the first under the second's name
is how a window stops meaning anything.

Positions where the arm's top-1 token differs from the reference's.

Only positions BOTH sides scored are compared, and that count is
returned as the denominator: a flip rate is meaningless without the
number of comparisons behind it, and a position one side left
unscored is not agreement.

One prompt's KL and argmax flips over the pinned window.

Returns ``scored: False`` with a stated reason whenever the window
could not be filled on both sides — a run that could not score has
to say so, because a zero KL and a zero flip count are exactly what
perfect agreement looks like.

Whether the harness reproduces the baseline against itself.

The scoring path is teacher-forced over a fixed sequence, so a
baseline re-scored against the baseline must read exactly zero flips
and zero KL. Anything else means the measurement is not attributable
to any flag, so this returns a REFUSAL and not a small number: the
right response is to fix the harness, never to subtract the offset
and carry on.

The measurement method behind the console's one-click flag A/B.

Each decision here exists because the naive version of it measures
something other than the flag.

Canonical prompts, fixed content
    Both arms decode the SAME prompts, built by
    :func:`~arbi_serve.server.playground_workloads.typed_prompt` at a
    fixed seed, in the same order. Drafter acceptance is a property of
    the text and decode throughput moves with acceptance, so a prompt
    drawn afresh per request puts the workload's own variance inside the
    delta. The sampler is left exactly as served — stochastic, thinking
    regime on — which is why a verdict is built from agreement across
    reps rather than from one delta.

A prefill target stated in TOKENS
    The word budget is bisected against the target's own ``POST
    /tokenize``, which renders the chat template the chat route renders.
    A word count is not a token count, and shipping one as the other
    names a prefill after a number the engine never sees.

TPOT divided by TOKENS, never by chunks
    A streamed chunk carries the text of ONE engine wakeup, so
    speculative decoding puts a whole accepted step in one chunk and an
    off-loop detokenizer coalesces several steps into one. Counting
    chunks therefore measures the grouping, which the flags under test
    move, instead of the decode rate. The probe asks the target for the
    tokens each chunk carried
    (``stream_options.include_token_count``) and
    :func:`steady_tpot_ms` divides by those.

TTFT and steady TPOT from the SAME request
    Decode after a deep prefill attends over the whole KV, so a decode
    rate measured behind a short prompt answers a different question.
    The leading streamed tokens are dropped so the ramp out of prefill
    does not sit inside the steady-state mean.

EOS respected
    ``max_tokens`` is a safety bound, never a target. Past the natural
    stop the distribution is degenerate and drafter acceptance stops
    resembling the served one.

Aggregate rate DERIVED, not timed
    Aggregate output rate is not an independent quantity. For one
    request it is ``n / (TTFT + (n - 1) * TPOT)``; for a rung it is that
    summed over the concurrent streams. TPOT is measured AT that
    concurrency, so batching contention is already inside it and nothing
    extra is assumed. Deriving the aggregate gives it the precision of
    TTFT and TPOT instead of the wall clock's, which also absorbs queue
    wait, scheduler stalls and stragglers. The wall-clock aggregate is
    reported beside it: when the two disagree, time is going somewhere
    neither TTFT nor steady-state TPOT accounts for, and that gap is
    itself the finding rather than something to average away.

One committed baseline, reused
    The default schedule measures arm A ONCE per baseline configuration
    — the flag set, the model and the prefill depth — commits it to the
    store, and then measures only arm B for every flag that shares it.
    A sweep pays for one A arm instead of one per flag.

    The alternative, :data:`SCHEDULE_ABBA`, alternates which arm runs
    first inside every rep. It cancels the within-rep ordering bias that
    a plain A-then-B schedule folds into arm B, at the cost of measuring
    both arms every time. It stays available behind an explicit opt-in,
    for a flag whose effect is small enough that the ordering term is
    worth paying to cancel, and every result records which schedule
    produced it — a reader cannot otherwise tell whether drift was
    cancelled or merely watched for.

Drift is watched for, since a reused baseline cancels none
    A committed baseline is a measurement from another time.
    :func:`drift_verdict` re-measures it and calls DRIFTED when it has
    moved by more than its own spread; every run that reused it since
    the last clean verification is then marked suspect rather than left
    standing. A baseline whose model, engine build or admission width
    differs from the live target is refused outright, never adjusted.

A verdict needs spread on BOTH sides
    :func:`compare_arms` calls a WIN or a LOSS only when the separation
    between the arms clears the spread of each of them, and refuses to
    call anything at all when either side has a single rep — one number
    against one number has no spread to clear, and reading it as a
    result is exactly the false confidence this harness exists to
    prevent. The paired schedule keeps :func:`verdict_for`, which adds
    the sign-agreement bar that pairing makes available.

The store's workload key for a run at this prefill depth.

The depth is part of the key because a baseline measured at one
prefill says nothing about another: prefill depth changes both the
TTFT it is compared on and the KV every decode step attends over.

One, half the admission width, and all of it.

Resolved against the target's live ``max_batch`` so the ladder tracks
the scheduler the server was actually sized for. The middle rung
floors at two — a middle rung of one would silently duplicate the
first and be reported as a batch. With no readable width the ladder
is a single c1 rung, and the caller says so rather than inventing a
batch the target may not admit.

The prefill depths this target can actually serve, and which it cannot.

A depth past the live ``max_context`` is not measurable, and quietly
dropping it would report a two-rung ladder under a three-rung name.
The unreachable depths are returned beside the reachable ones so the
caller states what it could not measure instead of presenting a
shorter ladder as the whole one. ``headroom`` keeps room for the
decode the rung still has to run.

Per-request prefill for a cell: the rung's TOTAL, split evenly.

The rung names an amount of prefill work, and concurrency says in how
many pieces it arrives — c1 sends one request of the whole total, cmax
sends ``total / cmax`` each. Holding the total fixed is what makes a
row comparable: the server does the same prefill work at every rung,
so a difference along the row is the batch SHAPE and not the amount.

Giving every stream the full depth instead would multiply the work by
the concurrency, ask the KV pool for more resident context than it
holds at the deep rungs, and leave the widest cell measuring queue
wait.

Which (total prefill x concurrency) cells the KV pool can hold.

Because the rung's prefill is a total, the prompt side costs the same
at every concurrency and only the decode side grows with it — each
stream needs room for what it generates. A cell is skipped by name,
with the arithmetic, rather than run into a pool that would serialise
it and report the queue wait as latency.

An unreadable capacity measures every cell: a missing number is not a
capacity of zero, and refusing the whole ladder on it would be a worse
answer than trying.

Which arm runs first in rep ``rep`` (zero-based): ABBA, not AB always.

Only the :data:`SCHEDULE_ABBA` path uses this. The default schedule
measures one committed baseline and then arm B alone, so there is no
ordering inside a rep to counterbalance.

Trim an over-long prompt's BODY until it hits ``target`` exactly.

A word budget cannot land on a token count: one word is under two
tokens, so the closest reachable count straddles the target and a
ladder rung named "1k" arrives as 1051. Bisecting on characters
inside the body closes that gap, and the body is what shrinks so the
lead and the closing marker stay intact — both arms must still decode
the same shape of prompt, not a truncated one.

Returns the best text, its count, and the character cut that produced
it (``None`` for the untrimmed text) -- the cut is what
:mod:`~arbi_serve.server.playground_prompt_sizes` freezes, so that the
same prompt can be rebuilt later without repeating the search. A token
can span several characters, so the target may sit between two
reachable counts; the closest is returned and the caller reports it as
off-target rather than letting the rung's name stand in for the count
it got.

Grow prompt ``idx`` until its chat-templated token count hits the target.

``count_tokens`` is awaited with candidate prompt text and returns the
count the target's own tokenize surface reports for it, so the search
converges on the prefill the engine will actually see.

Returns the prompt, the token count reached, the word budget that
produced it and the target — the achieved count is reported rather
than assumed, because a finite corpus can run out before the target.

Rebuild a prompt from what the bisection found, without searching.

The inverse of :func:`canonical_prompt`'s result and the reason the
frozen table can be a handful of numbers instead of a megabyte of
text: the builder is deterministic, so the word budget and the
refinement's character cut are the whole answer. Byte-for-byte -- a
test rebuilds every frozen entry and compares, because a rebuild that
drifted would size a prompt to a count nothing had measured.

Prompt ``idx`` from the frozen table, or ``None`` if it is not in it.

Same arguments and same result shape as :func:`canonical_prompt`, and
no tokenize call at all: this is the catalogue's launch path. ``None``
means nobody froze this size -- a custom depth, a moved seed, a
reworded task -- and the caller measures it live instead.

Steady-state time per output TOKEN from a stream's chunk timestamps.

``stamps`` are monotonic seconds, one per streamed chunk, and
``tokens_per_stamp`` is how many completion tokens each of those
chunks carried (``stream_options.include_token_count``). It defaults
to one per stamp, which is what a chunk carries only when nothing
batches tokens behind it.

A chunk is NOT a token and is not reliably a step either. The engine
stages the text of one wakeup, so speculative decoding puts every
token a step accepted into one chunk, and an off-loop detokenizer
coalesces several steps into one. Both make a chunk denominator
measure something that moves with the flags under test rather than
with decode speed.

The window opens at a chunk boundary and the denominator is the
tokens delivered STRICTLY AFTER it: those are the tokens the elapsed
time actually bought, and counting them makes the result identical
however the same tokens were grouped into chunks. With one token per
chunk it is the interval count, unchanged.

The ramp out of prefill is skipped in TOKENS: whole leading chunks
are dropped until ``skip`` tokens are behind the window, because a
chunk cannot be cut at a timestamp it does not have. A stream too
short to spare the ramp keeps everything after the first chunk, which
is still a decode window and never includes prefill.

One request's output rate, derived from its own measured latencies.

``n / (TTFT + (n - 1) * TPOT)``. A single-token completion has no
decode window, so its rate is the first token's alone.

Prefill rate for a whole rung: ALL its prompt tokens over the time
until every stream has produced a first token.

Aggregate, not per stream, because the ladder holds the rung's total
prefill fixed while the concurrency splits it — so this is the number
that is comparable along a row, and a per-stream rate would fall with
concurrency purely because each stream carries a smaller share.

The denominator is the slowest TTFT in the rung, which carries queue
wait and the first decode step as well as the prefill. That is what a
client experiences, and it is stated rather than called a kernel
rate: an isolated prefill number needs the engine's own span, not a
client's clock.

Decode rate for ONE stream, from its steady-state time per token.

Per stream, not per rung: at concurrency N the rung's rate is this
summed over the streams, and the two must never be confused. The
ramp out of prefill is already excluded by the TPOT window.

Fold one rung's per-request measurements into the rung's row.

Reports BOTH aggregates and the gap between them: ``tps_agg_derived``
summed from each request's own TTFT/TPOT/token count, and
``tps_agg_wall`` from the rung's wall clock. A notable gap means the
rung spent time outside first-token latency and steady decode — queue
wait, a scheduler stall, a straggler.

The gap is only computed when EVERY served request contributed a
derived rate. Comparing a derived aggregate over some of the streams
against a wall-clock aggregate over all of them would report a
shortfall of coverage as a shortfall of time.

Call a metric WIN, LOSS or NOISE from its per-rep deltas.

A single delta is never a verdict. The sign has to hold across a
clear majority of the reps AND the mean has to clear the spread of
those same reps; anything else is reported as NOISE with the numbers
that made it one, so the caller can add reps instead of reading a
number that is not there.

Pair each rep's two arms on one metric and hand back its verdict.

Pairing is per rep, not arm-mean against arm-mean: the two arms of a
rep ran adjacent under the same thermal and background conditions,
which is what the ABBA schedule is for.

Mean, spread and count over the values that are present.

Spread is the population standard deviation of the reps actually
measured — this reports the spread that was seen, it does not infer a
population from it.

Compare arm B's reps against a committed baseline's reps.

The two samples are independent — the baseline was measured at
another time, possibly for another flag — so there is no per-rep
pairing to exploit and no sign agreement to count. What stands in
for it is separation: a WIN or a LOSS requires the gap between the
two means to exceed the spread of EACH sample, so a difference has
to be larger than the noise on either side of it.

A sample with fewer than two reps has no spread, which would let any
gap clear a zero bar. That case is refused outright rather than
called: one number against one number is not a comparison.

Whether this run should re-measure the committed baseline.

Drift is what a reused baseline cannot cancel, so it is watched for
on a cadence: every ``every`` reuses the next run pays for one more A
arm and compares it against the committed one.

Has the committed baseline moved since it was recorded?

Each metric is DRIFTED when the re-measured mean sits further from
the committed mean than the committed sample's own spread — the
baseline has moved by more than it ever varied. Any metric drifting
makes the whole baseline drifted, because a result compared against
it was compared against a number that has stopped describing the
target.

Which flag-truth counters moved across an arm, and by how much.

Only counters that actually moved are returned: an arm's receipt is
the set of contract paths that fired under it, not the whole
registry.

Whether the flag's own declared contract counter fired under an arm.

A flag that is engaged, whose declared counter shows no fires and a
recorded refusal, did nothing — and the refusal text says why. That
is the receipt a measured delta needs before it is attributed to the
flag.

How one validated flag A/B is folded, named and written down.

The console's flag ablation is ONE click and one server-side run. Every
element of the method lives in :mod:`arbi_serve.admin_console.ab_probe`,
which documents why each is there, and the run that drives a live target
through it lives in :mod:`arbi_serve.admin_console.measure`. This module
is the shapes in between: which rungs a ladder runs, what each reported
number is divided by, and the record a run leaves on disk.

The default schedule commits ONE arm-A baseline per (flag set, model,
prefill depth) and then measures arm B alone for every flag that shares
it, flipping back to the original value at the end and nowhere else. The
paired ABBA sweep stays available behind an explicit opt-in. Reusing a
baseline cancels no drift, so drift is watched for instead: the baseline
is re-measured on a cadence and, when it has moved by more than its own
spread, every run that reused it since the last clean check is marked
suspect — which is why the record shapes below carry the baseline
reference and the drift verdict alongside the numbers.

A comma list of rungs, filtered to the ones the target offered.

A rung the caller invented is dropped rather than attempted: the
allowed list was resolved against the live admission width and
context, and honouring a value outside it would run a ladder the box
cannot serve and report the failure as a measurement.

One rung's four speed readings, each with its denominator named.

``tps_decode_per_stream`` and ``tps_decode_agg`` are different
quantities at every concurrency above one, so both are reported
rather than one of them being called "the" decode rate.

The per-rep, per-rung rows a committed baseline recorded.

``None`` when the file cannot stand in for an A arm on this ladder:
a baseline measured over different rungs is not comparable rung for
rung, and quietly lining the two up by index would compare a c1
reading against a full-batch one.

The method a stored record was measured under.

Recorded on every file this harness writes, because a measurement whose
method is not stored beside it cannot be compared to anything later —
and because a reader has to be able to tell whether drift was cancelled
by pairing or only watched for.

The store payload for a COMMITTED baseline: arm A, on its own.

It is its own file rather than a column inside some flag's comparison,
so every later flag that shares this baseline configuration reuses one
measurement instead of digging an A arm out of another flag's run.
``detail.arm_reps`` keeps the per-rep, per-rung rows, which is what lets
the reuse be a real sample with its own spread rather than a mean.

The store payload for a drift check on a committed baseline.

Written whether or not it found drift: a clean check is what opens the
next reuse window, and a check that is not recorded cannot open one.

The store payload for one flag's comparison.

``cells`` carry every rep of both arms at every rung; ``detail`` carries
where the A arm came from and, when one ran, the drift check on it — a
result read against a reused baseline has to travel with the evidence
that the baseline still described the target.

FastAPI app for the standalone admin console. See the package docstring.

This file is the composition root and nothing else: it owns the app object,
the order the routers are included in, and the static mount. Every concern
lives in its own module so it stays independently testable and so a reader
looking for one surface does not have to page past the others.

Establish what this console is, then say it in one line.

The checks and their reasoning live in
:mod:`arbi_serve.admin_console.local_engine`; they run here because
uvicorn installs its logging config after ``__main__`` returns, so a
line emitted earlier goes nowhere. They run as a TASK rather than
inline because the engine beside this console is usually still
starting: the console holds every request behind
:func:`_refuse_until_verified` while it waits, and starts serving the
moment the engine answers, without a restart.

The watchdog started beside them is the same first check, kept running:
the engine restarting strands this console in a dead network namespace
with no signal anywhere, and exiting is the only thing that re-attaches
it (see :class:`local_engine.EngineWatchdog`). It arms itself only in the
deployment that has that failure and states which, so it is never
silently absent.

The refusal, as a page an operator can act on.

Every word an operator needs is IN the page: which check failed, which
address or port it tried, and what changes the answer. A console that
refuses without saying which of its three preconditions failed sends
someone to read compose files.

Serve nothing until this console knows it is the right console.

The alternative is what shipped before: a console that starts, renders
a full UI, and is pointed at nothing / unreachable by anyone / built
from another tree. Each of those looked like a working console. None of
them said anything.

Serve the admin UI with ``Cache-Control: no-cache`` so the browser always
revalidates (cheap 304 when unchanged) instead of running a stale cached
page. Without this, a UI update only shows after a manual hard-refresh —
every ``index.html``/asset change silently served the old bytes from cache.
``no-cache`` (revalidate), not ``no-store`` (never cache): the ETag makes an
unchanged reload a 304, so it's free when nothing changed.

What an arm IS: the described difference between the two sides.

The measurement method is the same whatever varies between arm A and arm
B; only the realisation differs. Describing that difference as a value —
an axis plus a pair of arms — is what lets one runner serve every axis
instead of one runner per axis, each duplicating the method.

A served endpoint an arm can be measured against.

``model`` is optional because a target reports its own resident id;
it is accepted for an endpoint serving several. ``token`` is optional
because the second endpoint need not share the console's bearer.

One side of a comparison, described rather than implied.

Exactly one of the three realisation fields is read, chosen by the
run's axis. Carrying all three on one type is what lets the driver
take an arm without knowing which axis built it.

How one arm is put on the target, and what that took.

``overrides`` is empty for an axis whose arms are already distinct
without writing anything — which is the whole difference between
comparing two configurations and comparing two servers.

Refuse a pair the axis cannot realise, naming the missing field.

An arm that carries nothing its axis reads would still run — against
a target nobody changed — and report a delta between two identical
configurations as a measurement.

Where this arm runs and what it writes to get there.

The one place an axis is visible. Everything downstream consumes a
:class:`Realisation` and behaves identically whichever axis built it,
which is what makes the ladder, the fold, the verdict and the drift
check one implementation rather than three.

The store's workload key for a run over these prefill depths.

The depths are in the key because a baseline measured at one is not a
baseline at another: prefill depth changes both the TTFT the arms are
compared on and the KV every decode step attends over. Its own
prefix, because this runner sizes a rung's prompts as a TOTAL split
across the rung's streams — a baseline recorded under a per-stream
sizing describes different work under the same name.

What this arm sets, as one readable value.

A realisation that writes exactly one param is that param's value; one
that writes several has no single value and is named by its label
instead, because a set of params rendered as a value is a value nobody
can set.

What the stored record calls one arm's value.

A run that varies exactly one param stores that param's value, so a
sweep over one flag groups the way its reader expects; anything else
stores the arm's label, because a set of params has no single value.

The persisted A/B benchmark store, served to the browser.

One JSON per run, on disk rather than in a JS object that dies with the
tab, so the A arm is measured ONCE per flag set and every later run at
those flags only has to measure B. See :mod:`bench_store` for the file
layout, the baseline key, and the reuse window a drift verdict applies to.

Is there already an A baseline for this (workload, model, flag set)?

A hit means the caller can skip the A arm; ``mismatches`` is non-empty
when the stored baseline does not describe the live target, and the
caller must then say so rather than compare against it.

On-disk store for the admin console's playground A/B benchmark runs.

One run is one JSON file, named ``<name>_<UTC>_<sha>.json`` (baselines carry
a ``_BASELINE_`` marker before the sha), holding a superset of the
``{"args": ..., "cells": ...}`` shape the repo's own bench harnesses write:
``args`` records how the run was configured, ``cells`` the measurements.

The store lives OUTSIDE version control, so every file is written to be
interpretable on its own: it carries the full flag dict (not only its hash),
the model path, the engine identity that produced it, the resolved admission
width and the concurrency rungs actually run.

Baseline identity
-----------------
A baseline is recorded once per BASELINE KEY, which covers three things:

* the workload key — a baseline for one workload says nothing about another;
* the resident model id — a baseline measured on a different model is not a
  baseline, so a different model simply finds none and records its own;
* the flag fingerprint — an order-independent hash of the runtime flag dict.

The engine build (git sha), the target URL and the wall-clock time are
deliberately NOT part of the key. Keying on them would silently fork the
store into buckets nothing ever matches again; recording them instead lets a
lookup RETURN the baseline together with the ways it differs from the live
target, so the caller can refuse it loudly rather than compare against it
blindly.

Location
--------
``ARBI_CONSOLE_BENCH_DIR`` selects the directory; it defaults to
:data:`DEFAULT_BENCH_DIR`. The directory must be writable — a run that cannot
be stored raises :class:`BenchStoreUnavailable` naming the directory, it is
never dropped and never redirected to a temporary location.

Rename the ambiguous throughput/latency keys to self-describing ones.

The workload runner reports ``tps`` end-to-end (it includes prefill) and
``tpot_ms`` decode-only. A stored file read months later has no runner to
ask, so the store spells the difference out.

Assemble the stored record from a client payload.

The fingerprint and the baseline key are recomputed here from the
payload's own flags/workload/model — a client cannot assert a run belongs
to a flag set it was not measured under.

``directory/filename``, suffixed until it names nothing yet.

Two runs in the same second — or a demotion whose new name is already
taken by a later run — must never overwrite an existing measurement.

Look up the stored baseline for (workload, model, flag set).

Returns the fingerprint, the baseline key, the newest matching baseline
entry (or ``None``), and the identity mismatches that make it unusable.

Demote every baseline for a (workload, model, flag set) key.

Demotion rewrites the file with ``baseline: false`` and drops the
``_BASELINE_`` marker from its name — the measurement is kept, only its
role as the A arm is withdrawn.

Every stored run that was compared against one committed baseline.

Newest first. This is the window a drift verdict applies to: each of
these runs read its A arm off that file rather than measuring one.

The newest clean drift verification of one committed baseline.

A verification that found drift does NOT reset the window: the runs
behind it are the ones the drift calls into question, so the next
window starts only once a check comes back clean.

How many runs have REUSED this baseline since it was last verified.

The run that measured the baseline compared against a number it took
itself in the same window, and a verification is not a reuse either —
neither one is drift waiting to be found.

The runs that reused one committed baseline after ``since``.

This is exactly the set a drift verdict calls into question, so the
counter and the suspect-marking read the same definition rather than
two that can disagree.

Stamp stored runs as suspect, in place, with the reason.

The measurement is kept — it happened. What is withdrawn is the
claim that the baseline it was compared against still described the
target, so a reader who finds the file later finds the doubt with
it.

Can this pair be measured that way at all, asked before it is configured.

Scoring needs a boot flag no override reaches, so an operator who
configures a KL run and only then learns the box cannot serve one has
spent that time on a choice that was never available. The probe answers
up front and names the bring-up that would arm it.

Can KL be measured on these sites, and if not what clears it?

Asked by SENDING a scoring request, never by reading a schema: the
surface is declared on every build and refused at runtime on most of
them. Every site is probed, because a pair whose second endpoint
cannot score is a pair that cannot be compared however well the first
one answers.

What this console can measure against these sites, before a run is set up.

The fidelity prerequisite is reported here rather than only as a
mid-run refusal: an operator who learns after configuring a sweep
that the KL half needs a different boot has spent that time on a
choice that was never on offer.

What a live target can actually be asked to do, read off the target.

Every answer here is resolved against the running engine rather than a
static schema: a flag the build ships can still be inert on this box, and
a drafter mode the code knows about can still have no weights behind it.
Reporting the reachable set up front keeps an operator from configuring a
run that was never available.

Can flipping ``flag`` change anything on THIS boot?

The UI asks before offering to benchmark, and the run endpoint asks
again and refuses: an A/B on a param whose path this configuration
cannot reach measures noise, and noise given a number is read as a
result.

Every overridable param, with its reachability, category and swap tier.

One round trip for the whole Config tab: the flag list needs all three
facts about all the params at once, and asking per row would be N
classifications of the same :class:`LiveConfig`.

A param the oracle calls INERT is still returned, carrying its reason.
Hiding it entirely would make the flag look as though it does not exist
on this build, which is a worse lie than showing a knob that does
nothing — the caller filters, and can always reveal the rest.

Which drafter mode the target is in right now, from its own values.

Reported as ``None`` when the live values match neither mode: a
deployment can sit at a K the presets do not name, and calling that
one of them would put a label on a configuration nobody selected.

The drafter A/B's opening state: modes, ladders and what can be scored.

Every ladder is resolved against the LIVE target — the admission width
the scheduler was sized for and the context the KV pool actually
holds — so a rung the box cannot serve is reported as skipped instead
of being offered and then failing mid-run.

What the page cannot know on its own, and nothing it must not know.

The console process is started beside the engine and the observability
stack it drives; the page it serves does not know where any of them are.
A page that guesses puts each address in a second place -- one that has to
be configured per browser, and that nobody edits when a port moves -- and
a wrong guess does not fail loudly: an unreachable datastore and an empty
one both render as an empty panel.

WHAT IS NO LONGER PUBLISHED HERE, and why it is the point of this module:
the ENGINE'S address. It used to be, as a URL the page then handed back to
the console on every call. That put the one address the console derives
for itself into the browser, where it was ``localhost`` for a viewer whose
localhost is their own laptop; it also made the proxy dial whatever URL a
page asked it to, which is a request-forgery surface with no upside. The
page now names a server by ID from a server-side list
(:mod:`arbi_serve.admin_console.registry`) and never sees an engine
address at all -- except as text, in the picker.

WHERE EACH REMAINING ADDRESS IS DIALED FROM decides its shape, and the two
are not interchangeable:

* ``loki``, ``tempo`` and ``prometheus`` are dialed by THIS PROCESS, so
  they are full URLs valid from the console's own network position --
  inside a container that is the bridge gateway, never ``localhost``.
* ``grafana_port`` and ``prometheus_port`` are for the BROWSER: it follows
  the Grafana iframe and queries Prometheus directly (Prometheus is the
  one address both sides need). The host has to be whatever address
  reached the page, which only the browser knows, so the deployment states
  the port and the page supplies the host.

An exported-but-empty variable is not a configured value.

``${ARBI_CONSOLE_LOKI:-}`` in a recipe leaves exactly that shape, and
treating "" as deliberate would hand the page an empty address and call
it the deployment's answer.

What the page is told. ``None`` means "not configured" -- the page
keeps its own default for that one, which is built from the address the
browser used to reach this console.

The observability addresses the page cannot know on its own.

Unauthenticated on purpose, and safe to be: it discloses addresses the
operator configured on this process, to a page this process just served
from the same origin.

This console's own verdict on itself.

Answered even while the console is refusing to serve -- it is the
surface that SAYS what it is refusing about, and it is also how the
port check recognises its own reflection through the docker host.

Whether one registered server answers, and as what.

The picker draws a per-server state from this, so the three outcomes an
operator acts on differently stay separate: it is down, it is up and
refusing this console's credential, or it is serving. A remote that
401s is a CREDENTIAL fact and must not present as a dead UI.

The run: both arms in order, streamed, with every site restored.

Progress is streamed because a run takes minutes and an operator watching
a spinner cannot tell a slow rep from a wedged one. The restore sits in a
``finally`` so a browser that closed mid-run still leaves each target on
the configuration it started from.

A run that measured nothing says so — in an ``error`` frame, in
``run_done.errors``, on every arm that has no rows and on the record it
writes. Folding zero reps yields a NOISE verdict with a real-looking
shape, and that is indistinguishable from two arms that genuinely tied,
so an empty run reported as a verdict is a wrong answer rather than a
missing one.

Why this run has no result at all, or ``None`` when it has one.

A cell is a result only where BOTH arms produced a rep. Where none did,
every verdict is computed over zero measurements and comes out NOISE —
which is exactly what two arms that genuinely tied also look like. The
two must not read alike, so an empty run is reported as the error it is
rather than as a verdict.

Has the committed baseline moved since it was recorded?

Every cell, not just the cheapest one: the bar at each is that cell's
OWN committed spread, so a noisier cell raises its own bar rather
than manufacturing drift.

Commit the run — and, when it measured one, the baseline behind it.

A run that measured nothing is still written down, because the operator
who has to work out why wants the plan and the receipts; it is written
down carrying ``incomplete``, so a later reader cannot mistake it for a
result. A baseline is never committed from such a run.

Matched OpenAI-compatible endpoint comparison for the admin console.

Unlike :mod:`arbi_serve.admin_console.measure`, this runner deliberately uses
only ``GET /v1/models`` and streaming ``POST /v1/chat/completions``.  The
measure endpoint axis is an arbi-to-arbi scientific sweep: it asks each site
for admin configuration, exact tokenizer counts and KV capacity.  A vLLM,
SGLang or Ollama endpoint has none of those surfaces, so presenting that sweep
as a generic comparison would make the UI promise a run it cannot start.

This runner is the portable counterpart.  Both arms receive the same prompt
bytes, request order, sampler settings, concurrency and output ceiling.  It
records TTFT, TPOT, wall-clock aggregate throughput, errors and output length,
then writes a credential-free receipt to the existing benchmark store.

Canonical OpenAI base URL, accepting either host root or ``/v1``.

Credentials in a URL are rejected: unlike a bearer field, they leak into
browser history, proxy logs and the downloadable receipt.

Published evals (AIME26, IFEval, MMLU-Pro, GPQA, C-Eval) from the console.

Runs a benchmark against a served endpoint and stores the score as a row, so
"is the thing we just deployed still as good as the thing we published?" is a
button rather than an afternoon of assembling harness flags by hand.

WHY THIS MODULE OWNS NO SCORING
-------------------------------
Every eval parameter -- which task, thinking on or off, the generation budget,
the sample cap -- is READ AT CALL TIME out of ``arbi_serve/evals/catalogue.yaml``
by :func:`arbi_serve.evals.catalogue`, and the run shells out to
``python -m arbi_serve.evals``. Nothing here decides how a model is scored.

That indirection is the point. A score is only worth looking at if it is
comparable with the published baseline, and a copy of "thinking on, 32768
tokens, greedy" in this file would go stale the first time someone tuned the
catalogue -- silently, while still rendering a confident number in the UI. The
console asks the catalogue what the protocol is every time it draws the list,
and each row carries the ``catalogue.yaml`` line it was read from, so the
protocol is auditable from the screen.

The same reasoning is why a missing harness yields an EMPTY catalogue and a
stated reason rather than a built-in default set: a fallback protocol is
indistinguishable, on screen, from the real one.

ENDPOINT-AGNOSTIC ON PURPOSE
----------------------------
This module knows nothing about the console's managed target. It takes a base
URL, so the same button scores the server this console manages, a colleague's
vLLM on another box, or ollama on a laptop -- and the UI is what decides which
of those "the deployed server" means today.

The model to score, when the operator did not name one.

Only ever resolved when the endpoint serves exactly ONE model. Taking the
first of several would score whichever model the endpoint happened to list
first and label the row with it -- and a gateway (LiteLLM, ollama) lists
dozens, so the wrong answer is the usual case, not the edge one. A run
against the wrong model is worse than a run that refuses to start, because
it produces a plausible number under the right-looking heading.

The driver invocation -- the one thing this module assembles.

Runs the harness through ``sys.executable`` so it lands in the very
interpreter the console is running in, which is the one that has this
package (and therefore the ``arbi-openai`` adapter) importable.

``--include_path`` is what makes this project's own tasks
(``aime26_chat``) resolvable; without it the harness only knows its
built-ins.

Stream one eval run: progress while it works, rows when it finishes.

Streamed rather than awaited because a real run is minutes to hours, and
an operator watching a spinner cannot tell a slow endpoint from a wedged
one -- the same reason the workload runner streams.

Which part of the engine a flag belongs to, and what flipping it costs.

Two questions the admin console's flag list cannot answer from the
override catalogue alone, and both are answered by DERIVING from the
registries rather than by a hand-written table — so a flag added
tomorrow lands in a category and a tier by itself.

Category
    Read off the param's ``target`` and its name, in a declared order,
    because the first rule that matches wins and the rules overlap on
    purpose: ``mtp_capture_full_k_ladder`` is a drafter flag that
    happens to mention capture, and ``dflash_capture`` likewise, so the
    drafter rules are consulted before the capture rules. A handful of
    params do not classify from their name at all; those carry an
    explicit override, and anything still unmatched lands in ``other``
    rather than being guessed into a category it may not belong to.

Swap tier
    What a flip actually costs, before the operator clicks. The three
    live paths ``POST /v1/admin/config_override`` distinguishes are a
    capture-affecting delta (build and capture), a runtime-only delta
    that needs a member built with the value (an instant swap onto a
    prepared variant), and a fresh-read flag delta applied live on the
    active member. A param the target does not list as overridable is
    fixed at boot.

    The line between LIVE and VARIANT is the same fact
    :mod:`arbi_serve.flag_reachability` uses to refuse a live A/B: a
    runtime-scope flag whose contract fires in a build-time window has
    already had its one chance by the time an overlay installs, so it
    needs a member built with the value. Both readings come from
    :data:`~arbi_serve.flag_reachability.BUILD_TIME_PHASES`, so the tier
    a row advertises and the verdict the oracle gives can never
    disagree.

Which part of the engine ``name`` belongs to.

``target`` is the param's catalogue target when the caller has it;
it resolves the params whose own name says nothing (a dotted
``cache.*`` path is KV whatever it is called).

What flipping ``name`` costs on a live engine.

``entry`` is the param's row from the target's own
``GET /v1/admin/config_override`` catalogue, or ``None`` when the
target does not list it — which is itself the answer, since a param
outside the override registry is fixed at boot.

The target's catalogue wins where it speaks (it describes the build
that is actually serving); the local registry supplies the scope and
the contract phase it does not carry.

GPU clock lock, so a bench repeats.

The single biggest source of A/B run-to-run drift is unlocked GPU clocks.
The browser can't run nvidia-smi, but this console process can (its
container is granted cap_add: SYS_ADMIN + a GPU device — see
docker-compose.08b.yaml). Co-located box only: this pins THIS host's GPU,
which is the engine's GPU here.

Whether this console can read and pin clocks, and the clock if it can.

A capability question, so a console packaged without NVML or running off a
host with no card answers it — ``available: false`` with the reason —
instead of raising. A 500 here would say the console is broken, when the
truthful answer is that this deployment simply cannot do it; the UI hides
the control on that answer rather than rendering one that cannot act.

One HTTP client per console process, not one per request.

WHAT THIS FIXES. Every proxied read built its own
``httpx.AsyncClient``, and building one is not cheap: httpx's default
transport calls ``ssl.create_default_context()``, which loads and parses
the whole CA bundle. MEASURED in the shipped image, ``cProfile`` over ten
constructions: 9.6 ms each, 9.3 ms of it inside
``SSLContext.load_verify_locations`` — paid on every panel refresh, and
paid identically for an ``http://`` target that will never use TLS.

That is the ~11 ms floor under every row of #2255's cost table. The table
was taken through this proxy, so it was reading the console's own
construction cost as if it were the engine's: a read of a cached constant
(``compile_cache``) measured 10.8 ms there and **0.4 ms** against the
engine directly. The admin RPC has no such floor.

WHY A SHARED CLIENT RATHER THAN A SHARED SSL CONTEXT. Hoisting only the
context would remove the 9.3 ms and leave the rest: a per-request client
also means a fresh TCP connection every poll, no keep-alive, and a
connection pool that is discarded before it can be used. The console
dials the same handful of targets over and over — that is exactly the
workload a pooled client exists for.

WHAT STAYS PER-CALL. Timeouts and credentials. They differ per call (a
health probe is not a 900-second completion; a target's token is the
target's), and httpx takes both per request, so nothing is lost by not
baking them into the client.

WHAT DELIBERATELY DOES NOT USE THIS. ``endpoint_bench`` builds its own
client with its own connection limits, because a benchmark that shares a
pool with the console's own polling is measuring the pool.

The console's process-wide async HTTP client.

Carries no base URL, no default headers and no default timeout: every
call site passes an absolute URL and its own ``timeout=`` /
``headers=``, so the shared object holds nothing that could leak one
target's credential onto another target's request.

A generated ``arbi_serve_client.Client`` for one target, reused.

The generated client builds an ``httpx.AsyncClient`` of its own on
first use and caches it on the instance — so a fresh ``Client`` per
call paid the same CA-bundle load as a fresh raw client. Cached on
``(base_url, token)`` because the credential is baked into the
client's headers: keying on the URL alone would hand one caller's
token to the next.

The three facts a console must establish about itself before it serves.

A console that starts cleanly and is wrong is this component's whole
failure history. Each of the three below was hit in one bring-up, each
produced no error anywhere, and each presented as a UI that looked
finished and was not:

1. THE ENGINE IS THERE. The console shares the engine's network
   namespace, so the engine is at loopback on its own port by
   construction -- but "by construction" is a claim about a deployment,
   and a claim is checked, not assumed. Unverified, a console renders
   every panel as "unreachable" and blames the engine. The claim is also
   only made where the deployment makes it: a console that does not run
   beside an engine is not refused for the absence of one, it is told
   that nothing about a local engine was verified (see :meth:`run_once`).
2. ITS PORT IS REACHABLE. Under ``network_mode: service:<engine>`` the
   console cannot publish a port; the ENGINE'S container does. A console
   that binds a port the engine does not publish starts, logs
   ``Uvicorn running on ...``, and is reachable by nobody. There is no
   error, on either side.
3. IT IS THE ENGINE'S OWN CODE. Console and engine are one repository,
   and the console renders the engine's vocabulary. A console serving an
   older tree does not fail, it misreports.

Each check ends in one of three verdicts and never in two: it PASSED, it
FAILED (and then the console refuses to serve, saying exactly what it
tried), or it COULD NOT RUN -- which is reported as its own state and
never folded into "passed", because a guard that reports success on a
check it did not perform is worse than no guard.

AND THEN FACT 1 EXPIRES. The three above are established once, at boot,
and the first of them stops being true the moment the ENGINE restarts:
Docker builds the new engine container a NEW network sandbox, this console
keeps holding the old one, and its listener is left in a namespace nothing
routes to -- while the published console port now belongs to the new
sandbox, where nothing is listening. The console process is still running,
still healthy by every signal Docker has, and reachable by nobody, forever,
until someone restarts it by hand. :class:`EngineWatchdog` is what keeps
fact 1 checked for as long as the console serves, and its cure is to EXIT:
re-running the container is what re-resolves ``network_mode:
container:<engine>`` onto the live namespace, and nothing inside this
process can.

The port this console bound.

Falls back to the published constant rather than to zero: an app
imported by something other than ``__main__`` still binds that port by
default, so answering 0 would make the publication check ask about a
port nobody uses.

Whether this process is containerized.

The port check only means something in a container: outside one there
is no publication step to get wrong, and probing a host's own gateway
would be asking a question nobody asked.

This container's default gateway, i.e. the docker host.

The host is where a published port is actually bound, and from inside
the container that address is the only way to see the publication at
all. The deployment already depends on this route being open -- it is
the same address the console dials Loki, Tempo and Prometheus on.

Wait until something is accepting on ``host:port``, or the deadline.

Both parties this module cares about are asked through this one
function, because it is one question: is a listener there yet?

THIS CONSOLE'S OWN PORT. The publication check dials the docker host and
expects itself to answer, so it has to run after the listener exists. It
does not: the checks start from the lifespan, which uvicorn completes
BEFORE it accepts anything -- so without this the very first probe would
find nothing listening and refuse the boot it was meant to protect. That
is the shape of bug this whole module is about, and it would have been in
the guard itself.

THE ENGINE'S PORT, for :class:`EngineWatchdog`. A TCP connect and not an
HTTP request on purpose: the watchdog asks whether the namespace this
console holds still has the engine in it, and a listener answering the
handshake is exactly that question. Anything richer would also fail on a
busy or half-built engine and read it as a dead namespace.

Returns as soon as a connection is accepted, and otherwise keeps trying
until ``deadline_s`` -- so a caller that passes its own probe interval
gets that interval's worth of retries and no sleep of its own.

Is ``port`` published by the container whose namespace we are in?

Asked the only way it can be answered from inside that namespace: dial
the host's own address on that port and see whether THIS process is
what answers. The reply carries a per-boot instance value, so
"published and mine", "published and somebody else's" and "not
published" are three distinct answers rather than one ambiguous one.

It asks about the SAME number on the host, which is what every shipped
recipe publishes (``N:N``, held there by
``tests/test_console_port_publication.py``) and the only mapping under
which the question has one answer. A hand-rolled ``-p 8888:8899`` is
reported as unpublished, and the message says so rather than leaving
the operator to guess which of the two it meant.

Returns ``(check, detail)``. A check that could not run says so; it
never returns ``passed``.

The host and port the engine's listener is at, from its base URL.

Read off the address the console already derived rather than restated:
the URL is the only place that fact lives (``registry.local_target``
builds it from the shared namespace), and a second constant here would
be a second thing to get wrong.

End this process from a background task, having flushed the reason.

``sys.exit`` here raises inside the watchdog's task and uvicorn keeps
serving; asking uvicorn to shut down gracefully means draining a
listener in a network namespace nothing routes to. ``os._exit`` after
``logging.shutdown`` is the one form that both delivers the message and
ends the process. Non-zero: this is a failure the restart policy is
being asked to act on, not a clean stop.

Exit when the engine stops being in this console's network namespace.

THE FAILURE. ``docker restart`` on the engine strands this console
permanently and silently. The console joined the engine's network
sandbox (``network_mode: container:<engine>``); the restart builds the
engine a NEW one; this process keeps holding the old. Its listener is
now in a namespace with no route to it, and the console port published
by the engine's container points at the new sandbox, where nothing
listens. Docker sees a running container with a passing health check.
An operator sees a page that never loads and no error anywhere.

WHY EXITING IS THE FIX and not a workaround: re-resolving
``container:<engine>`` onto the live namespace happens at container
start and nowhere else, so no code inside this process can reattach it.
``restart: unless-stopped`` -- already on the console service in every
shipped recipe -- does exactly that, the moment the process ends.

WHY IT IS RIGHT EVEN WHILE THE ENGINE IS MERELY BOOTING. There is no
state in this deployment where a console should keep running with an
unreachable engine: either the namespace is dead (exiting is the only
cure) or the engine is restarting, in which case this console has to
re-attach to its new namespace anyway. It also serves nothing while the
engine is unreachable (:func:`app._refuse_until_verified`), so the
restart costs no availability, and Docker's own backoff bounds a loop
against an engine that is genuinely down.

IT ONLY ARMS IN THE OP MODE THAT HAS THE FAILURE, and says which mode it
decided it is in. The console has two, and they are not variations of one
thing:

* SIDECAR to a local engine -- the shipped recipes, a container inside
  the engine's network namespace. The local engine is derived rather than
  configured, this console already refuses to serve anything at all while
  it is unreachable (:func:`app._refuse_until_verified`), and it is the
  only mode where a namespace can go dead underneath the process. This is
  the mode the watchdog is for, and exiting in it costs no availability
  that the refusal was not already costing.
* CONTROL PLANE for REMOTE servers -- ``pip install arbi-serve[ui]`` on a
  laptop, driving engines added through the registry. There is no shared
  namespace, no local engine is expected to answer, and the absence of
  one is not a fault: :meth:`LocalEngineCheck.run_once` serves that
  console deliberately. Exiting it would replace the old footgun with a
  new one -- a control plane that kills itself for the state it is
  designed to run in.

The two are told apart by the container, and then by ``__main__`` having
set the listen port -- the app is importable, and under a test client or
any other host process it is a guest. A guest does not kill its host.

Is the engine running the same ``arbi_serve`` source as this console?

Returns ``(check_or_state, detail)``. An engine that does not
answer the question yields ``unavailable`` and NOT a pass: the
console still serves, and says out loud that the comparison did
not happen.

Re-check until the console is serving, then stop.

UNREACHABLE is transient by nature -- the engine is still starting
-- so it is retried. A source mismatch and an unpublished port are
not: neither can become true again without a redeploy, and a
console that quietly started working after refusing would teach an
operator to ignore the refusal.

ONE line stating what was derived, from what, and what changes it.

The engine states its own provenance this way and the console is
held to the same bar: an operator reading the boot log must be
able to see which server this console drives and why, without
opening a compose file.

Range-query Loki and return the newest lines, newest first.

Proxied rather than fetched from the browser: Prometheus ships permissive
CORS headers and is queried directly, but Loki does not, so a direct XHR
from the console origin is blocked.

Returned as plain ``{ts, line, level}`` records so the console can RENDER
logs itself. That is the point — an embedded Grafana Explore makes every
row individually expandable to show its label set, which for our OTLP
label cardinality is the same handful of labels on every line and carries
nothing an operator wants. Owning the render means the panel shows the log
text and nothing else.

One A/B runner whose arms differ along a SELECTABLE axis.

An A/B is a measurement method plus a description of what varies between
the two arms. The method — canonical prompts, the feasible cell grid, the
warmed ladder, the per-cell fold, the verdict's spread bar, the drift
check on a reused baseline — is the same whatever varies. Only the
realisation differs: putting the target on arm B means writing a runtime
override, or writing the drafter's three params together, or dialling a
different base URL.

Splitting the runner per axis therefore duplicates everything that
matters and shares only the part that does not. This module keeps the
method in one place and makes the varying part a described thing:

:class:`Arm`
    ``{label, flag_overrides?, drafter_mode?, endpoint?}``. What an arm
    IS, rather than a set of positional query parameters whose meaning
    lives in the handler that reads them.

:data:`AXES`
    ``flags`` — arms differ by runtime-flag overrides.
    ``drafter`` — arms differ by drafter mode.
    ``endpoint`` — arms are two different served endpoints, which is what
    lets the console compare two servers rather than two configurations
    of one.

:func:`realise_arm`
    The ONLY axis-dependent step: where an arm runs and what it writes
    before it is measured. Everything downstream of it takes a
    :class:`Realisation` and cannot tell which axis produced it.

Fidelity is asked on EVERY axis, and refused by name where it cannot run
    A speed-only A/B lets a numerics regression ship looking like a win,
    so KL and argmax flips are attempted on every axis. Where the
    comparison cannot be made the run says which prerequisite is missing
    and what clears it. It never emits a zero: zero KL and zero flips are
    exactly what perfect agreement looks like, and a run that could not
    score must not be readable as a run that scored perfectly.

The prerequisite is a CAPABILITY, readable before a run is configured
    Scoring needs ``ARBI_SERVE_LOGPROBS=1`` at boot — no override reaches
    it, because it reserves the logprobs scratch pool out of KV capacity.
    Discovering that only after configuring a run wastes the operator's
    time on a choice that was never available, so
    :func:`fidelity_capability` answers it up front and names the
    measurement overlay that arms it.

The method is split by the part of a run it answers for:
:mod:`arbi_serve.admin_console.arms` describes what varies between the two
sides, :mod:`~arbi_serve.admin_console.capability` answers whether a pair
can be measured that way at all, :mod:`~arbi_serve.admin_console.sweep`
measures one arm across the cell grid, and
:mod:`~arbi_serve.admin_console.driver` runs both and streams the result.
This module is the endpoint: it turns a request into a plan and hands the
plan to the driver.

:mod:`arbi_serve.admin_console.targets` owns the target I/O primitives the
method drives (the OpenAI client, the admin GETs and POSTs, the tokenizer
count, the scoring request) and :mod:`~arbi_serve.admin_console.ab_records`
the shapes a run is folded and written down in. Both are reached through
the module object rather than imported by value, so a test that
substitutes a fake target substitutes it for the whole console.

The deprecated aliases below are the only definitions of
``/playground/flag_ab`` and ``/playground/drafter_ab``; ``test_measure_ab``
asserts both resolve here, because the router is included ahead of every
other one and a route re-declared elsewhere would be silently shadowed
rather than reported.

Resolve a run against the LIVE sites, before any stream opens.

Every rung, every mode and every prerequisite is checked here so a
bad request is an HTTP error the caller can act on rather than an
error frame inside an otherwise-successful stream.

A short name for what this run varied, for the stored record.

One override is named by its param, so a flags run keeps the record
name a sweep already groups by; anything else is named by the axis and
the two arms, because a set of params has no single name.

One A/B, on the axis the caller names, streamed as it runs.

``arms`` is a JSON array of exactly two ``{label, flag_overrides?,
drafter_mode?, endpoint?}`` objects. The axis decides which of those
fields is read; everything after the arms are realised is identical
whichever axis was picked.

``switch_ceiling_s`` bounds ONE capture-affecting swap — the wait itself
tracks the target's readiness, so this only says when a build that never
finished stops holding the run open. It scales with the model, which is
why it is the caller's to set; ``0`` takes the default.

The admin and OpenAI surfaces of a target, relayed to the browser.

Every call here is dialled server-side (see the docstring on
``targets._unreachable_detail``): the target only has to be reachable from
the console's host, not from the operator's browser. The realtime relay at
the end is the one exception a WebSocket forces, and it is a genuine
bidirectional pump for that reason.

Stop one in-flight request reasoning on the target.

``request_id`` is the caller's own ``X-Request-ID`` or the engine's
per-process counter; the target resolves the caller's id first. The console
stamps an ``X-Request-ID`` on every request it launches, so a bar it
launched is addressable by the id it minted rather than by a counter it
would have to look up.

Stop one in-flight request on the target by forcing its stop token.

The request ends the way it would have ended anyway, which a dropped
connection does not achieve — and on this stack a dropped connection is not
reliably noticed at all (see ``workloads._RUNS``).

The target's live-tunable knobs, at the values that GOVERN.

Distinct from ``/proxy/config_override``, which is the boot-parameter
registry: a knob can be absent there and still have a value the engine
enforces. The KV admission watermark is one, and the console draws the
KV lane's thresholds from it rather than choosing them here.

Everything one swimlane tick needs, in ONE request to the engine.

The tick used to cost three: ``serving_state`` is itself ``/health/ready``
plus ``server_info``, and the timeline is a third. Three round trips to
draw one panel is the wrong number whatever each costs — and the reason it
is wrong is that they were three different INSTANTS, so the panel could
render a queue depth from one moment against rows from another.

404 on an older target that has no such route; the page falls back to the
two-call path, which is why that path is kept rather than deleted.

The engine's readiness verdict in the vocabulary the strip renders.

ONE translation, here, because this module already owns
:data:`_READY_REASON_NOTE` — the map from a machine reason to the sentence
an operator acts on. Teaching the page a second vocabulary for the same
fact is how the two drift.

Binary serving availability for the swimlane's load/uptime strip: is the
target ACCEPTING requests right now, and if not, a one-line reason note.

Source of truth is the load-balancer readiness probe (``/health/ready``),
so the strip agrees with what the LB would do. Sleep is read separately from
``server_info`` because readiness deliberately keeps an asleep pod "ready"
(it wakes on demand) — but for the operator strip, asleep = not serving.

The calibration bundles the target can reload from — seeds the reload
dialog's calibration picker, so an operator chooses a bundle the server
can actually see instead of typing a path into a blank box. 404 if the
backend predates the route (the UI then keeps free-text entry).

Forward a backend variant swap, carrying its calibration and basis.

POST, unlike its GET neighbour, because this REBUILDS: the target parks a
prepared variant beside the live one. The basis travels WITH the backend
rather than being a separate switch, which is what makes re-selecting a
prepared Hadamard or OSCAR variant instant instead of a recalibration --
the codebooks and the rotation they were fit in are one unit, and a
surface that let them be set independently would let an operator pair a
bundle with a basis it was not fit in.

A 900 s timeout: the first selection of a variant builds and captures it.

Point-in-time GPU memory snapshot by named pool. Not what the A/B
ablation panel's memory chart uses for its time series — it queries
Prometheus directly for that (the same scraped gauges the VRAM Grafana
dashboard shows), since there is no server-side time-series endpoint
here (``memory_record_start/stop/dump`` only toggle PyTorch's
alloc/free history and pickle it to disk for the external memory_viz
tool, never inline JSON). This route stays for a quick single-snapshot
read without going through Prometheus at all.

Flat per-stage timeline of recent + in-flight requests — the same
live scheduler read that powers the Grafana swimlane panel. Useful on
its own for spot-checking in-flight request bookkeeping (e.g. how many
Request objects a single connection actually admitted) without going
through Grafana at all.

Whether the engine is exporting spans at all, and if not, why.

An empty trace panel has two causes an operator acts on differently: this
request has no trace yet, or this ENGINE has never exported one. Only the
engine can tell them apart, and it says so exactly once in a boot log line
nobody is reading by then. It is also recorded in the flag-truth counter
that gates the channel, which is a surface, so the page can ask.

Fetch a request's trace from Tempo and return a flat span list for the
console's NATIVE waterfall — no Grafana. Two hops: TraceQL-search by
request_id (spans carry it as an attribute; no trace_id is exposed anywhere
in arbi-serve), then fetch the winning trace by id. ``tempo`` is the Tempo
base URL from the console's Settings; ``start``/``end`` (unix seconds) scope
the search under Tempo's max-search-duration.

Per-(backend, layer) dequant relative-L2 error for every registered
TKV backend, at whatever config is CURRENTLY active on the target — the
A/B ablation panel calls this once per column (value already applied)
to report a correctness signal alongside perf, not just perf alone.

Watermark detection: score pasted text (or token ids) against the
target's configured watermark key. Pure math on the target; the key
never travels to the browser — the console only relays the verdict.

Flush the target's radix prefix-cache tree.

Pages held by in-flight requests stay live until those finish; the tree is
detached so future matches miss, and anything already evictable returns to
the pool. ``?tenant=`` scopes it to one tenant.

Ephemeral A/B teardown: drop parked config variant(s) on the target and
reclaim their VRAM (regrowing the active member's KV back to full capacity).

``?key=<variant>`` drops one parked variant; no key drops EVERY non-active
variant, resetting residency to the single active baseline member — the
ephemeral guarantee (no permanent KV/memory tax after an A/B). Raw httpx
forward: the generated client has no binding for this route yet.

Live model reload: drain in-flight, free the old model, rebuild on
the new one. Same weight class as a full restart minus the process
bounce — every in-flight request is dropped, so the UI should confirm
before calling this, not fire it on a stray click.

Drain in-flight, then exit the target process.

Whether the exit becomes a RESTART is the target deployment's supervisor
policy — the response carries what it declared, and the UI labels the
action from that rather than assuming one.

Relay one streamed completion, keyword-strict SDK be damned.

``body`` may carry arbi-serve's own vendor extensions (``top_k``,
``mtp_k``, ...) that the ``openai`` SDK's typed ``create()`` signature
doesn't recognize and rejects as unexpected kwargs. ``extra_body`` is
the SDK's own documented escape hatch for exactly this — it merges
verbatim into the outgoing JSON regardless of whether the SDK's
Python signature knows the key, so routing every field through it
(not just the non-standard ones) sidesteps the whole problem at once.

WHICH servers this console may dial, and the credential for each.

The console is a browser page in front of a server-side proxy, and which
engine a proxy call goes to is chosen per request. There are two ways to
say which, and only one of them is safe.

THE UNSAFE ONE, which this replaces: the page names a URL and the console
fetches it. That is a request forgery hole with a UI on top -- anything
that can reach the console can make the console fetch any address the
console can reach, and a console is deliberately placed where it can
reach engines, dashboards and metric stores that are not otherwise
exposed. It is also unable to hold a credential, because the only place
left to keep one is the browser.

THE ONE HERE: the addresses live server-side, in this registry, and the
page names an ENTRY by id. An id that is not registered is refused, so the
set of addresses the console will dial is exactly the set someone
deliberately added -- a closed list, not an argument. Each entry carries
its own credential, sent server-side and never rendered back to the page.

ENTRY ZERO IS DERIVED AND FIXED. The console runs inside the engine's own
network namespace (compose ``network_mode: service:arbi-serve``), so the
engine is at loopback on its own port BY CONSTRUCTION -- there is nothing
to configure and nothing to get wrong, which is the point. It is also why
the local engine needs no credential: the request arrives at the engine as
loopback, which its admin auth already permits, so nothing about the
default deployment ever requires opening admin routes to non-loopback
peers. It cannot be edited or removed, because a console with no engine is
not a state an operator meant to reach.

PERSISTENCE is a JSON file under the cache volume the console already
mounts read-write (the same volume its TLS cert persists in), so a
registry survives a restart without inventing a store or a mount. The file
holds credentials, so it is written 0600 and never read by the page.

The console's server list: one derived entry plus the added ones.

Reads and writes are serialized under a lock. The console is a single
async process, but the store is a file and two concurrent adds racing
on read-modify-write would lose one of them silently.

The process-wide registry.

A module-level singleton rather than a FastAPI dependency because the
proxy resolves targets from plain helper functions that no route hands
a request-scoped object to; tests substitute it with
:func:`set_registry`.

What the page is allowed to know about this entry.

The credential is reported as a BOOLEAN. A page that must show
whether a server has one still learns nothing it could replay, and
an operator can see at a glance which entry is missing the thing a
401 is about to be caused by.

The added entries, or an empty list when nothing was ever added.

A store that exists and cannot be parsed is NOT an empty registry:
answering "no servers" for a file full of servers is the console
reporting success on a read it failed. It refuses instead, naming
the file.

Write the added entries, atomically and readable only by us.

Atomic because a half-written list is a list that loses servers on
the next read; 0600 because the file holds credentials. A write
that fails is raised, never swallowed -- an entry that seemed to be
added and is gone after a restart is the silent failure this whole
module is arguing against.

The entry ``ref`` names, or a refusal.

This is the whole security boundary and it is deliberately a
lookup with no fallback: an unknown id is refused, and an id that
happens to look like a URL is refused by the same rule, with the
reason spelled out -- because "it used to accept a URL here" is
exactly the regression that reopens the hole.

Edit one added entry. ``token=None`` keeps the stored credential.

Keeping it on ``None`` rather than on empty-string is what lets the
page edit a label without being handed the credential first: it
never has the value, so it cannot send it back.

One arm, measured across the cell grid, and scored against the other.

This is the method every axis runs: put a site on the arm, prove it can
still serve, warm the ladder, walk the feasible (prefill, concurrency)
cells and fold each one. The fidelity half at the end is attempted on
every axis too and REFUSES by name where it cannot run, because a KL of
zero and a KL that was never measured must not read alike.

One SSE frame.

Both spellings of the frame's kind ride on every frame: ``type`` is
what the flags panel reads and ``event`` what the drafter panel
reads. One stream readable by both is what lets the UI migrate a
panel at a time instead of in lockstep with this module.

Prefill depths from a comma list, filtered by what the target can hold.

A depth past the live context is dropped BY NAME rather than clamped:
a rung silently served at another depth reports a measurement under a
label nothing measured.

Put one site on ``overrides``, waiting on the target's own progress.

Returns whether anything was written. A set of values the target is
already on is NOT written: the same value in another spelling still
signs as a delta on the way in, and a capture-affecting delta is a
member build plus a recapture that changes nothing while it shrinks the
KV pool the run is measured against.

A capture-affecting write is quiesced first, because the config swap
parks the outgoing member and the park's co-residency guard refuses
while prefill transients are still in the allocator; and it is then
waited out against ``/health/ready`` rather than a fixed nap, because
the build behind it takes as long as the model is big and no constant
is right for two model sizes at once.

Write an arm's overrides to its site and let the change settle.

A realisation with no overrides writes nothing — an endpoint arm is
already itself — and neither does one whose values the site already
carries.

One prompt set per cell, sized so the rung's TOTAL prefill is fixed.

A cell at concurrency ``c`` gets ``c`` prompts of ``total / c`` tokens,
so every concurrency in a row prefills the same number of tokens and a
difference along the row is the batch shape rather than the amount of
work. Sets are cached by (per-request size, count) because two cells
can ask for the same prompts and each costs a bisection against the
target's own tokenizer.

Whether the second site tokenizes the shared prompts to the same depth.

Two endpoints can render different chat templates or carry different
tokenizers, in which case the same bytes are not the same prefill and
a TTFT comparison is between two different amounts of work. The
disagreement is reported rather than corrected: correcting it would
give the arms different bytes, which the fidelity half cannot align.

Why this site cannot currently serve, or ``None`` when it can.

An engine that hits an infrastructure error stops running forwards and
sheds every request until it clears, so every cell after that point
measures the shedding. Only a readiness surface that ANSWERS stops the
run: a surface that is not there (404/405) is an unreadable instrument
rather than a diagnosis, and refusing on it would stop every run
against a build that never shipped the route.

Did this arm's realisation actually take effect?

Axis-specific by necessity and by nothing else: a runtime flag proves
itself through its declared contract counter, a drafter mode through
a counter FAMILY plus tokens-per-step, and an endpoint through the
identity the site itself reports. All three answer the same question
in the same shape, so the caller renders one thing.

One arm's pass over the cell grid, bracketed by the truth counters.

The bracket is a plain before/after DELTA with no reset: the reset
endpoint also zeroes GRAPH_RECORD counters, which fire only at
capture and can never be refilled by traffic, so zeroing one would
make a capture-affecting path read dormant for the life of the
process.

Every cell's verdict, from the reps both arms produced there.

THE fold, for every axis. The paired schedule compares each rep's own
two arms; the reused-baseline schedule compares two independent
samples and needs spread on both sides. Nothing here can tell which
axis produced the rows, which is the property that keeps one bug from
having to be fixed in three places.

The folded cells as the panel's four speed readings per arm.

Derived from the same rows the verdicts are, not measured again: two
renderings of one measurement can disagree only if one of them is
computing something else.

KL and argmax flips between the arms, plus the null control behind them.

Scoring runs with speculation OFF because that is the only state the
scorer produces anything in. On the drafter axis that makes the
comparison structurally empty and this says so; on every other axis
the arms remain distinct with speculation off, so the KL is real.

A refusal always names the prerequisite. It never returns a zero: a
zero KL and a zero flip count are what perfect agreement looks like,
and a run that could not score must not be readable as one that did.

Everything this console does TO a target arbi-serve instance.

One module because the console is a client of a live engine and nothing
else: every route, and the A/B runner in :mod:`measure`, reaches an engine
only through the calls here. Reaching them through this module object
rather than importing them by value is what lets a test substitute a fake
target for the whole console at once.

The cache tenant the PAGE asked this request to be served under.

Read off the incoming header and forwarded verbatim on the outgoing call.
The console never invents one and never stores one: a tenant it chose
would be a tenant the operator could not test against, and the whole point
of the field is to see whether the header the browser sent is the
namespace the engine reports having served under.

Reject a blank/whitespace base URL loudly instead of handing httpx an
empty base_url (which fails with an opaque, unrelated error deep in
the transport).

Takes a base URL, not a page-supplied name: everything a BROWSER names
goes through :func:`dial` first, which is where the registry lookup and
its refusal live. The endpoint-benchmark arms are the callers that
legitimately arrive here with a URL.

Resolve a server the PAGE named into an address and a credential.

The single seam between "the browser said which server" and "the
console dialled one", and the reason there is exactly one: a second
place that turned a page-supplied string into an address is a second
place that could accept a raw URL, and accepting a raw URL is what
makes a proxy into an open relay. Every route the page reaches goes
through here; nothing else resolves.

The credential is the one stored WITH the registry entry. A browser
``Authorization`` header is honoured only for an entry that stores
none, which keeps the Settings-tab token working for a self-managed
engine without letting a page override a stored credential.

Turn a raw connect/timeout exception into an actionable message.

A target can legitimately be any remote arbi-serve instance — this
is a fleet-management console, not a same-host-only tool — so the
fix here is never "assume localhost". It's making the failure mode
itself informative: name the target, name the underlying reason,
and call out the one mistake this shape of error can't distinguish
from a genuinely-down remote — that ``target`` is resolved and
dialed from the console process's own network position, not from
the browser viewing this page.

Name the LAYER and the ADDRESS a panel's data did not come from.

An empty panel has three causes an operator acts on differently -- the
store is down, the store is up and holds nothing yet, or this console
was told the wrong address -- and once the message is "failed" they are
indistinguishable, so the operator learns to distrust the banner rather
than read it. Every store this console dials answers through here so
the three stay separable, and so the answer names which store, at which
address, dialed from whose network position.

The generated client for one target, REUSED across calls.

The generated ``Client`` builds an ``httpx.AsyncClient`` on first use,
and building one loads the CA bundle into a fresh SSL context — so a
fresh ``Client`` per proxied read paid that on every panel refresh.
:func:`arbi_serve.admin_console.http.admin_client` caches on
``(base_url, token)``; the credential is part of the key because it is
baked into the client's headers.

Await a generated ``*_detailed`` admin call and hand back its parsed body.

The generator's own ``.parsed`` is only populated when the target route
declares a ``response_model`` AND the call succeeded — read ``.content``
off the ``Response`` object directly instead, so this works uniformly
whether or not that route happens to be typed.

Connect/timeout failures are translated to a 502/504 with an
actionable message (:func:`_unreachable_detail`) instead of bubbling
up as an unhandled ``httpx`` exception, which FastAPI otherwise turns
into a bare, detail-free 500.

arbi-serve validates ``model`` against the actual resident id (the
model path, per ``GET /v1/models``) — it does NOT accept an arbitrary
placeholder string, so every synthetic request needs the real id.

The target's serving ``max_context``, or ``None`` when unavailable.

Used by workloads that size their prompt as a fraction of the live context
window; callers fall back to the workload's fixed word count.

Apply ``prefill_capture`` via the same config-override path the Config
tab's own A/B panel uses. Capture-affecting (rebuild+recapture on the
target) — same cost as flipping it by hand; called once, before the
workload's requests fire, not per-request.

Is ``wanted`` the value the target is already serving on?

Compared through the registry's own parser, because the two sides spell a
value independently: the live read returns whatever the target's config
field holds, and an override carries whatever the caller typed. An equal
value in another spelling posted as a delta costs a member build and a
recapture for no change at all.

Seconds spent waiting for the target to report ready, or ``None``.

``None`` means the ceiling passed with the target still not serving —
the caller's evidence that the swap failed, rather than an assumption
made from a clock. A readiness route that is not there (404/405) is an
unreadable instrument, not a diagnosis, so it reads as ready: refusing
on it would stop every run against a build that never shipped the route.

Read one ``completion_tokens_details`` field off a usage payload.

The field arrives as an attribute on a typed usage object or as a key
in a passthrough dict depending on the client, and a target that does
not draft sends neither.

GET one admin surface off the target, or ``None`` when unreadable.

An unreadable surface is a fact the reachability verdict carries as
UNKNOWN; it is never a crash and never a guess.

POST to one target surface, returning its JSON body when it has one.

``timeout`` overrides the client's own read deadline for this one call.
A config override that rebuilds and recaptures a pool member answers only
when the build is done, so the deadline a swap is given has to be the
ceiling the caller is willing to wait for that build — not the deadline
that suits an admin GET.

Assemble the reachability oracle's :class:`LiveConfig` from the target.

Every fact comes from a surface the target actually answered. A surface
it did not answer leaves its fact UNKNOWN, and the params that needed it
say so — none of them is refused on the strength of it.

The target's own chat-templated token count for ``text``.

The tokenize surface renders the same template the chat route renders,
so the count is the prefill the engine will actually see rather than a
word count wearing a token count's name.

Completion tokens carried by one streamed chunk, or None if unstated.

``token_count`` is arbi-serve's ``stream_options.include_token_count``
extension and rides as an extra field on the chunk. ``None`` means the
target does not carry it, which is a different fact from a chunk that
carried no new token, and the caller resolves the two differently.

Tokens per stamp, and the name of where the counts came from.

A target that answered ``token_count`` states them. One that did not
leaves only the total, which is spread evenly over the chunks: a chunk
is neither a token nor reliably a step, so an even share is the only
unbiased apportionment available and it still makes the window's
denominator TOKENS. With no total either, one token per chunk is the
pre-speculation shape and the last thing left to assume.

One measured request: TTFT and steady-state TPOT from the same stream.

Sampling is left exactly as served — no sampler fields, thinking regime
on, EOS respected with ``max_tokens`` only as a bound — and the prefix
cache is bypassed so both arms do identical full-prefill work.

``include_token_count`` asks the target how many tokens each chunk
carried. A chunk is one engine wakeup: speculation puts a whole
accepted step in one and an off-loop detokenizer coalesces several,
so a stamp counted as a token would divide by the grouping — and the
grouping is what the flags under test move. :func:`_stamp_weights`
resolves a target that will not state it.

Can this target actually score prompt logprobs, right now?

Asked by SENDING one, never by reading a capability off a schema. The
surface is declared on every build and refused at runtime on most of
them: it needs a boot flag that reserves scratch out of KV capacity,
and it scores nothing at all while a drafter is armed. Both refusals
are things an operator has to act on, so each returns the action that
clears it rather than a bare False.

Per-position top-K distributions for one prompt, teacher-forced.

``max_tokens=0`` with ``echo`` scores the prompt and generates
nothing, so both arms are scored over byte-identical input and the
positions align by construction rather than by a matching step.

The pool's total KV token capacity, or ``None`` when unavailable.

Sized at boot from whatever the card had left, so no static source has
it — a client can only learn it from the running server.

One failure's reason, in a form that is never empty.

``str()`` of a transport timeout is the empty string, and an empty reason
tested for truth reads as "nothing went wrong" — which is how a run that
measured nothing came to report itself as a clean result. The exception
type is therefore always part of the reason.

A drafter switch the target accepted but could not actually serve.

Raised so the run aborts at the switch instead of measuring whatever
the engine does afterwards. A capture-affecting swap that fails part
way can leave the engine unable to serve the PREVIOUS configuration
too, so the numbers behind a failed switch are not merely wrong for
the new arm — they describe a broken engine.

Error text if the target cannot currently generate, else ``None``.

One short generation, run after every switch. The config-override
call returning 200 says the request was accepted; only a token coming
back says the member it built can serve.

Hardware telemetry on the request swimlane's own time axis.

The swimlane (``GET /v1/admin/request_timeline`` → ``rows[].start_ms`` /
``end_ms``, anchored on the payload's ``now_ms``) is drawn against UNIX
wall-clock milliseconds taken from the ENGINE process
(``arbi_serve.server.request_timeline.now_s`` is ``time.time()``). Anything
meant to be read straight down that axis — "four requests were in flight
here; what was the GPU doing?" — has to arrive on the same epoch, in the
same unit, or the answer is a guess.

WHY PROMETHEUS AND NOT A SAMPLER
--------------------------------
The gauges this draws already exist and are already scraped: dcgm-exporter
publishes the device counters, and arbi-serve publishes its own memory
gauges through OTEL. Prometheus stamps every sample with its own UNIX
wall clock, which is the same epoch the engine stamps requests with — so
the two are directly comparable, and the only thing between them is NTP
skew between two hosts. That skew is not assumed to be zero: it is
measured against the engine's own ``now_ms`` and reported in every
response (:func:`clock_report`), so a real offset shows up as a number
rather than as a silently mis-registered overlay.

The alternative — a poller inside the console that samples NVML while a
browser tab is open — would produce a series that only exists while
someone is watching it, and would be stamped with the CONSOLE's clock
rather than the engine's. Both are disqualifying.

WHY AN ABSENT SERIES IS NOT A ZERO
----------------------------------
Not every counter exists on every device. DRAM-bandwidth activity is a
DCGM profiling (DCP) field, and GeForce-class parts do not expose it, so
on such a box the honest answer is "this device cannot report that",
not a flat line at the bottom of the chart — which is indistinguishable
from an idle memory system. Every series therefore carries
``available``; an unavailable one carries a ``reason`` and NO points at
all, and :func:`capabilities` lets the UI find that out before it offers
a toggle.

The vector for one series: a gauge as-is, a counter as its rate.

The rate window is the step, so each point is the exact mean over the
interval it covers rather than a sample taken somewhere inside it.

One PromQL expression covering every requested series.

``or`` unions vectors by label set, and every arithmetic operation drops
``__name__`` — so each disjunct is tagged with its own ``arbi_series``
label and scaled, which both converts the unit and makes the tag the only
thing distinguishing one disjunct from another. Without that tag the
disjuncts would collapse onto identical label sets and ``or`` would keep
only the first.

One expression means one round trip regardless of how many lines the
overlay draws.

Census of which of these gauges exist for this GPU, and how many series
each matches. Answers "absent" and "ambiguous" in the same single call:
a metric missing from the result is not exported; a count above 1 means
the selector matches several exporters and needs narrowing.

The step actually used: the caller's, widened just enough that the
window fits inside :data:`MAX_POINTS`.

Prometheus includes both endpoints of a range, so a span of ``n`` steps
yields ``n + 1`` samples — the budget is spent in intervals, not points.

Prometheus ``[unix_seconds, "value"]`` pairs → ``[unix_ms, float]``.

Milliseconds on the UNIX epoch is exactly what the swimlane's
``start_ms`` / ``end_ms`` / ``now_ms`` are, so a point can be mapped to x
by the same expression a request bar is.

Turn one ``query_range`` result into the per-series payload.

A LIST, in catalogue order, and the same shape the capability report
hands back: a client reads the catalogue and then asks for data, so two
shapes across those two calls is a special case on every consumer, and
the one that already bit is a caller that indexed the list and sent
array positions back as ``series=``. Every entry carries its own ``id``,
so a key in front of it says nothing the entry does not; an ordered list
additionally keeps the catalogue's order, which a JSON object does not
promise.

``presence`` is the metric census (:func:`build_presence_query`), or
``None`` when it was not needed or could not be taken — the distinction
matters, because "I know this metric is absent" and "I could not find out
why this series is empty" are different reasons and are reported as
different text.

How far apart the two clocks that stamp this overlay actually are.

Both are UNIX wall clock, so they are the same epoch by construction and
no conversion is possible or needed. What is NOT guaranteed is that the
engine host and the Prometheus host agree to the millisecond, and a
disagreement shifts every hardware point sideways relative to the request
bars. Measured rather than assumed, and reported either way; when the
target cannot be read the skew is reported as unknown, never as zero.

Which of the catalogue's series this box can actually draw, and why not
for the rest — so the overlay labels a toggle it can honour instead of
offering one that yields an empty chart.

One Prometheus instant query (plus the cached clock probe). Meant to be
read once when the overlay mounts, not on the poll.

Hardware series over ``[from_ms, to_ms]`` on the swimlane's clock.

``from_ms``/``to_ms`` are UNIX epoch milliseconds — the same numbers the
swimlane's ``start_ms``/``end_ms`` carry — and every returned point is
stamped the same way, so the overlay maps a point to x with the swimlane's
own expression and nothing needs to be reconciled at render time.

No ``now`` is returned. The caller already holds the authoritative one —
the swimlane's own ``now_ms`` — and a second, differently-aged "now" here
would be a third clock to reconcile. ``clock`` carries the measured skew
and the instant it was measured, so its age is visible.

Steady-state cost is one Prometheus ``query_range`` per call whatever the
number of series (:func:`build_range_query`); the clock and metric-census
probes behind it are cached for :data:`_PROBE_TTL_S`.

Side surfaces the audio playground needs that are not the target.

The MCP endpoint and the tide lookup are third-party services the browser
cannot reach directly (no permissive CORS, and the API key must not leave
this process), so they are proxied here rather than called from the page.

Expose the served voice default so the editor can show it before connect.

The import is deferred AND optional: the console ships as a lean HTTP
client without torch or numpy, which :mod:`arbi_serve.realtime.session`
needs. Deferring alone only moves the failure to call time, where it
surfaces as a 500 on a console that is working exactly as designed, so the
absent dependency is reported as the capability it is.

The canonical workload catalogue and the runner that drives it.

A workload is defined once, server-side, so the shape an operator ran is
recorded with the numbers rather than reconstructed from whatever the
page happened to send.

Iterate ``stream`` and close it on every exit, cancellation included.

``contextlib.aclosing`` would do it for an async generator; an OpenAI
``AsyncStream`` is not one, so its own ``close()`` is called -- and the
``finally`` is what makes this hold under CancelledError, which is the
only path that matters here.

Stop a run this console launched, by the id its ``run_start`` event carried.

Returns ``cancelled: false`` with a reason when the id names no live run --
it has already finished, or it belongs to a different console process --
rather than reporting a success that stopped nothing.

The ``X-Request-ID`` sent for request ``idx`` of one workload run.

``<prefix>-<workload>-<run_id>-<idx>``: ASCII, no whitespace, so it is a
legal header field value and a legal HTTP token. The workload key is
sanitised because a separator inside it would blur the fields apart.
The engine carries this value onto the request's timeline rows
(``client_request_id``), so a bar is attributable to one launched request
rather than inferred from timing overlap.

Create sinusoidal timestep embeddings.
:param t: a 1-D Tensor of N indices, one per batch element.
                  These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an (N, D) Tensor of positional embeddings.

Capture the chunk graph at an arbitrary set of chunk sizes.

Generalizes :meth:`_init_cuda_graph_chunk` (which hardcodes the
reference {30,48,96}) so the arbi-serve streaming path — whose
effective mel chunk size is 25 codes * up_rate = 50, not a
reference bucket — can engage the graph. ``max_size`` bounds the
padded KV length; it must exceed the driver's capped att-cache
length (prompt_mels + 100).

Args:
    x: shape (b, dt, c)
    mu: shape (b, dt, c)
    t: shape (b,)
    spks: shape (b, c)
    cond: shape (b, dt, c)
    cnn_cache: shape (depth, b, c1+c2, 2)
    att_cache: shape (depth, b, nh, t, c * 2)

Args:
    token: shape (b, t), with look ahead tokens
    mel: shape (b, t, c), groundtruth mel
    spk: shape (b, 192), speaker embedding
Returns:
    cache: dict {
        'conformer': {'cnn_cache': xxx, 'att_cache': xxx},
        'estimator': {'cnn_cache': xxx, 'att_cache': xxx}
    }

Fixed euler solver for ODEs.
Args:
    x (torch.Tensor): random noise
    t_span (torch.Tensor): n_timesteps interpolated
        shape: (n_timesteps + 1,)
    mu (torch.Tensor): output of encoder
        shape: (batch_size, n_feats, mel_timesteps)
    mask (torch.Tensor): output_mask
        shape: (batch_size, 1, mel_timesteps)
    spks (torch.Tensor, optional): speaker ids. Defaults to None.
        shape: (batch_size, spk_emb_dim)
    cond: Not used but kept for future purposes

Fixed euler solver for ODEs.
Args:
    x (torch.Tensor): random noise
    t_span (torch.Tensor): n_timesteps interpolated
        shape: (n_timesteps + 1,)
    mu (torch.Tensor): output of encoder
        shape: (batch_size, n_feats, mel_timesteps)
    mask (torch.Tensor): output_mask
        shape: (batch_size, 1, mel_timesteps)
    spks (torch.Tensor, optional): speaker ids. Defaults to None.
        shape: (batch_size, spk_emb_dim)
    cond: Not used but kept for future purposes
    cnn_cache: shape (n_time, depth, b, c1+c2, 2)
    att_cache: shape (n_time, depth, b, nh, t, c * 2)

Args:
    mu(torch.Tensor): shape (b, c, t)
    spks(torch.Tensor): shape (b, 192)
    cond(torch.Tensor): shape (b, c, t)
    cnn_cache: shape (n_time, depth, b, c1+c2, 2)
    att_cache: shape (n_time, depth, b, nh, t, c * 2)

Transform query, key and value.

Args:
    query (torch.Tensor): Query tensor (#batch, time1, size).
    key (torch.Tensor): Key tensor (#batch, time2, size).
    value (torch.Tensor): Value tensor (#batch, time2, size).

Returns:
    torch.Tensor: Transformed query tensor, size
        (#batch, n_head, time1, d_k).
    torch.Tensor: Transformed key tensor, size
        (#batch, n_head, time2, d_k).
    torch.Tensor: Transformed value tensor, size
        (#batch, n_head, time2, d_k).

Compute attention context vector.

Args:
    value (torch.Tensor): Transformed value, size
        (#batch, n_head, time2, d_k).
    scores (torch.Tensor): Attention score, size
        (#batch, n_head, time1, time2).
    mask (torch.Tensor): Mask, size (#batch, 1, time2) or
        (#batch, time1, time2), (0, 0, 0) means fake mask.

Returns:
    torch.Tensor: Transformed value (#batch, time1, d_model)
        weighted by the attention score (#batch, time1, time2).

Compute scaled dot product attention.

Args:
    query (torch.Tensor): Query tensor (#batch, time1, size).
    key (torch.Tensor): Key tensor (#batch, time2, size).
    value (torch.Tensor): Value tensor (#batch, time2, size).
    mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
        (#batch, time1, time2).
        1.When applying cross attention between decoder and encoder,
        the batch padding mask for input is in (#batch, 1, T) shape.
        2.When applying self attention of encoder,
        the mask is in (#batch, T, T)  shape.
        3.When applying self attention of decoder,
        the mask is in (#batch, L, L)  shape.
        4.If the different position in decoder see different block
        of the encoder, such as Mocha, the passed in mask could be
        in (#batch, L, T) shape. But there is no such case in current
        CosyVoice.
    cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
        where `cache_t == chunk_size * num_decoding_left_chunks`
        and `head * d_k == size`


Returns:
    torch.Tensor: Output tensor (#batch, time1, d_model).
    torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
        where `cache_t == chunk_size * num_decoding_left_chunks`
        and `head * d_k == size`

Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
    query (torch.Tensor): Query tensor (#batch, time1, size).
    key (torch.Tensor): Key tensor (#batch, time2, size).
    value (torch.Tensor): Value tensor (#batch, time2, size).
    mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
        (#batch, time1, time2), (0, 0, 0) means fake mask.
    pos_emb (torch.Tensor): Positional embedding tensor
        (#batch, time2, size).
    cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
        where `cache_t == chunk_size * num_decoding_left_chunks`
        and `head * d_k == size`
Returns:
    torch.Tensor: Output tensor (#batch, time1, d_model).
    torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
        where `cache_t == chunk_size * num_decoding_left_chunks`
        and `head * d_k == size`

Relative positional encoding module (new implementation).

Details can be found in https://github.com/espnet/espnet/pull/2816.

See : Appendix B in https://arxiv.org/abs/1901.02860

Args:
    d_model (int): Embedding dimension.
    dropout_rate (float): Dropout rate.
    max_len (int): Maximum input length.

For getting encoding in a streaming fashion

Attention!!!!!
we apply dropout only once at the whole utterance level in a none
streaming way, but will call this function several times with
increasing input size in a streaming scenario, so the dropout will
be applied several times.

Args:
    offset (int or torch.tensor): start offset
    size (int): required size of position encoding

Returns:
    torch.Tensor: Corresponding encoding

Encoder layer module.
Args:
    size (int): Input dimension.
    self_attn (torch.nn.Module): Self-attention module instance.
        `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
        instance can be used as the argument.
    feed_forward (torch.nn.Module): Feed-forward module instance.
        `PositionwiseFeedForward` instance can be used as the argument.
    feed_forward_macaron (torch.nn.Module): Additional feed-forward module
         instance.
        `PositionwiseFeedForward` instance can be used as the argument.
    conv_module (torch.nn.Module): Convolution module instance.
        `ConvlutionModule` instance can be used as the argument.
    dropout_rate (float): Dropout rate.
    normalize_before (bool):
        True: use layer_norm before each sub-block.
        False: use layer_norm after each sub-block.
    enable_cuda_graph (bool): Control whether to enable CUDA Graph.

Compute encoded features.

Args:
    x (torch.Tensor): (#batch, time, size)
    mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
        (0, 0, 0) means fake mask.
    pos_emb (torch.Tensor): positional encoding, must not be None
        for ConformerEncoderLayer.
    mask_pad (torch.Tensor): batch padding mask used for conv module.
        (#batch, 1，time), (0, 0, 0) means fake mask.
    att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
        (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
    cnn_cache (torch.Tensor): Convolution cache in conformer layer
        (#batch=1, size, cache_t2)
Returns:
    torch.Tensor: Output tensor (#batch, time, size).
    torch.Tensor: Mask tensor (#batch, time, time).
    torch.Tensor: att_cache tensor,
        (#batch=1, head, cache_t1 + time, d_k * 2).
    torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).

Positionwise feed forward layer.

FeedForward are appied on each position of the sequence.
The output dim is same with the input dim.

Args:
    idim (int): Input dimenstion.
    hidden_units (int): The number of hidden units.
    dropout_rate (float): Dropout rate.
    activation (torch.nn.Module): Activation function

Input x.

Args:
    x (torch.Tensor): Input tensor (#batch, time, idim).
    x_mask (torch.Tensor): Input mask (#batch, 1, time).

Returns:
    torch.Tensor: linear input tensor (#batch, time', odim),
        where time' = time .
    torch.Tensor: linear input mask (#batch, 1, time'),
        where time' = time .

A 1D upsampling layer with an optional convolution.

Parameters:
    channels (`int`):
        number of channels in the inputs and outputs.
    use_conv (`bool`, default `False`):
        option to use a convolution.
    use_conv_transpose (`bool`, default `False`):
        option to use a convolution transpose.
    out_channels (`int`, optional):
        number of output channels. Defaults to `channels`.

Args:
    xs: shape (b, dt, c)
    last_chunk: bool. If last chunk, will pad input with lookaheads
    att_cache: shape (depth1+depth2, b, nh, 2*t1, c).
    cnn_cache: shape (b, c, t1+t2). Where t1=2 (pre_lookahead_layer), t2=4 (up_layer)

Perform padding for the list of tensors.

Args:
    xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)].
    pad_value (float): Value for padding.

Returns:
    Tensor: Padded tensor (B, Tmax, `*`).

Examples:
    >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)]
    >>> x
    [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])]
    >>> pad_list(x, 0)
    tensor([[1., 1., 1., 1.],
            [1., 1., 0., 0.],
            [1., 0., 0., 0.]])

Make mask tensor containing indices of padded part.

See description of make_non_pad_mask.

Args:
    lengths (torch.Tensor): Batch of lengths (B,).
Returns:
    torch.Tensor: Mask tensor containing indices of padded part.

Examples:
    >>> lengths = [5, 3, 2]
    >>> make_pad_mask(lengths)
    masks = [[0, 0, 0, 0 ,0],
             [0, 0, 0, 1, 1],
             [0, 0, 1, 1, 1]]

Implementation of a sine-based periodic activation function
Shape:
    - Input: (B, C, T)
    - Output: (B, C, T), same shape as the input
Parameters:
    - alpha - trainable parameter
References:
    - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
    https://arxiv.org/abs/2006.08195
Examples:
    >>> a1 = snake(256)
    >>> x = torch.randn(256)
    >>> x = a1(x)

Args:
    in_features: shape of the input
    alpha: trainable parameter
    alpha_trainable: whether alpha is trainable
    alpha_logscale: whether to use log scale for alpha
        alpha is initialized to 1 by default, higher values = higher-frequency.
        alpha will be trained along with the rest of your model.

Definition of sine generator
SineGen(samp_rate, harmonic_num = 0,
        sine_amp = 0.1, noise_std = 0.003,
        voiced_threshold = 0,
        flag_for_pulse=False)
samp_rate: sampling rate in Hz
harmonic_num: number of harmonic overtones (default 0)
sine_amp: amplitude of sine-wavefrom (default 0.1)
noise_std: std of Gaussian noise (default 0.003)
voiced_thoreshold: F0 threshold for U/V classification (default 0)
flag_for_pulse: this SinGen is used inside PulseGen (default False)
Note: when flag_for_pulse is True, the first time step of a voiced
    segment is always sin(np.pi) or cos(0)

SourceModule for hn-nsf
SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
             add_noise_std=0.003, voiced_threshod=0)
sampling_rate: sampling_rate in Hz
harmonic_num: number of harmonic above F0 (default: 0)
sine_amp: amplitude of sine source signal (default: 0.1)
add_noise_std: std of additive Gaussian noise (default: 0.003)
    note that amplitude of noise in unvoiced is decided
    by sine_amp
voiced_threshold: threhold to set U/V given F0 (default: 0)
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
F0_sampled (batchsize, length, 1)
Sine_source (batchsize, length, 1)
noise_source (batchsize, length 1)
uv (batchsize, length, 1)

Definition of sine generator
SineGen(samp_rate, harmonic_num = 0,
        sine_amp = 0.1, noise_std = 0.003,
        voiced_threshold = 0,
        flag_for_pulse=False)
samp_rate: sampling rate in Hz
harmonic_num: number of harmonic overtones (default 0)
sine_amp: amplitude of sine-wavefrom (default 0.1)
noise_std: std of Gaussian noise (default 0.003)
voiced_thoreshold: F0 threshold for U/V classification (default 0)
flag_for_pulse: this SinGen is used inside PulseGen (default False)
Note: when flag_for_pulse is True, the first time step of a voiced
    segment is always sin(np.pi) or cos(0)

SourceModule for hn-nsf
SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
             add_noise_std=0.003, voiced_threshod=0)
sampling_rate: sampling_rate in Hz
harmonic_num: number of harmonic above F0 (default: 0)
sine_amp: amplitude of sine source signal (default: 0.1)
add_noise_std: std of additive Gaussian noise (default: 0.003)
    note that amplitude of noise in unvoiced is decided
    by sine_amp
voiced_threshold: threhold to set U/V given F0 (default: 0)
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
F0_sampled (batchsize, length, 1)
Sine_source (batchsize, length, 1)
noise_source (batchsize, length 1)
uv (batchsize, length, 1)

NemotronLabs-VoiceChat-11B RVQ-VAE audio codec — encode + decode.

Ports NVIDIA's ``RVQVAEModel``
(``nemo.collections.speechlm2.modules.ear_tts_vae_codec.RVQVAEModel`` — the
class ``duplex_ear_tts.py`` actually imports and runs at serving time; the
sibling ``rvq_ear_tts_vae.py`` module in the same package is an unused
duplicate and is not the source of this port), both directions: raw
waveform → discrete multi-codebook indices (encode, ``Wav2Latent`` +
the quantizer's forward residual-VQ snap) and indices → 22.05 kHz
waveform (decode, already ported). The encode direction exists to
support offline speaker-voice-conditioning prep
(the voice preparation step) — TTS generation
itself only ever runs decode.

Architecture, verified against ``config.json``'s
``model.speech_generation.model.codec_config`` and the real checkpoint's
``tts_model.audio_codec.*`` tensor shapes (214 tensors: 76 ``encoder.*`` +
76 ``decoder.*`` + 62 ``prvq.*``; this module implements all 214):

* ``num_quantizers=31``, ``codebook_size=1024``, ``latent_size=512``.
  Every ``prvq.mus_list.{0..30}`` tensor in the checkpoint is exactly
  ``[1024, 512]`` — there is no reserved/mask row and no
  ``[num_quantizers, codebook_size+1, latent_size]`` buffer anywhere in the
  codec (dead or live module). ``prvq._variance_list.{i}.variance`` is a
  scalar EMA-variance per quantizer level, unrelated to codebook width; the
  decode path does not read it, it is only carried so a strict
  ``load_state_dict`` against the real checkpoint succeeds.
* ``n_fft=16``, ``hop_length=4``, ``base_hidden_size=384``,
  ``channel_mult=(1, 2, 4)``, ``rates=(7, 7, 9)`` (encoder order; the
  decoder reverses both). ``wav_to_token_ratio=1764 = hop_length *
  rates[0] * rates[1] * rates[2]`` samples per code frame. The encoder's
  three ``Conv1d`` downsample stages (stride == kernel_size, "valid"
  convolution) turn ``rates[0]*rates[1]*rates[2]`` STFT spectral frames
  into one code frame; the decoder's three ``ConvTranspose1d`` stages do
  the exact inverse, and the STFT's own forward/inverse hop is
  ``hop_length`` samples per spectral frame either way.
* The decoder's last layer (``decoder.layers.12``, checkpoint shape
  ``[18, 384, 1]``) projects to ``n_fft + 2 = 18`` channels, split 9/9 into
  magnitude and phase; waveform synthesis is an inverse STFT — Hann-window
  overlap-add built from ``torch.fft.irfft`` + ``F.fold``
  (:func:`_spec_to_wav`), not ``torch.istft``. The encoder's first layer
  (``encoder.layers.0``, checkpoint shape ``[384, 18, 1]``) is the mirror
  projection FROM 18 channels — but those 18 channels are the forward
  STFT's raw real/imaginary parts (:func:`_wav_to_spec` +
  ``torch.view_as_real``), not a magnitude/phase parameterization; the
  encoder and decoder use two different 18-channel conventions on purpose
  (verified from the reference source, not a port inconsistency).
* The codec module carries no sample-rate constant; ``target_sample_rate``
  is a caller-side setting (``config.json``'s ``data.target_sample_rate`` =
  22050 Hz), reproduced here as the driver class's ``SAMPLE_RATE``.

Two call surfaces, mirroring :class:`arbi_serve.audio.token2wav.Token2Wav`:

* :meth:`NemotronAudioCodec.synthesize` / ``synthesize_waveform`` — whole-turn
  codes → audio, one-shot (decode).
* :meth:`NemotronAudioCodec.encode_waveform` — whole-turn waveform → codes,
  one-shot (encode); used offline by the voice-prep script, not on the
  per-token generation hot path.
* :meth:`NemotronAudioCodec.open_stream` / :meth:`StreamHandle.feed` /
  ``feed_waveform`` — chunked decode with an internal lock serializing
  access to the shared decoder, replicating ``duplex_ear_tts.py``'s
  ``decode_one_audio_step``: each call decodes a trailing window of at most
  ``number_prev_tokens`` code frames from scratch (no persistent conv
  cache — the reference always calls ``decode(cache=None)`` here) and
  returns only the newest frames' worth of waveform, discarding the rest of
  the window's output.

Forward STFT via zero-padding + ``torch.stft``: mono waveform → complex spectrogram.

Mirrors :func:`_spec_to_wav`'s manual centering (``center=False`` plus
explicit edge padding, so encode/decode framing agrees) but pads with
zeros rather than reflecting — matches the reference's own forward
``spectrogram()``, which is plain ``F.pad`` with no ``mode=`` override.

Per-layer trailing-context cache for incremental causal-conv decode.

Holds each ``_ConvNeXt1d``'s ``kernel_size - 1`` trailing activations plus
the decoder's trailing spectrogram frames, so successive one-frame decode
calls produce SAMPLE-CONTINUOUS waveform instead of each chunk restarting
from zero-padded convolutions and a fresh iSTFT overlap-add. This is what
NVIDIA's own production streaming server threads per stream
(``use_codec_cache=True`` by default,
``nemotron_voicechat_inference_wrapper.py``); the per-tick TTS path
(:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.generate_tts_frame`)
threads one of these per session for the same reason — see design doc
§7.20.

Not used by :class:`StreamHandle`, which replicates the reference's OTHER
(fallback) streaming shape: ``cache=None`` plus a windowed *code* history.

``codes``: ``[B, T, num_quantizers]`` long. Swaps every per-quantizer
entry equal to ANY of ``control_codes`` (BOS/EOS/PAD sentinels) for
``silence_tokens``' matching-depth entry; ordinary revealed codes pass
through untouched. Direct port of the reference's
``replace_control_speech_codes`` (``duplex_ear_tts.py``,
``silence_tokens is not None`` branch).

Lives beside the codec rather than beside the turn driver because the
reference applies it at the decode call site unconditionally, on the
codes about to enter :meth:`NemotronAudioCodecDecoder.decode` — never
on the codes fed back as the next step's autoregressive input.

Latent frames → waveform: the ``decoder`` half of the RVQ-VAE vocoder.

``self.layers`` is a flat :class:`~torch.nn.ModuleList` matching the
checkpoint's ``decoder.layers.{0..12}`` indices one-for-one: each of the
3 upsample stages is one ``ConvTranspose1d`` (stride == kernel_size ==
the stage's rate) followed by ``num_blocks`` ``_ConvNeXt1d`` blocks, and
a final ``Conv1d`` projects to ``n_fft + 2`` spectrogram channels.

Waveform → latent frames: the ``encoder`` half of the RVQ-VAE vocoder.

``self.layers`` is a flat :class:`~torch.nn.ModuleList` matching the
checkpoint's ``encoder.layers.{0..12}`` indices one-for-one, in the
checkpoint's own (non-reversed) ``channel_mult``/``rates`` order —
the mirror image of :class:`NemotronAudioLatentDecoder`'s layout with
the per-stage order swapped: here each of the 3 downsample stages runs
``num_blocks`` ``_ConvNeXt1d`` blocks FIRST, then one strided
``Conv1d`` (stride == kernel_size == the stage's rate, "valid"
convolution); the decoder upsamples first and blocks after.

Full RVQ-VAE codec: waveform ↔ discrete multi-codebook indices.

Submodules are named ``encoder``/``decoder``/``prvq`` to match the
checkpoint's ``tts_model.audio_codec.{encoder,decoder,prvq}.*`` keys
one-for-one (214 tensors total) — see :meth:`load_from_full_state_dict`.
Despite the ``*Decoder`` name, the class implements both
directions.

RVQ-VAE codec driver: whole-turn one-shot encode/decode + streaming windowed decode.

Structure mirrors :class:`arbi_serve.audio.token2wav.Token2Wav` — a
lock-serialized synthesizer plus a :class:`StreamHandle` for chunked
decode — so a later integration pass can wire it into the same
off-critical-path ``asyncio.to_thread`` pattern.

One request's chunked decode state — replicates ``decode_one_audio_step``.

Each ``feed`` appends the new code frames to the running history,
decodes a window of at most ``number_prev_tokens`` trailing frames from
scratch (``cache=None``, matching the reference), and returns only the
newest frames' worth of waveform, discarding the rest of the window's
output. ``number_prev_tokens=None`` (or 0) decodes the full history on
every call.

``z``: ``[B, T, latent_size]`` float → ``[B, T, num_quantizers]`` long.

Per-depth nearest-neighbor residual quantization: at each depth,
snaps the running residual to its closest codebook entry (squared
Euclidean distance) and subtracts that entry before the next
depth. This is the codec's own native encode-time quantizer —
distinct from ``mog_head.depthsum_encoding_step``'s iterative
CFG/MoG sampling snap used during TTS *generation*, which reuses
the *decode*-side embedding table but never runs this loop.

``wav``: ``[B, T]`` or ``[B, 1, T]`` float, ``T`` a multiple of
``wav_to_token_ratio`` → (``[B, T // wav_to_token_ratio, num_quantizers]``
long codes, code lengths). Offline/prep-time use — not on the
per-token generation hot path (see module docstring).

Admission control for concurrent token2wav speech-out streams.

The vocoder's memory is bounded by ``ARBI_AUDIO_T2W_RESERVE_GB``. Each open
stream carries its own flow chunk cache and HiFT overlap caches, on top of a
resident base (weights + DiT cudagraph). So N concurrent streams need
``base + N x per_stream``.

Admitting more streams than the reserve can hold buys nothing anyway:
:class:`~arbi_serve.audio.token2wav.Token2Wav` serializes every decode behind
its own lock, so a second stream cannot vocode while the first is vocoding.
It only spends memory. This module turns that memory over-subscription into
a queue: a turn waits for a slot (its audio starts later) instead of OOM-ing
mid-SSE.

Capacity is ``ARBI_AUDIO_MAX_CONCURRENT_SPEECH`` (0 = auto). Auto derives it
from the reserve the operator actually granted::

    N = floor((reserve_gb - BASE_GIB) / PER_STREAM_GIB), clamped to >= 1

so raising the reserve raises concurrency, and the two knobs cannot drift.

Text generation is untouched — only the vocoder is gated. A speech turn holds a
slot from ``open_stream`` until its last PCM chunk (or until the client
disconnects and the generator is closed), never across the queue wait.

The process-wide slot semaphore, bound to the running loop.

Built lazily: capacity reads runtime flags, and an ``asyncio.Semaphore``
must be created on (or at least used from) the loop that awaits it. The API
child serves every speech turn on one uvicorn loop, so one instance suffices;
we rebuild if the loop identity ever changes (tests create fresh loops).

Async context manager holding one speech-stream slot.

``async with SpeechSlot():`` around the whole open_stream..last-chunk
lifetime. Releasing is unconditional (including on GeneratorExit when a
client disconnects mid-stream), which is the property that keeps the bound
honest under churn.

The ``scratch.audio_t2w`` pool — a bounded home for the speech vocoder.

The token2wav decoder (CosyVoice2 flow + HiFT) is a separate torch module that
the engine's KV profiler never sees. Under process mode it does not even live
in the engine process: the routes — and therefore the vocoder — run in the API
child, which has its own CUDA context and its own caching allocator. The engine
meanwhile locks ``gpu_memory_utilization`` of the device.

``ARBI_AUDIO_T2W_RESERVE_GB`` shrinks the engine's KV pool by that much, and
ordering (wired in ``server/app.py``) makes those bytes actually reach the
vocoder:

    engine builds -> profiles -> carves the reserve out of KV
    -> locks its plan ("memory locked — won't grow while serving")
    -> ReadyMsg -> the child builds the one vocoder into that hole

Only then are the reserved bytes something nothing else can take.

This module additionally routes every token2wav allocation — weights, voice
caches, the DiT capture, and the per-utterance transient — into a named
:class:`NamedMemPool` (``scratch.audio_t2w``). That buys isolation from
allocator fragmentation and an honest ``pool="scratch.audio_t2w"`` metric label
instead of the footprint hiding in ``unpooled.unregistered_pool``.

Does ``model_path`` carry a ``token2wav/`` weights directory?

The KV-pool profiler's/grow-floor's ``audio_t2w_reserve_gb`` carve-out
(``arbi_serve/engine/profile.py``, ``arbi_serve/engine/inprocess_capture.py``)
only matters for a served model that can ever build a token2wav decoder
in the first place — the same directory check
:func:`arbi_serve.server.engine_boot._preload_audio_decoder` already
uses as the actual gate on whether the decoder is built at all. A model
without this directory (e.g. NemotronVoiceChat, whose speech output is
its own in-process RVQ codec, not a separate token2wav pool/process)
never builds one, so reserving VRAM for it out of the KV pool would
shrink serving capacity for a subsystem that will never run —
``enable_audio`` alone (the input-tower-AND-speech-output modality
switch) cannot tell the two cases apart.

Create the ``scratch.audio_t2w`` pool, or ``None`` when not applicable.

``None`` for a CPU decoder or a non-positive reserve — callers then fall
back to the default allocator (untracked, unbounded: the old behaviour).

Longest utterance a single turn can produce, from the token budget.

Sizes the check on the worst reply this server can generate, rather than
a fixed duration, so it answers the question that matters.

Report, at boot, whether the vocoder can actually run.

Does not try to preallocate the reserve as one block: the weights are
already resident in the pool by the time we get here, so claiming
``reserve_gb`` on top of them asks for far more than the reserve. Freeing
such a block would hand the bytes back to the driver anyway.

What makes the reserve real is ordering, not preallocation: the engine
locks its plan first ("memory locked — won't grow while serving") having
carved ``reserve_gb`` out of its KV pool, and only then does the API
child build the vocoder into that hole. So the honest thing to do here
is verify the hole is big enough and say so.

Call this after ``warmup()``: before it, the DiT cudagraph capture
has not happened yet and the free-memory read is optimistic.

The transient is sized for the longest reply this server can generate
(max_tokens / 25 Hz).

Returns True when there is room for a worst-case vocode transient.

Build a :class:`Token2Wav` whose every allocation lives in the pool.

Loads the weights inside ``scratch.audio_t2w``, warms every prepared voice
(their caches belong in the pool too, not on the default heap), then checks
that the reserve still has room for a vocode transient. Returns the decoder;
the pool hangs off ``decoder._mem_pool`` for metrics.

CosyVoice2-style token2wav — speech codes → 24 kHz waveform.

Slim serving driver over the vendored CosyVoice2 flow-matching decoder
(:mod:`arbi_serve.audio.cosyvoice2`) and the HiFT vocoder
(:mod:`arbi_serve.audio.hifigan`), adapted from the Step-Audio-2
reference ``token2wav.py`` (Apache-2.0) with the deployment-hostile
pieces removed:

  * **No hyperpyyaml** — the flow module is constructed directly; the
    scalar hyperparameters are read from ``flow.yaml`` with a minimal
    text extractor and cross-checked by the ``strict=True`` state-dict
    load (a geometry mismatch fails loudly at boot, never silently).
  * **No onnxruntime / s3tokenizer / kaldi at serve time** — the
    speaker-prompt conditioning tensors (speech tokens, x-vector,
    prompt mels) depend only on the voice's prompt wav, so they are
    precomputed once by the voice preparation tool
    (which runs the ONNX models in a tooling env) and loaded here as a
    ``.pt`` per voice from the voices dir.

Two call surfaces, mirroring the reference:

  * :meth:`synthesize` — whole-turn codes → WAV bytes (Phase-3
    non-streaming path).
  * :meth:`open_stream` / :meth:`StreamHandle.feed` /
    :meth:`StreamHandle.finish` — chunked flow inference (25-code
    chunks + ``pre_lookahead_len`` lookahead) with the HiFT mel/source
    cache and Hamming-window overlap smoothing, yielding PCM16 chunks
    for SSE streaming.

Extract the scalar hyperparameters from the hyperpyyaml file.

The file is a hyperpyyaml object graph we deliberately do not
execute; the scalars we need are plain ``key: value`` lines. Any
key not found keeps the shipped default — and a genuinely wrong
value cannot mis-serve, because the ``strict=True`` state-dict load
fails on the resulting shape mismatch.

Speech codes → waveform, with per-voice prompt conditioning.

``model_dir`` is the checkpoint's ``token2wav/`` directory
(``flow.yaml`` + ``flow.pt`` + ``hift.pt``); voices load from
``voices_dir`` (default ``<model_dir>/voices``), one
``<name>.pt`` per voice as written by
the voice preparation tool.

Thread safety: one decode at a time (an internal lock) — the
Phase-3 sidecar serializes turns; per-request streaming state lives
on the :class:`StreamHandle`.

One request's chunked token2wav state.

Feed 25-code chunks (+3 lookahead codes handled by the caller's
buffering); each ``feed`` returns a PCM16 byte chunk with the HiFT
overlap tail held back; ``finish`` flushes the tail.

Bound the conformer att-cache to the last ``w_frames`` token-rate frames.

The cache packs two rates on dim 3: the first ``depth1`` layers
(``encoders``) use the token-rate half ``[:offset1]`` (their real data,
duplicated by a ``repeat(...,2,...)``); the remaining layers (``up_encoders``)
use the full mel-rate ``2*offset1``. Slicing must respect that split — keep
the last ``w_frames`` token frames and last ``2*w_frames`` mel frames, then
repack so ``offset1`` (= dim3//2) equals ``w_frames``.

Route allocations in the block into ``scratch.audio_t2w``.

No-op when the pool is absent (CPU decoder, tests, or a caller that
opted out). ``NamedMemPool.use()`` is re-entrant and no-ops while a
cudagraph capture is routing the allocator elsewhere, so nesting it
around ``enable_cuda_graph`` is safe.

A named voice (loaded + cached from disk) OR an ephemeral triple.

A ``str`` is a prepared voice from ``voices_dir``. A mapping is an
ad-hoc conditioning triple — the live-voice-clone path: the incoming
speaker's own ``{prompt_speech_tokens, spk_emb, prompt_mels}`` built by
:class:`arbi_serve.audio.voice_prompt.VoicePromptExtractor`. NOT
cached (they are per-turn), so they carry no name and never touch disk.

Capture the DiT diffusion inner loop as a CUDA graph (opt-in).

The dominant streaming cost is the DiT/CFM Euler loop (10 steps x
CFG-batch-2 over 16 blocks), which is launch-bound. The vendored
DiT ships a static-query / padded-masked-KV capture that the
reference only wires for chunk sizes {30,48,96}; arbi-serve's
25-code chunk upsamples to ``STREAM_CHUNK_CODES * up_rate`` mel
frames (50 by default), so we capture at that size. ``max_att``
bounds the padded KV length and must exceed the driver's capped
att-cache length (``prompt_mels + 100``); the default 1000 covers
prompts up to ~450 tokens.

Idempotent. Requires CUDA. GPU-validate audio correctness (the
counting/copy/ASR greedy probes) before relying on this.

Run a few synthetic feed cycles at boot to kill cold-start.

Warms cudnn/cublas autotune + the CUDA allocator (and, when
``enable_cuda_graph`` is set, captures the DiT graph at the exact
streaming chunk size discovered here) so the first real request
does not pay one-time costs. Uses a prepared voice if given/available,
else synthesizes conditioning of the correct shapes (perf/warmup is
shape-driven; the synthetic output is not audible).

Serve-time Step-Audio-2 voice-conditioning extractor (live voice cloning).

The CosyVoice2 decoder (:class:`arbi_serve.audio.token2wav.Token2Wav`)
conditions every synthesis on a speaker prompt: its 25 Hz speech tokens
(``s3tokenizer`` ONNX), a 192-d CAM++ x-vector (``campplus`` ONNX) and its
24 kHz 80-bin mel. A *named* voice caches that triple on disk, built once by
the voice preparation tool.

This module builds the same triple on the fly from an arbitrary PCM16 clip —
the incoming speaker's own audio — so a speech-to-speech turn can come back in
their voice, no predefined voice required. It is the serve-time twin of the
prep script: same math, same ONNX models (already in the checkpoint's
``token2wav/`` dir), run per turn instead of offline.

Deps beyond the base serving image: ``s3tokenizer`` (for the 25 Hz codes) and
``onnxruntime`` (already present for the diarizer). The 24 kHz mel filterbank
comes from ``torchaudio`` (Slaney-normalized — bit-for-bit the librosa default
the prep script used) so no ``librosa`` is pulled onto the serving path. When
``s3tokenizer`` is absent the extractor raises a clear error at build time, so
clone mode fails loud rather than mis-serving.

Memory / latency: the ONNX models run on CPU (they never touch the token2wav
GPU reserve); the returned tensors are CPU tensors the decoder moves onto the
device under its own pool. The conditioning clip is capped to
``max_prompt_s`` seconds so a long user turn can neither blow the DiT
att-cache reserve (sized off ``prompt_mels``) nor stall time-to-first-audio.

Build a token2wav conditioning triple from raw PCM16, on demand.

Construct once (loads the two ONNX models lazily on first use); call
:meth:`extract` per turn. Not safe for concurrent use — one per session.

Construct an extractor from the resident model's ``token2wav/`` dir.

Returns ``None`` when the model path can't be resolved or the ONNX models
are absent — clone mode then reports it rather than crashing the session.

PCM16 + rate → ``{prompt_speech_tokens, spk_emb, prompt_mels}`` (CPU).

Raises ``ValueError`` on too-short / empty audio and propagates a clear
error if the ONNX deps are unavailable (clone mode must fail loud).

Pluggable attention-backend registry.

A backend is a triplet (descriptor + metadata builder + per-layer attn
op). Each declares its :class:`StateKind` so the engine knows which
slot in ``active_backends: dict[StateKind, AttentionBackend]`` it
fills.

Internal backend specs use the ``kind:name`` form everywhere:

  - ``paged_kv:tkv-bypass``              — raw bf16 KV on the turbo-attn kernels
                                            (Turbo prefill + split-K paged decode,
                                            cudagraph-capturable decode; whole-model,
                                            mixed head dims, head_dim up to 1024) — the
                                            production bf16 default
  - ``paged_kv:tkv-k4v4``               — turbokv codec at 4-bit K, 4-bit V
  - ``mla_shared:mla``              — MLA shared-KV
  - ``mamba:mamba``                 — Mamba SSM
  - ``gdn:gdn``                     — Gated Delta Net

The ``kind:`` prefix is required; admin endpoints reject a missing
prefix at parse time.

The concrete backend classes are loaded lazily through a module
``__getattr__`` (PEP 562), so importing this package does not pull any
heavy backend submodule (torch + the turbo-attn / tkv CUDA closures)
until a backend class is actually accessed.

User-facing CLI:
  - ``--kv-cache-dtype tkv`` (with ``TKV_BITS`` / ``TKV_K_BITS`` /
    ``TKV_V_BITS`` env) selects the TKV codec. Translated by
    :mod:`arbi_serve.cli` into the internal
    ``paged_kv:tkv-k{K}v{V}`` spec before reaching the engine.
  - ``--backend kind:name`` (repeatable) registers non-tkv backends.

Adding a backend: implement the three protocols (declaring
``state_kind()``), register a factory below.

Build a :class:`TkvBackend` from a ``tkv-k{K}v{V}`` spec string.

Single source of truth: the ``(k_bits, v_bits)`` come from ``TKV_BITS``
(via :func:`resolve_declared_tkv_bits`), never from the spec name. The
bits encoded in the name are parsed only to validate that they agree with
the declared config; a mismatch is a hard error. This is what makes the
bit-width footgun structurally impossible: a hand-written / request-supplied
spec (e.g. ``paged_kv:tkv-k2v8`` on an HTTP ``attention_backend`` field, or
a directly-built ``ServerConfig``) cannot silently override the boot-time
``TKV_BITS`` declaration, and an asymmetric ``k != v`` — which a single
``TKV_BITS`` number can never produce — is rejected outright.

Parse ``mla`` or ``mla-k{K}`` into an :class:`MlaBackend` instance.

MLA shared-KV folds K and V into one quantiser; only ``k_bits`` is
user-tunable. ``mla`` (no suffix) reads ``TKV_K_BITS`` from env
(default 4); ``mla-k<K>`` pins ``k_bits`` directly. ``K`` any integer in ``[2, 8]``.

Factory: turn a bare backend name into an instance.

Recognized:
  - ``tkv-bypass``              — raw bf16 KV on the turbo-attn kernels (Turbo prefill
                                  + split-K paged decode; all
                                  head dims) — the production bf16 default
  - ``tkv-k{K}v{V}``           — turbokv codec (K, V any integer in [2, 8])
  - ``mla`` / ``mla-k{K}``      — MLA shared-KV (DeepSeek V2/V3, Kimi K2)
  - ``mamba``                   — Mamba SSM
  - ``gdn``                     — Gated Delta Net

Build a per-layer attn op from the selected backend.

The engine binds ONE backend per :class:`StateKind`; this is a thin
pass-through that forwards the per-layer geometry to the backend's
own ``make_attn_op``. Each backend handles its own head geometry —
``tkv-bypass`` serves every head dim (including Gemma 4's
``head_dim=512`` full-attention layers) natively on the turbo-attn
kernels, so no per-layer kernel substitution is needed.

Per-head attention-sink normalisation at the backend boundary.

A sink is one extra logit per query head that contributes only to the
softmax denominator. The model owns it as an ``nn.Parameter`` of the
activation dtype (bf16 in production); the turbo-attn kernels that
consume it — ``tkv.kernels.turbo_attn_simt.tq_splitk_batch_decode``
/ ``tq_splitk_batch_mtp_decode`` (``attn_sink=``) and
``tkv.kernels.cuda.prefill.turbo_prefill_dispatch.turbo_prefill``
(``learnable_sink=``) — all take **float32**.

:func:`normalize_attention_sink` is the single conversion site: every
paged backend stores its per-layer sink through it, so the cast, the
device move, the contiguity and the per-rank head count are checked
once.

Sink-carrying variant of turbo-attn's raw-bf16 paged attend.

:class:`tkv.runtime.bypass_bf16.BypassBf16Attend` routes scatter →
split-K decode / MTP verify / Turbo prefill for the ``tkv-bypass`` codec
mode, but its three kernel call sites do not forward a per-head
attention sink even though every kernel behind them takes one:

  * ``tkv.kernels.turbo_attn_simt.tq_splitk_batch_decode(attn_sink=)``
  * ``tkv.kernels.turbo_attn_simt.tq_splitk_batch_mtp_decode(attn_sink=)``
  * ``tkv.kernels.cuda.prefill.turbo_prefill_dispatch.turbo_prefill(learnable_sink=)``
    (re-exported as ``tkv.kernels.cuda.prefill.turbo_prefill_dispatch.turbo_prefill``)

This subclass overrides exactly those three leaf methods to pass the
sink; routing, KV scatter, the fused K/V views, the q-prescale and the
decode autotune picks are inherited unchanged. It is constructed ONLY
for layers that declare :attr:`LayerSpec.attention_sinks` — sink-free
layers keep the stock class, so the production bf16 path is untouched.

The sink is fp32 ``(num_heads,)``, normalised by
:func:`arbi_serve.backends._attn_sink.normalize_attention_sink`.

Cumulative K lengths, bounded by what the block table addresses.

Each row of ``block_table`` maps ``block_table.shape[1] * page_size``
tokens, so a ``seq_lens`` entry above that bound would send the kernel
past the pages the row maps. The bound comes from the table's static
column count, not from live page occupancy, so the call carries no host
sync and stays cudagraph-capturable. The kernel reads ``int32``.

Shared base for paged-KV attention backends + their metadata builders.

:class:`~arbi_serve.backends.tkv_backend.TkvBackend` (uint8 codec storage)
and :class:`~arbi_serve.backends.tkv_bypass_backend.TkvBypassBackend` (raw
bf16 bypass) are two **codec modes of one paged-KV backend family**: both
serve :attr:`StateKind.PAGED_KV` with a single fused KV plane, both thread
MTP intent the same way, both drive the turbo-attn kernel stack, and both
carry the same per-step CSR page metadata. The Turbo prefill / split-K decode
kernels they call into are policy-parametric on the KV format, so the two
backends differ only in the codec-specific geometry (uint8-packed vs
bf16-fused slot), the concrete attn-op / metadata-builder they construct,
and the custom-op dispatch name.

This module hoists what the two modes share so each subclass is thin:

  - :class:`PagedKvMetadataBuilder` owns the ``(CuSeqlens, BlockTable)``
    sub-builder pair and the shared :class:`AttnPagedKVMeta` construction
    (:meth:`_allocate_meta`, byte-identical between the modes). Each
    concrete builder adds its codec-specific ``_finalize`` wiring (TKV's
    ``TQRunState`` + Turbo prefill scratch; bypass's plain CSR triplet stash).
  - :class:`PagedKvBackend` owns the descriptor boilerplate: the
    ``PAGED_KV`` state kind, the single fused plane, MTP threading, and
    the ``make_attn_op`` head-slice.

Parallels :mod:`arbi_serve.backends._recurrent_base` (the recurrent
family's shared base) for the paged family.

Shared base for paged-KV per-step metadata builders.

Owns the parts identical across the codec modes: the
``(CuSeqlensSubBuilder, BlockTableSubBuilder)`` pair (shape +
block-table slices every paged backend needs), the
:class:`AttnPagedKVMeta` construction, and the persistent CSR
page-metadata refresh (:meth:`_fill_page_csr` + the
:meth:`cudagraph_capture_step` / :meth:`recompute_external_meta`
entry points that are pure CSR refreshes). Concrete subclasses:

  * expose their persistent CSR triplet via :meth:`_csr_buffers`
    (TKV's live in its ``TQBufferPool``; bypass owns plain tensors)
    and set ``self.block_size`` in ``__init__``;
  * add their codec-specific per-step wiring in :meth:`_finalize` —
    TKV's ``TQRunState`` / Turbo prefill scratch, or bypass's plain CSR
    triplet stash (``_bt_indptr`` / ``_bt_indices`` / ``_bt_lpl``).

Descriptor base for paged-KV attention backends (the codec modes).

Both codec modes serve ``StateKind.PAGED_KV`` with a single fused KV
plane and thread MTP intent identically. Subclasses set the
codec-specific geometry (:meth:`kv_cache_dtype` / :meth:`kv_cache_shape`
/ :meth:`bytes_per_token`), construct their concrete metadata builder +
attn op (``make_metadata_builder`` / ``make_attn_op``), and — via
:meth:`_paged_attn_op_heads` — share the TP head-slice their
``make_attn_op`` needs.

Return the persistent ``(indptr, indices, last_page_len)`` triplet.

The buffers must be pre-allocated at builder construction and
stable across steps (captured decode graphs bake their
``data_ptr``\ s), sized for the max capture batch.

Repopulate the persistent CSR triplet in place from live extents.

Runs :func:`tkv.runtime.page_metadata.compute_page_metadata_cuda`
into the pre-allocated :meth:`_csr_buffers` — no host sync, no
allocation, capture-safe. ``seq_lens`` / ``block_table`` come from
either the batch (build / capture path) or the meta (drafter
path). Returns the live batch size ``B``.

Re-populate the CSR page-metadata buffers from the (captured /
persistent) batch buffers.

Called once during build and again on every cudagraph replay (via
the capture path's ``getattr(builder, "cudagraph_capture_step",
None)`` hook). Reads ``batch.seq_lens`` + ``batch.block_table``
(persistent tensors the replay ``copy_()``\ s each step) and
writes the :meth:`_csr_buffers` triplet in place — no host sync.
Subclasses with additional per-replay scratch (TKV's Turbo prefill verify
extents) extend via ``super()``.

No-sync per-step CSR refresh for the CAPTURED drafter chain.

``decorate_external_meta`` runs once pre-capture to attach the
persistent buffer references to the meta. Inside the captured
region the per-step ``seq_lens`` / ``block_table`` change each
replay (``copy_()``'d by ``DrafterChainGraph.replay``), so this
re-derives ``indptr/indices/lpl`` in place into the SAME
persistent buffers (stable ``data_ptr`` the captured kernel
baked). Reads extents off the meta — no ``ScheduledBatch`` on the
drafter path. Capture-safe (no host sync). Identical for both
codec modes.

The MTP verify forward is cudagraph-capturable on both codec modes.

Both modes route the verify shape (B rows × (K+1) query tokens) to
a split-K paged MTP kernel — a plain CUDA launch with no host sync
— so the whole-forward verify graph captures cleanly. A codec mode
without a capturable verify kernel overrides this to False; the
engine's capture path reads the answer off the backend serving
``PAGED_KV``, so a backend that cannot answer refuses capture setup
rather than being captured on an assumed True.

Per-paged-layer slot byte sizes, or ``None`` when uniform.

``None`` means every paged layer resolves to the same slot size, so
the pool keeps its single stacked-slab layout. A codec mode that can
size layers differently (a per-layer bit table) overrides this and
returns one entry per ``paged_layer_indices`` entry, in that order.

Single fused plane per layer (K and V interleaved in one slot).

Both codec modes pack K and V into one fused slot (uint8-packed
for TKV, bf16 ``[K | V]`` for bypass); the pool treats the rank-3
``kv_cache_shape`` as a single plane and hands the raw slab to the
attn op as the ``state_view``.

Per-rank ``(num_heads, num_kv_heads)``, TP-sliced by the current
parallel config.

Shared prologue for both codec modes' ``make_attn_op``. Layers that
declare :attr:`LayerSpec.attention_sinks` carry the per-head sink
tensor separately: the engine binds it onto the constructed op via
:func:`arbi_serve.engine.attention_sinks.bind_attention_sinks`.

Shared base for block-owned recurrent-state backend descriptors.

GDN, Mamba and ShortConv are all "recurrent" backends: the layer's
forward runs inside the model block (``GDNBlock`` / ``Mamba2Block`` /
``ShortConvBlock``), the engine-side :class:`AttnOp` is a no-op
pass-through, and the backend keeps no paged KV slab — recurrent state
is per-request (owned by a state pool), not per-token.

This module hoists the boilerplate the three descriptors share. Each
backend module subclasses these and fills the arch-specific hooks: the
state kind, the metadata-builder and attn-op types, the default name,
and the builder's ``sub_builders`` / ``_allocate_meta`` / ``_finalize``.

Engine-side no-op AttnOp for block-owned recurrent layers.

The recurrent block runs the layer's forward directly; the engine
still constructs one ``attn_op`` per layer so the ``attn_ops`` list
is fully populated. Dispatching the layer through this op is a
programmer error, so :meth:`forward` always raises.

Subclasses set :attr:`arch` (display name, e.g. ``"GDN"``) and
:attr:`block` (owning block, e.g. ``"GDNBlock"``); the raised message
names the concrete subclass via ``type(self).__name__``.

Shared base for recurrent-backend metadata builders.

Concrete subclasses set the :attr:`sub_builders` tuple (a single
:class:`StateIndicesSubBuilder` for their state kind) and implement
:meth:`_allocate_meta` / :meth:`_finalize` for their per-step meta
object. The drain and the cudagraph capture hook are shared (see
:meth:`build`, :meth:`cudagraph_capture_step`).

Descriptor base for block-owned recurrent-state backends.

Recurrent state is per-request (owned by a state pool), so these
backends expose no paged KV slab: :meth:`kv_cache_shape` is empty
and :meth:`bytes_per_token` is zero. Subclasses set the class-level
hooks :attr:`default_name`, :attr:`state_kind_value`,
:attr:`metadata_builder_cls` and :attr:`attn_op_cls`.

Settle the recurrent slab's queued writes, then compose the meta.

A recurrent meta names the slab rows the forward will gather
``h0`` from, and the GDN prefill path gathers it UNMASKED
(``_gdn_fla._forward_prefill_fla``) — so an admission-queued
zero-clear that has not been applied by the time the meta is
built is a forward reading a freed request's state, silently.

The drain belongs HERE, on the base every recurrent builder
inherits, because the alternative is a precondition each entry
path has to remember: the runner's start-of-step, the capture
sweep, the boot readiness gate, the post-resize warmup and the
activation probe are all just callers of :meth:`build`, and an
entry path nobody has written yet will be one too. Draining
where the slab is read makes the invariant a property of the
read rather than of the caller.

Safe at this depth because :meth:`build` never runs inside a
cudagraph capture region — capture populates per-replay scratch
through :meth:`cudagraph_capture_step` instead — so the flush's
H2D copy and indexed zero-write cannot be baked into a graph.

No-op: recurrent backends need no per-replay scratch population.

The capture-path engine pre-binds ``batch.state_indices`` to a
persistent buffer before :meth:`build` runs, so the meta the
captured kernel sees already references the stable ``data_ptr``.
Replay-time updates flow through :meth:`CapturedGraph.replay`'s
``recurrent_state_indices`` parameter. The hook exists as a
no-op rather than being absent because the cudagraph capture
path's ``getattr(builder, "cudagraph_capture_step", None)``
lookup treats "absent" as "this builder's meta references batch
tensors only" — which is true here, so absence would also be
correct, but defining the hook explicitly documents the intent.

Attention-backend protocols.

A backend is a triplet:
  - :class:`AttentionBackend` — descriptor (name, slab shape / dtype /
    bytes-per-token, factories for the per-step metadata builder and
    the per-layer attn op).
  - :class:`arbi_serve.backends.metadata.MetadataBuilder` — composable
    per-step builder ABC; backends sub-class and declare which
    sub-builders (``cu_seqlens``, ``block_table``, ``lora_bucket``,
    ``state_indices``) they need. See that module's docstring for
    details.
  - :class:`AttnOp` — per-layer forward callable with the uniform
    signature ``(q, k, v, state_view, batch_meta, output,
    num_tokens)``. ``batch_meta`` is the kind-specific dataclass on
    :class:`ScheduledBatch`, never the whole batch.

Each backend declares ``state_kind() -> StateKind``; there is no
default. The engine builds one
``active_backends: dict[StateKind, AttentionBackend]`` slot per state
kind the model uses, and picks the per-layer attn op from that map
via :attr:`LayerSpec.state_kind`.

:class:`AttnOp` and :class:`AttentionBackend` are :class:`typing.Protocol`s
— duck typing keeps the surface light. :class:`MetadataBuilder` is
an ABC (in :mod:`arbi_serve.backends.metadata`) because the
composable shape is best expressed with a concrete ``build()`` glue
on the base class plus per-backend overrides of ``_allocate_meta`` /
``_finalize``.

Per-layer attention forward.

One instance is created per (layer × backend) and held by the
engine. On every step the engine calls :meth:`forward` once per
layer with the layer's per-state-kind slab view and the
kind-specific metadata view produced by the active
:class:`MetadataBuilder`.

Implementations must:
  - write per-step K / V into the layer KV cache view
    (compress-then-store for tkv; bf16 scatter for flash-attn;
    recurrent state update for Mamba / GDN);
  - run attention over the cumulative cache;
  - write per-token output into ``output[:num_tokens]``
    with shape ``(N, num_heads, head_dim)``.

Backend descriptor.

One instance per name in the engine registry. Stateless w.r.t. the
model — the engine asks the backend for shape / metadata-builder /
per-layer attn-op factories, then drives them.

Naming convention is load-bearing: the CLI parses ``--backend
kind:name`` (e.g. ``paged_kv:tkv-k4v4``) and the HTTP admin
endpoint swaps by ``(kind, name)``. A backend's ``name`` is its
registry key; ``state_kind()`` declares which slot it fills.

DeepSeek-V4 latent decode on turbo-attn's split-K decode kernel.

DSv4 scores every query head against ONE 512-wide latent per token and
that same latent is the value, which is the ``(v_aliases_k,
TQ_VEC_SIZE_R == 0)`` corner of turbo-attn's unified split-K decode
kernel — reached from Python through
``tkv.kernels.turbo_attn_simt_latent_mqa.tq_latent_mqa_batch_decode``.

The adaptation is one mapping. A DSv4 query reads a per-token list of
entry ids into a grouped slab; the kernel reads a paged KV cache through
a ``block_table`` at ``page_size = 1``. So a row's list becomes that
row's block table once three things are fixed:

  * **flat block index.** The slab's id space is per-slot, the kernel's
    is global, so id ``e`` of slot ``s`` is the block
    :meth:`~arbi_serve.models._deepseek_v4_kv.DSv4SlabAddress.flat` puts
    it at.
  * **backed rows.** An id is clamped to the rows its slot's group has
    physically mapped, so a block table never names reserved-but-unbacked
    VA. A clamped id is one the caller's own mask already discards.
  * **compaction.** ``-1`` is DSv4's pad and can sit anywhere in a list
    (before the start of a short sequence, after a short top-k); the
    kernel walks the first ``seq_lens[n]`` entries of a row. The valid
    ids move to the front IN ORDER and ``seq_lens`` counts them, so the
    walked prefix names exactly the ids the reference gathers.

``max_seq_len`` is the id list's STATIC width — a host-known upper bound
on ``seq_lens``. Reading the true max off ``seq_lens`` is a device sync
and is illegal under graph capture.

The latent's own shape has to suit the kernel's lane split — see
:func:`check_dsv4_kernel_geometry`, which refuses a geometry here rather
than letting a device-side check fire mid-step.

``sliding_window`` is 0: this arch's window is already resolved into the
id list, so the kernel must range over every entry the list names. The
per-head sink rides through to the phase-2 merge, which is where it
belongs — it enters the softmax denominator once per row, never per
split.

The torch path (:func:`~arbi_serve.models._deepseek_v4_ops.sparse_latent_attend`)
is the executable reference this mapping is diffed against, and the only
path off CUDA.

Refuse a latent shape the split-K kernel cannot address.

The kernel gives each lane of a warp ``head_dim // 32`` consecutive
channels and has it read a single E8M0 exponent for them, so a group
that is not a whole number of lane slices would scale part of a slice
with its neighbour's exponent. It also indexes the exponent by shift,
which needs a power-of-two group. Both are device-side ``TORCH_CHECK``
otherwise — a fault deep in a captured step rather than here.

The call ``tq_latent_mqa_batch_decode`` is made with, built but not made.

Split out from the call so the mapping is checkable where the kernel
cannot run: every field is a plain tensor or int, and together they
describe one attention problem that
:func:`sparse_latent_attend` describes too.

Map one decode step's id lists onto the kernel's paged arguments.

``ids`` must name CACHED rows only — an id at or past ``addr.rows``
addresses this step's own latents, which are not in the slab and which
no block table can reach.

Every step here is a device-side tensor op on shapes fixed by the
layer's capacities, so this runs unchanged inside a captured graph.

Attend a decode step whose id lists name cached rows only.

On CUDA this is the split-K kernel, always: the step's shape already
established that every id addresses the slab, so there is no case
left in which the kernel cannot serve it. Off CUDA there is no kernel
to select and :func:`dsv4_reference_attend` runs.

A ZERO-ROW step returns early. Attention-DP makes this reachable: an
attention-DP set with no requests must still enter the forward, because
the FFN exit reduces over the whole TP group and a rank that skipped it
would strand every peer in that collective. The caller's "every run
holds one token" predicate is vacuously true over an empty run list, so
a 0-row step arrives here rather than being filtered upstream, and a
launch with ``num_seqs == 0`` is a grid of zero blocks — not obviously
fatal, but not something the kernel is specified for either.

DeepSeek-V4 sparse-attention backend descriptor.

State kind: :class:`StateKind.DSV4_SPARSE`. Like the recurrent backends,
the attention itself lives inside the model's own block — it reads its
three per-request stores directly rather than through a per-layer op — so
the engine-side :class:`AttnOp` is a no-op and the whole job here is
metadata.

That metadata is deliberately small and deliberately HOST-derived:
per-request slab rows and per-row token counts. Both are facts the
scheduler already holds, and both would otherwise have to be read back off
``cu_seqlens_q`` inside the forward — a device sync that serialises the
step and is illegal mid-capture. The model refuses a step that arrives
without them rather than reading them back.

The decode attend the block calls is this backend's surface too:
:func:`~arbi_serve.backends.dsv4_attend.dsv4_decode_attend`, re-exported
here, which routes a cached-only decode step onto turbo-attn's split-K
DSv4 kernel.

Engine-side no-op AttnOp for DeepSeek-V4 layers.

The attention runs inside :class:`DSv4LatentAttention`, which owns the
window ring, both pooled streams and the id lists that address them.
Dispatching through this op is a bug.

Per-row token counts for this step, from the HOST.

``query_lens`` when the batch carries it, else the host mirror's
``cu_seqlens_q`` (already a CPU tensor), else the uniform-decode shape
— one token per sequence — which ``num_seqs`` and the token count
establish between them without reading any tensor's values.

Returns ``()`` when none of those describe the step; the model turns
that into a loud refusal, because the only remaining source is a device
read.

The ``DSV4_SPARSE`` pool view reachable from ``slab``, or ``None``.

Call sites hand the metadata builder either the
:class:`~arbi_serve.cache.multi_state_pool.MultiStatePool` or a bare
pool; ``None`` covers a synthetic capture / test batch built without
one, where there is no mapping to check against.

Per-row ``(slab row, full sequence length)`` from the HOST, or ``None``.

The slab row comes from the pool's own allocator, keyed by the batch's
request ids; the length from the pinned host twin of ``seq_lens``, or
from ``seq_lens`` itself when it already lives on the host. ``None``
means one of the two is unavailable — a synthetic capture batch, a
request the pool never allocated for — and the caller falls back to the
batch-level check rather than reading either off the device.

Build :class:`DSv4Meta` from a :class:`ScheduledBatch`.

Slab-row resolution is the shared
:class:`StateIndicesSubBuilder` (kind ``DSV4_SPARSE``) — the same
per-request row allocator the recurrent kinds use, which is what makes
a request's window, streams and accumulators share one row across the
whole stack.

Attach the host-side token counts and the verify flag, and check
the batch fits the pool.

The pool backs each slab row to what that row's request declared, so
a sequence longer than its own row's backing would address unbacked
memory and fault the device inside the forward. Every number the
check reads is a host int the scheduler already computed, so it
costs comparisons and no device read — this build runs downstream of
the start-of-step drain that applies the growth (the same
drain-before-build contract the GDN builder holds), so the mapping
it reads is this step's.

DeepSeek-V4's two per-step fused kernels, adapted to this repo's tensors.

Both live in ``tkv.kernels.cuda_trained_quant_fused`` and each replaces a chain of
torch ops the forward would otherwise run per layer per step:

  * :func:`dsv4_kernel_compress_store` — the compressed stream's write.
    Gate softmax pooling over the window, RMSNorm, interleaved RoPE on the
    tail, the UE8M0-scaled FP8 E4M3 round trip, the packed slot write and
    the scatter to one flat slot per pooled entry, in one launch over one
    read of the window. It reproduces the round trip TWICE, because the served path
    applies it twice — once in the forward and once inside
    :meth:`~arbi_serve.models._deepseek_v4_kv.DSv4KVLayout.encode`, which
    re-derives the group scale from already-quantized values. That is
    value-idempotent and NOT byte-idempotent, so a single-stage kernel
    would write different bytes for the same latent on the groups whose
    re-derived amax falls back a power of two.
  * :func:`dsv4_kernel_indexer_q` — the indexer query's interleaved RoPE,
    Sylvester Walsh-Hadamard rotation and FP4 E2M1 round trip.

The torch ops in :mod:`arbi_serve.models._deepseek_v4_ops` and the
compressor's own ``finalize`` remain the executable spec: they are the only
path off CUDA and what the adaptation here is diffed against.

Two things the geometry checks exist for. The compress kernel gives each
lane four channels and reduces the RMSNorm sum of squares as a per-warp XOR
butterfly, so the latent has to be a whole number of warps of lane slices
and a scale group a whole number of lanes inside one warp. The indexer
kernel puts ONE row on ONE warp, which fixes its head width at exactly
``32 * 4``. Both are ``static_assert`` on the turbo-attn side — a build
failure deep inside a step's first launch rather than a refusal here.

Both entries take their rope angles as a table plus an index. The angles
this repo holds are the model-dtype ``RoPECache`` tables; the torch ops
widen a gathered row to fp32 before rotating, so what the kernel must see
is that same widened row. The gathered rows ARE the table handed over, and
the index is then the identity — which keeps the widening to one row per
entry instead of a per-step fp32 copy of the whole table.

The rotated (Hadamard + FP4) compressed stream the indexer builds is a
different pipeline: no fused kernel covers it, so
:func:`dsv4_compress_store_args` refuses to describe one.

Whether a step's tensors sit on the device the fused entries run on.

Both entries are CUDA kernels. Off CUDA there is nothing to select and
the torch ops they were written from run instead — the same split the
DSv4 decode attention makes.

The call ``fused_pooled_compress_store`` is made with, built but not made.

Split out from the call so the mapping is checkable where the kernel
cannot run: every field is a plain tensor or a number, and together
they describe the same write the compressor's ``pool_entries`` /
``finalize`` / ``DSv4KVLayout.encode`` / masked scatter chain performs.

Map one stream advance's pooling window onto the kernel's arguments.

``window_kv`` / ``window_score`` are the pooling window flattened over
``(sequence, group)``; ``dst`` / ``valid`` are the scatter the torch
path performs with ``store[dst] = where(valid, ...)``, handed over
element-wise so the kernel leaves an invalid destination untouched
instead of reading and rewriting it. ``store`` is the flat slab as
``(flat_slots, 1, slot_bytes)``, so the kernel's ``row_stride`` is one
slot and its slot index is always zero — the group-major slab has no
``(row, slot)`` strided view for it to walk.

Every step is a device-side tensor op on shapes fixed by the layer's
capacities, so this runs unchanged inside a captured graph.

Gated Delta Net backend descriptor — Qwen3-Next.

State kind: :class:`StateKind.GDN`. Mirrors the Mamba backend's
shape — the heavy lifting lives inside :class:`GDNBlock`; the
engine-side :class:`AttnOp` is a no-op pass-through and the metadata
builder produces a degenerate :class:`GDNMeta` (per-row slab index +
seq lens, no slot mapping).

Resolve this step's GDN chunk-index table from the host twin.

The GDN layers reach ``prepare_chunk_indices`` from inside the opaque
``arbi_serve::gdn_attention_v2`` custom op, which carries no host twin.
Resolved there, FLA takes its device route, whose ``_segmented_arange``
sizes a ``repeat_interleave`` by ``counts.sum()`` and blocks the host on
a device read. Resolved HERE, on the host, before the forward, the
per-step memo hands every layer the same table and no layer reads the
device for it.

Prefill only: the chunk kernel is the only consumer, and the decode /
verify route is fused-recurrent, which resolves no chunk indices. On a
row-class mix the chunk kernel consumes the PREFILL class's boundaries
(:attr:`GDNMeta.cu_seqlens_q_split`), so that is the tensor primed: the
memo is keyed on the tensor's identity, and priming the fused
boundaries would leave every layer to resolve the table it actually
uses on the device.

Record whether this step's GDN prefill can take the host-twin route.

Mirrors :func:`~arbi_serve.models._gdn_fla_forward._resolve_cu_seqlens`'s
branch order against the SAME batch fields it reads, one step earlier and
on the host. The consumer compares ``cu.device`` against the hidden
states' device; the pool and the batch are built on one device, so the
batch's own ``seq_lens`` device stands in for it here.

Host-side attribute reads and one int increment — no device read, no sync.

Publish this step's fold-split plan for the GDN ops, or refuse by name.

The plan is the scheduler's (``Scheduler._arm_fold_split``); nothing is
re-decided here, and the derivation itself is
:func:`~arbi_serve.cache._fold_emit_staging.publish_step_plan` — shared
with the captured-prefill replay, which builds no metadata — so a step
is armed for the ops in exactly one way: the request id resolved to the
LAUNCH ROW the chunk kernel sees, the kernel's tables written host-side,
and the identity keys the ops check before they emit. Nothing about the
plan is a Python value inside a compiled forward
(:mod:`arbi_serve.cache._fold_emit_staging` says why that matters).

On a row-class-mixed step the chunk kernel consumes the prefill class
alone, rebased to start at 0 (:attr:`GDNMeta.cu_seqlens_q_split`), so
the plan is keyed on THAT tensor and its row is relative to it.

Build :class:`GDNMeta` from a :class:`ScheduledBatch`.

Per-row recurrent slab-row resolution is owned by
:class:`StateIndicesSubBuilder` (kind=GDN) — see its docstring for
the resolution order. :meth:`_finalize` derives
``has_initial_state`` from ``seq_lens`` and decides ``verify_pass``
from the batch's MTP context.

Verify pass: multi-token-per-row batch with ``mtp_meta`` set and
``is_prefill=False``. The block routes this through
``fused_recurrent_gated_delta_rule`` (T=K+1 input, varlen via
``cu_seqlens``) so bonus-slot numerics match K=1 decode — the
bundled MTP head was trained against fused-recurrent output.

Copy the batch's row-class split onto the metadata, with its own ``cu``.

The split is data the batch builder decided
(:func:`~arbi_serve.runtime.fused_mixed_step.fused_mixed_batch`); this
only validates it against the rows and tokens the metadata describes and
resolves the prefill class's boundaries ONCE, so the block's layers all
hand the chunk kernel one tensor. A split that does not leave both
classes non-empty is a builder defect and is refused loudly rather than
routed to a class-pure forward that would commit the verify rows' state.

Mamba backend descriptor — Mamba-2 hybrids.

State kind: :class:`StateKind.MAMBA`. The backend descriptor is
intentionally tiny — Mamba's "attention" lives inside the block
(the block owns the projections and the SSM forward). The engine-facing
:class:`AttnOp` is a no-op pass-through; a Mamba decoder layer calls
``mixer.forward(...)`` directly rather than dispatching through
``attn_ops[i]``.

The :class:`MambaMetadataBuilder` builds the per-step
:class:`MambaMeta` from a :class:`ScheduledBatch` — a per-row slab-row
index tensor + an ``has_initial_state`` mask + the seq-len vector.
There is no slot mapping (recurrent state is per-request, not per-token).

Publish this step's fold-split plan for the Mamba layers, or refuse by name.

The plan is the scheduler's (``Scheduler._arm_fold_split``); nothing is
re-decided here, and the derivation is
:func:`~arbi_serve.cache._fold_emit_staging.publish_step_plan` — the same
one the GDN builder and the captured-prefill replay take — so a step is
armed for the ops in exactly one way. It counts on the GDN builder's
``gdn_fold_split_delivered`` rather than a second counter: the fact is
"the decision reached the layers", one per step, and two counters for it
would be two answers to one question.

A pool that also holds a GDN view is left to the GDN builder, which keys
the plan on the chunk kernel's own boundary tensor (rebased on a
row-class-mixed step); publishing again here would key it on another
tensor and refuse GDN's launches by identity.

Engine-side no-op AttnOp for Mamba layers.

Mamba's forward lives inside the mixer block
(:class:`Mamba2Block`); the model calls the
block directly. The engine still constructs one ``attn_op`` per
layer (so the ``attn_ops`` list is fully populated) — Mamba's entry
is this no-op. Calling ``forward`` is a programmer error and we
raise to surface it.

Build :class:`MambaMeta` from a :class:`ScheduledBatch`.

Per-row recurrent slab-row resolution is owned by
:class:`StateIndicesSubBuilder` (kind=MAMBA) — see its docstring
for the resolution order (``batch.state_indices`` →
``RecurrentStatePool.row_for(req_id)`` → ``arange(B)`` default).
:meth:`_finalize` derives ``has_initial_state`` from ``seq_lens``
and decides ``verify_pass`` from the batch's MTP context.

Verify pass: multi-token-per-row batch with ``mtp_meta`` set and
``is_prefill=False``. The block routes this through the per-token
decode kernels so each of the K+1 positions advances the conv/ssm
slab one step at a time — the chunk scan commits only the final
post-K+1 state, which partial-accept rollback cannot undo.

Mamba SSM backend descriptor.

Stateless w.r.t. the model — every per-layer fact (intermediate
size, state dim, conv kernel, dt rank) lives on
:class:`LayerSpec.mamba`; the engine reads it directly when the
block is constructed. This descriptor only exists to fill the
``active_backends: dict[StateKind, AttentionBackend]`` slot.

Composable :class:`MetadataBuilder` ABC + sub-builders.

One shared ABC plus a small set of composable sub-builders backs the
per-backend metadata builders (tkv, tkv-bypass, mla, mamba,
gdn, short_conv) instead of each reinventing the ``(batch, slab) ->
Meta`` pattern. Each backend declares which sub-builders it needs;
:meth:`build` is glue, not logic.

  - ``_allocate_meta`` — backend constructs its kind-specific empty
    meta dataclass (``AttnPagedKVMeta`` / ``MLAMeta`` / ``MambaMeta``
    / ``GDNMeta``).
  - sub-builders run in declared order — each ``populate(meta, batch,
    slab)`` mutates ``meta`` in place with its slice (cu_seqlens,
    block_table, recurrent state_indices, etc.). Stateless; the same
    instance is shared across backends.
  - ``_finalize`` — per-backend wiring that doesn't fit the
    sub-builder pattern (TKV's ``tq_run_state`` + bypass_safe + Turbo
    prefill scratch; MLA's geometry-validated bf16 decompress scratch
    if any; recurrent backends' ``has_initial_state`` derivation).

Per-backend code shrinks to ~30 LOC: an ``_allocate_meta`` + a
``_finalize`` + a ``sub_builders`` tuple.

One reusable building block in a :class:`MetadataBuilder`.

Each sub-builder owns a narrow piece of per-step metadata
(cu_seqlens, block table, RoPE offsets, LoRA bucket, recurrent
state-row indices) and *populates that slice in-place* on the
backend-specific meta dataclass the parent :class:`MetadataBuilder`
constructed in :meth:`_allocate_meta`.

Sub-builders are stateless; the same instance can (and should) be
shared across backends. ``CuSeqlensSubBuilder`` is shape-only and
applies to every paged backend; ``BlockTableSubBuilder`` only to
paged backends; ``StateIndicesSubBuilder`` only to recurrent
(Mamba/GDN) backends.

The contract is ``populate(meta, batch, slab) -> None`` — implementations
must not return a new meta; mutate ``meta`` in place. The parent
:class:`MetadataBuilder._finalize` may then attach backend-specific
derived fields (e.g. TKV's ``tq_run_state``).

Composable per-step metadata builder.

Each backend subclasses this ABC, declares its
:attr:`sub_builders` tuple in execution order, implements
:meth:`_allocate_meta` (construct the empty kind-specific meta),
and :meth:`_finalize` (post-sub-builder wiring — TKV's
``tq_run_state``, recurrent backends' ``has_initial_state``).

The :meth:`build` glue is fixed: allocate, run sub-builders,
finalize, return.

Call sites pass ``(ScheduledBatch, slab/pool)``.

Re-derive this builder's persistent scratch inside the capture.

A captured graph bakes the addresses of the buffers it recorded.
Replay ``copy_()``s fresh contents into the persistent ``batch``
tensors, so any buffer this builder OWNS that is a function of those
tensors — the paged CSR triplet, for one — must be recomputed by a
call recorded INSIDE the graph, or every replay gathers against the
capture-time synthetic batch.

Every builder answers. A builder whose meta pointer-aliases the
persistent ``batch`` tensors owns no such scratch and says so with an
explicit no-op body: absence would otherwise mean both "nothing to
re-derive" and "scratch nobody wired up", and only one of those is
safe to replay.

MLA backend — DeepSeek V2 / V3, Kimi K2.

State kind: ``StateKind.MLA_SHARED``. Per-layer slab layout is a
``(num_pages, page_size, slot_bytes)`` uint8 paged cache;
``slot_bytes`` is computed via
:func:`tkv.runtime.mla.spec.mla_slot_bytes` from the MLA geometry on
``LayerSpec.mla`` and the per-backend ``k_bits``.

Three tkv entry points carry the layer:

  * **compress on cache write** —
    :func:`tkv.runtime.mla.compress_into_slot` packs
    ``(kv_c_normed, k_pe)`` into the uint8 slot at the indices given
    by ``slot_mapping``.
  * **decode** —
    :func:`tkv.kernels.turbo_attn_simt_mla.tq_mla_batch_decode` reads the
    packed uint8 slots directly and dequantises in-kernel. Runs on a
    query-length-1 step against the weight-absorbed query, launched once
    per head chunk so any head count is served at any tensor-parallel
    degree. The kernel scores in the codec's rotated space; the un-rotate
    applies ``k_R_inv`` *and* ``k_per_channel_scale``, so a calibrated
    codec is served exactly.
  * **prefill / mixed** —
    :func:`tkv.kernels.cuda_mla_dequant_gather.tq_gather_and_dequant_mla`
    gathers and dequantises only the batch's cached tokens into a
    bounded ``(total_kv_tokens, kv_lora_rank + qk_rope_head_dim)``
    bf16 workspace; K and V are unfolded from it through ``kv_b_proj``
    and attended by tkv's Turbo prefill varlen kernel
    (``turbo_prefill`` + ``BypassLoader``).

Neither serving path materialises a bf16 mirror of the cache, and
nothing here can reach tkv's whole-slab decompress — it is a test-only
parity oracle. The remaining refusals in
:meth:`MLAAttnOp._unsupported_geometry` are hardware / codec bounds, not
fallback triggers.

A single shared bit width ``k_bits`` (default 4) is used; MLA
shared-KV folds K and V into one quantiser so there is no separate
``v_bits``. Override via the backend name (``mla-k4``, ``mla-k8``, …)
or the ``TKV_K_BITS`` env var.

Per-step metadata builder for the MLA backend.

Same paged-KV shape as the bf16 paged-KV builders — the MLA
op consumes ``cu_seqlens_q``, ``seq_lens``, ``block_table``,
``slot_mapping``, ``max_seq_len``, ``max_query_len`` — but
interpreted against MLA slot bytes (see :class:`MlaBackend`'s
:meth:`kv_cache_shape`).

:class:`MLAMeta` doesn't carry ``cu_seqlens_k`` (MLA is full-KV
over the cumulative cache, not varlen-K), so the
:class:`CuSeqlensSubBuilder`'s ``cu_seqlens_k`` branch is a no-op
here — the meta dataclass simply lacks the field.

``MLAMeta.total_kv`` is taken from the batch's host-side
``total_kv_tokens``; a batch built without it is refused here rather
than served by a device→host sync in the attend op.

Per-layer MLA attention forward over the packed uint8 slot cache.

Every step compresses the new tokens' ``(kv_c_normed, k_pe)`` into
the slot at ``slot_mapping``
(:func:`tkv.runtime.mla.compress_into_slot`), then routes by query
length:

  * ``max_query_len == 1`` — absorb ``W_UK`` into the query, rotate
    it into the codec's basis, and call ``tq_mla_batch_decode`` over
    the packed slots; un-rotate the latent result through ``k_R_inv``
    and ``k_per_channel_scale``, then unfold it through ``W_UV``.
  * ``max_query_len == mtp_block_m > 1`` and ``N == B * block_m`` —
    the same absorption through ``tq_mla_batch_mtp_decode``, which
    gives every draft row its own causal bound.
  * otherwise — gather-and-dequant the batch's cached tokens into a
    bounded bf16 latent workspace
    (``tq_gather_and_dequant_mla``), unfold K and V through
    ``kv_b_proj``, and run tkv's Turbo prefill varlen kernel.

The :class:`MLAAttentionBlock` calls :meth:`forward` with the *full*
pre-projection tensors:
  ``q``: ``[N, num_heads, qk_head_dim]`` (NoPE+RoPE concatenated)
  ``k_compressed``: ``[N, kv_lora_rank]`` — the post-norm latent
  ``k_rope``: ``[N, qk_rope_head_dim]`` — the per-token RoPE'd K
The op handles cache write + read + attend + W_UV unfold internally.

No-op: this builder owns no buffer that replay could leave stale.

Every device field of :class:`MLAMeta` is a pointer ALIAS of the
matching ``ScheduledBatch`` tensor, which under capture is the
persistent buffer replay copies into — so the replay's writes are
already what the recorded kernels read. There is no derived
quantity here (no paged CSR analogue): the MLA kernels consume
``seq_lens`` / ``block_table`` raw and bound K by the capture-shape
constant, not by batch data.

An :class:`MLAMeta` field that ever becomes a copy rather than an
alias, or a derived buffer this builder allocates, invalidates this
and must be re-derived here instead.

Bind the per-layer ``kv_b_proj`` weight.

``kv_b_proj`` projects ``(kv_lora_rank,)`` to
``(num_heads * (qk_nope_head_dim + v_head_dim))``. Split into the
per-head ``W_UK`` / ``W_UV`` halves used by the decode absorption
and the prefill unfold.

Return the first tkv MLA kernel precondition this layer fails.

``None`` means both native kernels can serve this geometry on
this device with this codec. Every branch here is a device or
codec bound: there is no launch shape, head chunking, or codec
setting that satisfies it.

Build this layer's private codec on ``device`` if absent.

The codec is private per layer so calibration writes land on one
layer only. Idempotent; a codec already installed on a different
device is an error rather than a silent rebuild, because the
installed one may already carry calibration.

Run MLA attention for this layer into ``output``.

Splits the fused query into nope/pe parts and dispatches the
``arbi_serve.mla_attention`` op (always causal).

Every per-step value the op needs travels through the op's own
schema — the cache slab, the four scheduling tensors, and the
host ints (``max_query_len`` / ``max_seq_len`` / ``total_kv`` /
``mtp_block_m``). Nothing is read back off a Python attribute
inside the op, so Dynamo can trace this call and a captured
graph replays the values its buffers hold rather than whatever
the last eager forward happened to stash.

Real-impl side of ``arbi_serve::mla_attention``.

Writes the new tokens into the packed cache and routes the step to
:meth:`_native_decode` (query length 1), :meth:`_native_mtp_verify`
(uniform ``block_m`` rows per sequence) or :meth:`_gathered_attend`.

The step's :class:`MLAMeta` is rebuilt from the op's own
arguments; the op reads no per-step Python state.

Write these tokens' slots without attending.

The cache-write half of :meth:`forward`, reachable on its own:
the codec and the slot geometry resolve exactly as they do
there, and the bytes landed are the bytes that call would land
for the same ``(kv_c_normed, k_pe, slot_mapping)``. A second
write at the same ``slot_mapping`` supersedes the first.

Score the weight-absorbed query against the packed slots.

The kernel consumes a query pre-rotated into the codec's basis and
returns the latent output in that same basis; the un-rotate here is
``k_R_inv`` followed by ``k_per_channel_scale``, the K-side
counterpart of :meth:`TkvCodec.unrotate_output` (MLA shares one
quantiser between K and V, so the K scale is the V scale).

Returns ``(N, num_heads, kv_lora_rank)`` fp32.

Per-layer slab shape ``(num_pages, page_size, slot_bytes)``.

For MLA, ``head_dim`` is ``kv_lora_rank + qk_rope_head_dim``
(576 for V2-Lite, 576 for V3); the pool passes that through. We
derive ``slot_bytes`` from the MLA layout helper.

Build a codec private to one MLA layer, on ``device``.

Constructed directly on the target device so the rotation matrices
resolve to that device's cached Hadamard instead of a per-layer copy
of a host-built one.

Short-convolution backend descriptor — LFM2 / LFM2-MoE (Liquid AI).

State kind: :class:`StateKind.SHORT_CONV`. Like the Mamba / GDN
backends, the heavy lifting lives inside :class:`ShortConvBlock`; the
engine-side :class:`AttnOp` is a no-op pass-through and the metadata
builder produces a :class:`ShortConvMeta` (per-row conv-state slab
index + ``has_initial_state`` flag + seq lens, no paged slot mapping).

Build :class:`ShortConvMeta` from a :class:`ScheduledBatch`.

Per-row conv-state slab-row resolution is owned by
:class:`StateIndicesSubBuilder` (kind=SHORT_CONV). :meth:`_finalize`
derives ``has_initial_state`` from ``seq_lens`` (a row seeds from
the saved conv-state buffer iff it continues a prior step) and
carries ``is_prefill`` so :class:`ShortConvBlock` dispatches
``causal_conv1d_fn`` (prefill) vs ``causal_conv1d_update`` (decode).

Composable sub-builders for the :class:`MetadataBuilder` ABC.

Each sub-builder owns one slice of per-step backend metadata and
mutates the meta dataclass in place. Stateless; shared across backends.
See :mod:`arbi_serve.backends.metadata` for the parent ABC.

Block-table + slot-mapping sub-builder.

Populates ``meta.block_table`` and ``meta.slot_mapping`` from the
:class:`ScheduledBatch` scheduling tensors. Used by every paged-KV
backend (TKV, tkv-bypass, MLA).

Stateless — share one instance across every backend that needs it.

Populate block_table + slot_mapping + seq_lens + max_seq_len + max_query_len.

Sets the per-step paged-cache addressing fields:

  * ``meta.block_table`` — ``(B, max_pages)`` int32, per-request
    page IDs.
  * ``meta.slot_mapping`` — ``(N_tokens,)`` int64, flat slot index
    per token. Slot = ``page_id * page_size + offset``.
  * ``meta.seq_lens`` — ``(B,)`` int32 per-request total token count.
  * ``meta.max_seq_len`` — Python int, longest full seq.
  * ``meta.max_query_len`` — Python int, longest per-step query.

Pads to ``max_pages`` so the kernel grid is shape-stable for
cudagraph capture (the engine sizes ``block_table.shape[1]`` to
``ceil(max_context / page_size)`` at capture warm-up).

Like :class:`CuSeqlensSubBuilder`, fields are populated by direct
pointer-aliasing assignment from ``batch.*``, so the cudagraph-replay
``copy_()`` of ``batch.*`` is the only mutation between steps.

Cumulative-seq-length sub-builder.

Populates ``meta.cu_seqlens_q`` (and on paged-KV metas the matching
``cu_seqlens_k`` + ``query_start_loc`` alias) from the
:class:`ScheduledBatch` scheduling tensors.

Every backend that runs a varlen attention kernel consumes
``cu_seqlens_q`` (paged-KV via FA varlen / Turbo prefill / TKV;
recurrent-state via FLA's ``cu_seqlens`` prefill path); this is the
single most-shared metadata slice.

Stateless — share one instance across every backend that needs it.

Populate cu_seqlens_q (+ k for paged metas, + query_start_loc alias).

Sets:

  * ``meta.cu_seqlens_q`` — copy of ``batch.cu_seqlens_q``.
  * ``meta.cu_seqlens_k`` — copy of ``batch.cu_seqlens_k`` IF the
    meta dataclass declares the field (paged-KV metas do; MLA /
    recurrent metas don't).
  * ``meta.query_start_loc`` — alias of ``cu_seqlens_q`` for the
    AttendKernel wrappers' getattr-style introspection. Set IF the
    meta declares it (paged-KV via ``__post_init__``).

``query_start_loc`` is re-bound to ``batch.cu_seqlens_q`` here even
though :class:`AttnPagedKVMeta.__post_init__` already initializes
it from the constructor arg, because the constructor arg goes
through a ``field(init=False)`` and isn't a
:class:`ScheduledBatch`-side pointer alias. Re-binding keeps the
pointer stable across captured/replayed batches.

LoRA-bucket sub-builder.

LoRA per-step state is built by
:func:`arbi_serve.adapters.lora.build_lora_batch_state` and attached to
:attr:`ScheduledBatch.lora_state` (see ``arbi_serve/engine/batch.py``)
— it is not part of the per-backend :class:`MetadataBuilder` output.
Per-step LoRA bucket selection lives under
:mod:`arbi_serve.adapters.lora`.

This sub-builder is a no-op stub, reserved by name as one of the four
canonical sub-builders.

Stateless.

Recurrent-state-row index sub-builder.

For the recurrent backends (Mamba, GDN), each batch row maps to one
row in the per-layer recurrent slab. This sub-builder resolves that
mapping and populates ``meta.state_indices``.

Resolution order:

  1. If :attr:`ScheduledBatch.state_indices` is non-None, use it
     (engine-pinned mapping for MTP partial-accept replay forward
     and the captured-decode persistent buffer).
  2. Else if :attr:`ScheduledBatch.req_ids` is non-None AND a
     :class:`PerRequestRowPool` is reachable through ``slab`` for the
     sub-builder's :class:`StateKind`, look up each req_id's slab row
     via :meth:`PerRequestRowPool.row_for`. A missing
     ``alloc_for_request`` makes every ``row_for`` raise
     :class:`KeyError`; rather than silently falling through to (3)
     — where every request would decode against slab row 0, leaking
     state across sequential same-prompt requests — this path raises.
     The production path always pre-allocs, so a missing alloc gets a
     loud, actionable error.
  3. Sentinel-route default: every batch row maps to slab row 0
     (the permanent zero-sentinel; see
     :class:`RecurrentStatePool` module docstring) when ``req_ids``
     is None or no recurrent pool is reachable. Logs a one-time
     warning at the call site so a test fixture / capture-warmup
     synthetic batch hits a stable path without silently routing
     through a live request's slab row.

Each :class:`StateIndicesSubBuilder` is parameterized on
:class:`StateKind` (one for ``MAMBA``, one for ``GDN``) so the lookup
finds the correct recurrent pool view from the
:class:`MultiStatePool`. Stateless apart from the kind binding.

Populate ``meta.state_indices`` for recurrent (Mamba/GDN) backends.

Constructor arg :attr:`kind` selects which :class:`PerRequestRowPool`
view to look up by ``slab._views[kind]`` (the
:class:`MultiStatePool` keying). One instance per recurrent kind:

.. code-block:: python

    MAMBA_STATE_INDICES = StateIndicesSubBuilder(StateKind.MAMBA)
    GDN_STATE_INDICES = StateIndicesSubBuilder(StateKind.GDN)

Locate the per-:class:`StateKind` slab-row pool inside the
:class:`MultiStatePool`.

The contract this needs is :meth:`PerRequestRowPool.row_for`, which
is where the row allocator lives — NOT
:class:`RecurrentStatePool` specifically. Narrowing to the latter
would exclude :class:`ShortConvStatePool` and
:class:`DSv4StatePool`, which are siblings under the same base, and
silently route them to step (3) instead of the loud raise step (2)
exists to produce. ``resolve_recurrent_rows`` — the model-runner
mirror of this policy — already keys on
:data:`SLOT_STATE_KINDS` and duck-types ``row_for``, so a narrower
rule here makes the two sites disagree about which kinds are
covered.

Returns ``None`` when ``pool`` is ``None`` (test path), when
``pool._views`` is missing (legacy / mocked pool), or when the
kind isn't registered.

Return the smallest slab dim 0 across every recurrent pool view
registered on ``pool``, or ``None`` when no recurrent view is
registered.

Used by the profile-time forward to clamp synthetic
:attr:`ScheduledBatch.state_indices` into the slab's row range.
The recurrent backends (Mamba / GDN / ShortConv) read
``slab[state_indices[i]]`` per row; profiling against a scratch
pool with ``max_num_seqs=1`` AND a default ``arange(B)`` mapping
for ``B > 1`` raises :class:`IndexError` deep inside
:class:`GDNBlock` / :class:`Mamba2Block`. Clamping happens at the
profile call site, before building the batch; this helper is the
shared lookup so both
:meth:`EagerModelRunner.profile_activation_peak` and
:func:`arbi_serve.runtime.activation_profile.profile_engine` see
the same answer.

Returns the slab leading dim (``max_num_seqs + 1`` — one
more than the operator-facing usable capacity, the +1 being the
permanent zero-sentinel at row 0). The profiler clamps via
``.remainder(slab_rows)`` so synthetic indices land in
``[0, slab_rows)`` — including the sentinel, which is fine since
the sentinel is clean-zero and the profile forward's recurrent
update is meaningless either way.

Looks at every entry in :attr:`MultiStatePool._views` that is a
:class:`PerRequestRowPool` and returns the minimum slab leading
dim. Returns ``None`` for pools without a ``_views`` attribute
(legacy / mocked) and for pools whose every view is paged-KV / MLA.

The base class is the predicate on purpose. An enumeration of
concrete pool types here is a clamp that silently stops covering
the kind a new one adds: the profiler keeps synthesizing rows
against the OTHER pools' minimum, which is a valid index right up
until the new pool is the smaller one.

Per-row ``has_initial_state`` mask — the single preference point.

Precedence:

1. ``batch.has_initial_state`` when set — the precise per-row
   ``prompt_consumed > 0`` mask the engine batch build writes on
   persistent prefill steps (a view of the pinned
   ``PiecewiseBuffers`` slice, stable ``data_ptr``, no extra
   allocation here). A first chunk gets ``False`` — the conv path's
   ``prior * mask`` then does not depend on the slab row having been
   zero-cleared.
2. Derived ``seq_lens > 1`` otherwise (decode steps, non-persistent
   fallback, ad-hoc test batches). Exact for decode rows (any
   committed token ⇒ state is real); for a prefill first chunk with
   ``n > 1`` it reads ``True`` against a zero-cleared slab row,
   which is mask-equivalent (``prior_zeroed * 1 == prior * 0``).

The two branches are attributed on
``has_initial_state_precise_write`` (at the engine write site) and
``has_initial_state_seqlen_derived`` (here) so ``/v1/admin/flag_truth``
says which derivation a served config actually consumes.

Pulled out as a helper (not a sub-builder) because it's read by
:meth:`MetadataBuilder._finalize` — we don't want to spread one
derived field across sub-builder + finalize. See
:class:`MambaMetadataBuilder._finalize` for the call site.

Per-layer turbo-attn (tkv) attention op + slot-byte layout helpers.

Holds :class:`TkvAttnOp` (the per-layer :class:`torch.nn.Module` forward
that builds + caches a :class:`tkv.runtime.attention.TKVCore`), plus the
duck-typed :class:`_TkvImplShim` the AttendKernel wrappers introspect,
the pure-function :func:`compute_layout_offsets` slot-byte math, and the
calibration named-buffer roster :data:`_CAL_BUFFER_NAMES`.

Re-exported from :mod:`arbi_serve.backends.tkv_backend` so
``arbi_serve.backends.tkv_backend.{TkvAttnOp,compute_layout_offsets}``
keep resolving.

Minimal duck-typed "impl" object the AttendKernel wrappers introspect.

The decode_attend / turbo_prefill_attend / bypass wrappers were authored against the vLLM
backend's TKVAttentionImpl. Most fields there are irrelevant to
the standalone serve path; we expose the fields the wrappers
actually read.

Pure-function slot-byte layout math for the decode read side.

Returns ``(k_packed_off, k_norm_off, v_packed_off, v_norm_off,
k_packed_dim, v_packed_dim, total_bytes_per_token)``.

Delegates to :func:`tkv.kernels.cuda_compress_store.fused_cs_layout`
-- the SAME function the compress-store WRITE kernel bakes its
compile-time section sizes from (see that kernel's ``_get_module``).
Delegating instead of reimplementing the padding math makes
write/read layout drift structurally impossible: any independent
reimplementation can fall out of sync with ``fused_cs_layout``'s
padding rule, and a write/read offset mismatch makes decode read
V (and some K-norm bytes) from the wrong offsets, which manifests
as decode collapsing into degenerate repeated-token output.
Whatever padding rule the write kernel uses is exactly what this
returns, always, because it is the same call.

Move a codec's geometry tensors onto ``device``, preserving aliasing.

These tensors ALIAS on the host: the codec derives every rotation matrix
from one process-global Hadamard, so migrating attribute by attribute
forks one host tensor into four device copies per layer. Each distinct
host tensor is moved once and every attribute naming it is re-aliased onto
the result — the same read-only aliasing the host side and the MLA codec
path already rely on.

``shared`` extends the dedup across layers, which collapses a whole
model's rotation matrices onto one device tensor. Its entries pin the host
tensor so its identity stays valid for the map's life; scope it to one
build sweep so a device tensor can never outlive the pool it came from.
Returns the map, so a caller that passed none can inspect what was moved.

Per-layer TKV attention forward.

Builds + caches a :class:`TKVCore` lazily on first forward (so we
have a real device tensor in hand for the rotation-state build),
then routes every subsequent call through ``core.forward`` which
handles compress-store + bypass + decode internally.

Calibration named-buffer contract (compile-prep).
    AWQ + EXL3 quantized linears register their per-channel scales
    as :meth:`torch.nn.Module.register_buffer` entries (see
    ``arbi_serve/weight_quant/awq/linear.py`` L136-150 and
    ``arbi_serve/weight_quant/exl3/linear.py`` L60+). This op
    follows the same pattern: TKV centroids / boundaries /
    per-channel scales (see :data:`_CAL_BUFFER_NAMES`) are
    registered as named buffers on this :class:`torch.nn.Module`,
    with the underlying :class:`tkv.TkvCodec`'s same-named
    attributes aliased onto the buffer storage. Calibration reload
    (:func:`arbi_serve.engine.build.apply_calibration_to_tkv_ops`)
    is :meth:`torch.Tensor.copy_` into the registered buffers —
    the codec attribute and the buffer share storage, so the
    kernel reads the new values immediately and a future compiled
    trace's pointer-to-buffer-storage stays stable across reload.

    Rotation matrices and WHT sign vectors are deterministic from
    ``head_dim`` and a fixed seed, so they are not re-bound on
    calibration reload — they live on the codec as plain
    attributes. ``_build_core`` carries them onto the right device.

Bind this layer's per-head learned attention sink.

Stored fp32 on the op's device and published on the impl shim the
tkv attend wrappers introspect: decode routes it to
``tq_splitk_batch_decode(attn_sink=...)`` /
``tq_swa_splitk_batch_decode(attn_sink=...)``, MTP verify to
``tq_splitk_batch_mtp_decode(attn_sink=...)``, and Turbo prefill to
``turbo_prefill(learnable_sink=...)``.

Install per-layer codec ops, mirroring calibration tensors as buffers.

On install we (a) replace each registered calibration buffer's
storage with the codec's same-named tensor (so the buffer's
``data_ptr`` points at the codec's memory) and (b) alias the
codec attribute onto the buffer (so any future reassignment
via ``self.k_centroids = ...`` would bind both sides). After
this call, ``op.k_centroids is ops.k_centroids`` — both names
refer to the same tensor object — and a subsequent
``op.k_centroids.copy_(new)`` updates the value the codec
kernel reads.

Install this layer's group-2 VQ state on the codec + mirror it
as named buffers.

Registers (or re-registers) ``vq_k_centroids`` / ``vq_v_centroids``,
which are ``None`` on a plain :class:`tkv.TkvCodec` and so have
no placeholder to :meth:`torch.Tensor.copy_` into the first time a
vq2 bundle lands. Registration allocates a fresh tensor, so this
CHANGES ``data_ptr`` and is only safe before the core is built and
graphs are captured — the boot install path.
:func:`arbi_serve.engine.build.apply_calibration_to_tkv_ops` refuses
to reach this from a live reload (the resident KV pages were encoded
by the old codec) and uses :meth:`update_calibration`'s in-place
copy for a same-config codebook refresh instead.

``vq_mode`` is the derived family switch the compress/decode/prefill
kernels select on; ``vq_mode_k`` / ``vq_mode_v`` say which side's
dequant reads the VQ table.

Reload calibration tensors via in-place :meth:`torch.Tensor.copy_`.

Each kwarg, when not ``None``, is copied into the matching
registered buffer in place. The buffer's ``data_ptr`` does not
change — a future compiled trace that captured the buffer
pointer at trace time still reads the new values without
retracing. ``set_codec`` must have been called first (the
buffers must be sized).

Sources may be Python lists / NumPy arrays / CPU tensors / GPU
tensors of any compatible dtype; ``copy_`` handles the cast.
Shape must match the live buffer.

Return True iff a calibration reload since the last consume
invalidated this op's captured-graph dependencies; resets the flag.

Called by :func:`engine.swap_admin.areload_calibration` after the
per-layer apply loop. A True return triggers a one-shot drop of
every captured graph (whose data_ptrs reference the OLD per-layer
:class:`tkv.runtime.rotation.RotationState`'s derived tensors).

Update the per-step MTP token-count expected by the verify kernel.

The :class:`DecodeAttend` decode wrapper checks
``tokens_per_seq == core._mtp_block_m`` before delegating to the
:class:`MTPFusedAttend` split-K kernel. ``tokens_per_seq`` is the
per-step ``K + 1`` for each row in the verify batch (K drafts
plus the bonus / last-committed token), which the engine's
:class:`MtpStepPlan` knows but the static constructor arg
cannot — different requests run at different K and the same
attn op is reused across steps.

This setter mutates both the ``TkvAttnOp``'s own field (so a
first forward after this call constructs the underlying
:class:`TKVCore` with the right value) and the live
``_core._mtp_block_m`` when ``_core`` already exists. The engine
flips this from :func:`arbi_serve.engine.run_step._mtp_verify_step`
once per verify forward, before the forward call.

Eagerly construct the per-layer :class:`TKVCore` at boot.

The decode/prefill hot path runs the model under
``torch.compile`` (Inductor): the model's ``forward`` is traced
and the ``arbi_serve::tkv_attention`` custom op is emitted
directly into the compiled graph. The Python-level
:meth:`forward` wrapper that lazily builds ``_core`` (the
``if self._core is None: self._build_core(...)`` branch) is
therefore never executed at steady state — Inductor calls the
op's real-impl (:meth:`_dispatch_through_custom_op`) straight
from generated code, which asserts ``_core is not None``.

Building the core eagerly at boot (right after codecs are
installed, before the first compiled forward) makes the core a
boot-time invariant rather than a forward side effect. This is
the correct place: ``set_codec`` has run, calibration (if any)
has been applied, and ``self.device`` is known. Idempotent —
a no-op once ``_core`` exists.

Publish the per-step ``batch_meta`` onto the op's side-channel.

The TKV attention runs through the ``arbi_serve::tkv_attention``
custom op, whose Tensor-only signature cannot carry the rich
:class:`AttnPagedKVMeta` (TQRunState + the ``_tq_*`` mirrors the
AttendKernel wrappers introspect). :meth:`forward` stashes
that meta on ``self._call_batch_meta`` just before dispatching,
but under ``torch.compile`` the model's forward is traced and
the custom op is emitted directly into the Inductor graph, so the
Python ``forward`` wrapper (and its ``self._call_batch_meta =
batch_meta`` side effect) never runs at steady state. The op
real-impl then asserts ``_call_batch_meta is not None`` and dies.

The engine therefore publishes the per-step meta here eagerly,
from :meth:`arbi_serve.runtime.model_runner.ModelRunner._run_model_forward`,
outside the compiled region. All PAGED_KV layers in a step share
the same meta object (only ``layer_idx``, passed as an op arg,
differs), so the engine calls this once per TKV op with
``batch.attn_meta`` before the forward.

Decorate and publish a sparse block-table call over ``num_seqs`` rows.

A sparse call gives every query token its own row, so ``num_seqs``
and ``num_tokens`` are the token count of the step. The pool is
sized on first use and re-sized only outside a capture region; a
caller that already filled the CSR triplet on the metadata has it
copied into the pool the TKV decode wrapper reads.

Drop runtime objects that may retain views of a KV slab.

A rebuilt attention op constructs a fresh core before serving. The
codec and calibration buffers remain attached to this module.

Quantize-then-store K/V into the paged TKV slab and run attention into ``output``.

Stashes ``batch_meta`` (carrying the TQRunState and ``_tq_*``
mirrors) on a side-channel, then dispatches the registered TKV
attention op. The core is built eagerly at boot
(:meth:`ensure_core_built`) — under ``torch.compile`` this
``forward`` is traced away and does not run at steady state, so
the engine also publishes the per-call meta via
:meth:`publish_call_meta` before the compiled forward. The eager
guards below keep the non-compiled (test / eager) path correct.

Real-impl side of ``arbi_serve::tkv_attention``.

Dispatches into :class:`tkv.runtime.attention.TKVCore`'s
``forward`` with the per-call ``batch_meta`` recovered from the
:class:`TkvAttnOp` instance stash. The custom op signature
only carries Tensors + Python scalars; the rich attn-meta
dataclass (TQRunState mirrors, Turbo prefill scratch ints) goes
through this side-channel.

Boot invariant: the per-layer core is built eagerly, never here.
Under ``torch.compile`` Dynamo traces the Python :meth:`forward`
wrapper away and emits a direct ``torch.ops.arbi_serve.
tkv_attention`` call into the Inductor graph, so this real-impl
is what runs at steady state. Both the core and the per-step
``_call_batch_meta`` are therefore established outside this
boundary:

  * ``_core`` — built by :func:`ensure_tkv_cores_built` inside
    ``build_active`` (run at boot and on every backend swap via
    ``critical.py``), after codec/calibration install and before
    the o_proj fold. We do not build it here: allocating a core
    inside the op real-impl can fire mid-cudagraph-capture, which
    is illegal (``cudaErrorStreamCaptureInvalidated``).
  * ``_call_batch_meta`` — published per step by
    :meth:`publish_call_meta` from ``ModelRunner._run_model_forward``
    (the single forward call site), eagerly and outside the
    compiled region.

A ``None`` for either here means that boot/per-step invariant was
violated — fail loud rather than silently papering over it.

``{"ancestor_mask": ...}`` for a tree verify block, else empty.

Per-route binding of
:func:`~arbi_serve.spec_decode.tree_spec.verify_block_mask_kwargs`,
which owns the shape gate, the empty-dict contract and the
address-stable tensor. What belongs here is the capability probe
for THIS route: the codec reaches the kernel through
``TKVCore.forward``, a different call site into the shared Turbo prefill
mainloop than the bf16 bypass's attend, so a mask threaded to one
does not travel to the other.

The batch width comes off the published per-step meta rather than
a stored attribute: it is the same ``seq_lens`` the core's own
verify-shape test reads, so the shape gate here and the dispatch
there cannot disagree.

turbo-attn (tkv) attention backend — paged uint8 codec storage.

Wraps :class:`tkv.runtime.attention.TKVCore` (compress + bypass +
decode) behind the engine-side :class:`AttentionBackend` /
:class:`AttnOp` / :class:`MetadataBuilder` triplet. Kernel dispatch
uses :class:`DecodeAttend` for decode and :class:`TurboPrefillAttend`
for both prefill and the first-chunk bypass fast path (native inline
attend — no host flash_attn dependency).

This module owns:
  - per-step :class:`TQRunState` construction;
  - Turbo prefill scratch buffer management;
  - per-layer :class:`TKVCore` build + lazy device migration of codec
    tensors.

The compress + attend math lives in :class:`TKVCore`.

Refuse to boot on a stride-blind ``compute_page_metadata`` kernel.

Runs the kernel once on a narrow view of a wide block_table
((4,1) view of a (4,8) buffer, rows [1..4]) and requires the exact
page list back. The engine's persistent-decode batches pass exactly
this layout (``pb.block_table[:B, :max_pages]``); a kernel that
addresses rows by ``size(1)`` instead of ``stride(0)`` returns
``[1, 0, 0, 0]`` — every non-first sequence of a multirow batch
silently reads the null page (zero dequant norms -> NaN or
silently-zero attention; the "TP2 GDN prefill-capture NaN").
Behaviour-verified, not version-string-verified, so a stale tkv can
never serve. One tiny launch at builder construction; never on the
hot path.

turbo-attn (tkv) attention backend descriptor (the quantized codec mode).

KV cache layout: paged uint8 slab of shape ``(num_pages, page_size,
slot_bytes)`` per layer. Slot-byte budget is determined by
``(num_kv_heads, head_dim, k_bits, v_bits)``.

State kind (``PAGED_KV``), the single fused plane, and MTP threading
are inherited from :class:`PagedKvBackend`.

Resolve ``(k_bits, v_bits)`` for a global layer index.

Returns the per-layer smart-mix entry from the calibration
bundle when present, else the backend's fallback ``(k_bits,
v_bits)`` (uniform mode, or a layer the bundle doesn't list).

True iff per-layer bits actually vary across the loaded table.

Uniform tables (or no table) report False so the pool keeps the
single stacked-slab fast path; heterogeneous tables report True
so the pool builds per-layer slabs sized to each layer's bits.

Byte size of one quantized TKV slot for the configured k/v bit-widths.

Routed through :func:`compute_layout_offsets` (its ``total_bytes``
element) rather than ``tkv.runtime.layout.slot_bytes`` directly —
see that function's docstring: this must always agree with what
the compress-store write kernel actually writes, and duplicating
the padding formula independently lets the two drift (e.g. an
allocator sizing 544 B/slot while the write kernel writes
576 B/slot, corrupting the KV cache).

Byte size of one quantized TKV slot for a specific layer.

Uses that layer's smart-mix ``(k_bits, v_bits)`` so per-layer
slabs in the composite pool are sized to their actual footprint
(a K4V4 layer's slot is smaller than a K6V6 layer's slot). See
:meth:`_slot_bytes` for why this goes through
:func:`compute_layout_offsets` rather than
``tkv.runtime.layout.slot_bytes``.

Per-paged-layer slot byte sizes, or None when uniform.

Returns one ``slot_bytes`` value per entry in
``paged_layer_indices`` (same order) when the smart-mix table
makes them differ; returns ``None`` when every paged layer
resolves to the same slot size (the pool then keeps the single
stacked-slab fast path). Drives the composite per-layer-offset
cache layout.

Per-layer per-token byte cost == one quantized TKV slot.

When ``layer_idx`` is given, the slot is sized at that layer's
smart-mix ``(k_bits, v_bits)`` — matching what the pool's
per-layer slab actually allocates (``slot_bytes_for_layer``).
The budget path (``per_page_bytes_for_paged_kv``) passes it so
the KV page-count is derived from the real per-layer footprint
instead of charging every layer at the bundle's max bits.
Without it, smart-mix over-reserves per page and the engine
gets fewer KV pages than the compression permits.

``layer_idx is None`` (legacy callers) keeps the uniform /
max-bits fallback for backward compatibility.

Raw-bf16 attention on the turbo-attn kernels (``tkv-bypass``).

Uncompressed bf16 KV stored in the same fused single-plane paged slab the
TKV codec uses, but with no codec: no rotation, no quantization, no
calibration. Attention runs on the turbo-attn kernels — the Turbo prefill kernel
(:func:`tkv.kernels.cuda.prefill.turbo_prefill_dispatch.turbo_prefill`)
and the split-K paged flash-decode
(:func:`tkv.kernels.turbo_attn_simt.tq_splitk_batch_decode`) — through
the dep's raw-bf16 bypass loaders
(:class:`tkv.kernels.loaders.BypassLoader` /
:class:`tkv.kernels.loaders.BypassDecodeLoader`).

This is the whole-model bf16 backend for ``head_dim > 256`` models (Gemma
4: 28 sliding head_dim-256 layers + 7 full head_dim-512 layers): every
paged layer runs on it, mixed head dims handled by the pool's
``per_layer_geometry`` driving per-layer fused-slab geometry. The split-K
paged kernel is a plain CUDA launch (no host sync). MTP verify uses split-K
through its register-resident width ceiling and the Turbo prefill kernel above that ceiling. Both
routes are cudagraph-capturable at every supported head_dim, including the
head_dim=512 full-attention layers (see
:meth:`TkvBypassBackend.supports_capturable_mtp_verify`). The Turbo prefill
request-prefill route runs eager by design.

Fused per-slot bf16 layout (single plane, ``kv_num_planes() == 1``):
``[K(H_kv, D) | V(H_kv, D)]``, ``total_bytes_per_slot = 2*H_kv*D*2``,
``k_off = 0``, ``v_off = H_kv*D*2``. A fused
``(num_pages, page_size, 2*H_kv*D)`` bf16 slab viewed as
``(num_pages, page_size, 2*H_kv, D)`` splits into
``K = [:, :, :H_kv, :]`` and ``V = [:, :, H_kv:, :]`` — non-contiguous
slices the scatter and both turbo-attn loaders read correctly.

State kind: :attr:`StateKind.PAGED_KV`.

Widest decode-class ``block_m`` the q-prescale scratch must cover.

``mtp_n_draft`` is the served step width — the chain's depth K, or a
tree's NODE count; verify runs ``tokens_per_seq = that + 1``. ``None``
(or a non-positive value) means the width is not known here, and the
historical floor applies.

Sizing from the served depth rather than a literal is what keeps the
scratch and the admission gate in agreement — they disagreed before,
and the disagreement was invisible because the overflow path is a
correct-but-slower fallback rather than an error.

Served MTP step width for this boot, read from the runtime flag.

The chain's depth K, or a TREE's node count — verify lays out one
query row per node, so the chain flag under-counts a tree's block and
everything sized from it comes out short.

Read from the flag rather than threaded through the backend protocol:
that protocol is shared with every other backend and none of them
needs K. ``None`` when the flag is unreadable, which is what
:func:`q_prescale_block_m_ceil` turns into its floor.

The BUDGET that reserves for this scratch resolves the width here too.
A literal mirrored into the budget is what let a rename move the
allocator's meaning while the reserve kept the old number.

Per-step metadata for the tkv-bypass backend.

Inherits the shared ``(CuSeqlens, BlockTable)`` sub-builder pair and
:class:`AttnPagedKVMeta` construction from
:class:`PagedKvMetadataBuilder`; the pair covers ``cu_seqlens_q/k``,
``seq_lens``, ``block_table``, ``slot_mapping``, ``max_*``. This
builder adds the persistent no-sync CSR page-metadata buffers the
split-K decode loader consumes (``indptr`` / ``indices`` /
``last_page_len``). Those are built once per step via
:func:`tkv.runtime.page_metadata.compute_page_metadata_cuda` into
pre-allocated tensors (mirrors
:meth:`TkvMetadataBuilder.cudagraph_capture_step`) — never per-layer,
never with ``.item()``.

It also owns the persistent q-prescale scratch pair
(``q_scratch_f32`` / ``q_scratch_bf16``) the decode / MTP-verify
attend consumes (:meth:`tkv.runtime.bypass_bf16.BypassBf16Attend._prescale_q`)
— the bypass analog of TKV's budgeted ``TQBufferPool`` q tiles. Flat
buffers sized for the widest decode-class step
(``min(max_num_tokens, max_num_seqs × q_prescale_block_m_ceil(K)) ×
max(H_q_per_rank × head_dim)`` over paged layers — per-layer views are
sliced by the attend), allocated under the engine's named mem-pool
when provided so bypass transients stay off the default heap.

The persistent triplet + scratch pair are stashed on the meta
(``_bt_indptr`` / ``_bt_indices`` / ``_bt_lpl`` / ``_bt_q_scratch_*``)
so the per-layer decode op reads them off the published meta with zero
host sync.

Whether THIS attend instance dispatches its verify branch to the Turbo prefill kernel.

Older tkv reached the paged Turbo prefill kernel only from ``attend``'s PREFILL
branch, which is why a verify block wider than the split-K ceiling is
disguised as a prefill below. Newer tkv dispatches its own verify
branch to ``mtp_verify_prefill`` — and that branch is the only one that
accepts an ancestors-only tree mask, because the mask indexes
block-local node ids and a prefill chunk has no block. There the
disguise must stop.

Per INSTANCE and not per class: the same tkv declines the Turbo prefill route for a
sliding-window layer and for a subclass that owns ``mtp_verify`` (the
sink-carrying one), and those layers still need the prefill-branch
disguise to reach the Turbo prefill kernel at a width split-K cannot serve.

Per-layer raw-bf16 attention on the turbo-attn kernels.

Step:
  1. Scatter new K / V into the fused paged slot (bf16 views; skipped
     for Gemma-4 YOCO ``num_kv_shared_layers`` layers).
  2. Decode (one query per row, ``N == B``) → split-K paged
     flash-decode via :class:`BypassDecodeLoader` (q pre-scaled,
     ``sliding_window`` passed); else prefill → Turbo prefill kernel via
     :class:`BypassLoader`.
  3. Copy the result into ``output[:N]``.

The decode path issues no host sync, so it is cudagraph-capturable.

Raw-bf16 turbo-attn paged backend (the bypass codec mode).

Per-layer KV layout: one fused rank-3 bf16 plane
``(num_pages, page_size, 2*num_kv_heads*head_dim)`` per layer —
``[K | V]`` interleaved per slot (``kv_num_planes() == 1``, like
TKV's single packed slot, but uncompressed bf16). Mixed head dims
(Gemma 4: 256 + 512) drive per-layer fused-slab geometry through the
pool's ``per_layer_geometry``. State kind (``PAGED_KV``), the single
fused plane, and MTP threading are inherited from
:class:`PagedKvBackend`.

Attach the persistent triplet (sliced to live extents) + the
q-prescale scratch pair to ``meta``.

The published meta carries the ``_bt_*`` references into the
decode op's custom-op real-impl with zero host sync.

Populate the CSR triplet on a meta built outside :meth:`build`.

The MTP drafter chain (and the seed-drafter forward) build bare
:class:`AttnPagedKVMeta` objects (one query per row) that never
route through :meth:`build` / :meth:`_finalize`, so ``_bt_indptr``
/ ``_bt_indices`` / ``_bt_lpl`` stay ``None`` and the decode op's
:class:`BypassDecodeLoader` ``from_bf16`` crashes on
``None.contiguous()``. This runs the same no-sync
``compute_page_metadata_cuda`` sequence :meth:`_finalize` runs, but
reads the live extents off the meta's own ``seq_lens`` /
``block_table`` (a :class:`ScheduledBatch` is not in hand on the
drafter path). One query per row (``N == B``). Mutates + returns
``meta``.

``bypass_safe`` is accepted for interface parity with
:meth:`TkvMetadataBuilder.decorate_external_meta` and ignored —
raw bf16 has no compress-store, so there is no first-chunk
bypass-eligibility flag to thread. The drafter dispatches to
whichever PAGED_KV builder exposes this method with one uniform call.

Publish a sparse call whose CSR triplet is already populated.

The bypass decode path reads ``_bt_indptr`` / ``_bt_indices`` /
``_bt_lpl`` straight off the metadata and derives its row count from
``seq_lens``, so the shape arguments carry no work here.

Bind this layer's per-head learned attention sink.

Swaps the stock attend for
:class:`~arbi_serve.backends._bypass_sink_attend.SinkBypassBf16Attend`,
which forwards the fp32 sink into ``tq_splitk_batch_decode`` /
``tq_splitk_batch_mtp_decode`` / ``turbo_prefill``. Layers
without a sink keep the stock attend.

Record the per-step MTP verify token count (``K + 1``).

Driven by :func:`arbi_serve.spec_decode.mtp_verify._mtp_block_m_scope`
once per verify forward (and by the capture sweep), exactly as it
drives :meth:`TkvAttnOp.set_mtp_block_m`. This value is what lets the
op tell the MTP verify shape (``N == B * mtp_block_m`` with
``mtp_block_m > 1``) apart from a chunked prefill (also ``N != B``):
the verify shape routes to the split-K paged MTP kernel
(:meth:`~tkv.runtime.bypass_bf16.BypassBf16Attend.mtp_verify`), a chunked prefill to the
Turbo prefill kernel.

Publish the per-step ``batch_meta`` onto the op's side-channel.

Under ``torch.compile`` the Python :meth:`forward` is traced away
and the custom op is emitted directly into the Inductor graph, so
the engine publishes the per-step meta here (eagerly, outside the
compiled region) and the op real-impl reads it back via
:func:`get_layer_dispatch`. Mirrors :meth:`TkvAttnOp.publish_call_meta`.

Scatter K/V into the fused slab and run turbo-attn attention into ``output``.

Stashes ``batch_meta`` on the side-channel (eager / test path) and
dispatches the registered ``arbi_serve::tkv_bypass_attention`` op.

``output`` is written through a per-head view, so it must be rank-3
``(num_tokens, num_heads, head_dim)``. A flat buffer of the same
element count is not interchangeable and is rejected here: the kernel
would otherwise fill it through a mismatched view and return
non-deterministic garbage with nothing raised.

Real-impl side of ``arbi_serve::tkv_bypass_attention``.

Scatters new K/V into the fused bf16 slot, then attends *in place*
over the paged cache: decode (``N == B``) via the split-K paged
flash-decode (:class:`BypassDecodeLoader`, q pre-scaled), prefill
via the Turbo prefill kernel (:class:`BypassLoader`). No host sync on
the decode path.

The persistent CSR triplet (``indptr`` / ``indices`` /
``last_page_len``) is recovered from the per-step meta published
on this op via :meth:`publish_call_meta`.

``{"ancestor_mask": ...}`` for a tree verify block, else empty.

Per-route binding of
:func:`~arbi_serve.spec_decode.tree_spec.verify_block_mask_kwargs`,
which owns the shape gate, the empty-dict contract and the
address-stable tensor. What belongs here is the capability probe
for THIS route: the bypass reaches the kernel through
``BypassBf16Attend.attend``, and the codec route reaches it
through a different entry whose signature is a separate question.

Builder-owned q-prescale scratch for this call, or empty.

The metadata builder stashes its persistent flat scratch pair on
the meta (``_bt_q_scratch_*``); this layer's slice needs
``num_tokens × H_q × D`` elements. A meta without the stash
(bare test metas) or a rare step wider than the sized scratch
(eager block_m past the sized ceiling) falls back to
the attend's allocating prescale — correct either way, just not
pooled. Host-side ints only; capture-safe.

Per-layer fused slab shape ``(num_pages, page_size, 2*num_kv_heads*head_dim)`` bf16.

Rank-3 (pages on dim 0). The trailing dim packs ``[K(H_kv, D) |
V(H_kv, D)]``; the attn op views it as ``(num_pages, page_size,
2*H_kv, D)`` and splits K / V by head slice.

Per-layer per-token KV byte cost: fused K+V, 2 bytes/bf16 elem.

``layer_idx`` is accepted for interface parity with the smart-mix
TKV backend and ignored — bypass stores every layer at the same
uniform bf16 footprint, so page cost does not vary per layer.

Construct the tkv-bypass metadata builder (sub-builders +
persistent CSR + pooled q-prescale scratch).

The scratch pair is sized to the widest per-layer q extent —
``max(H_q_per_rank × head_dim)`` over paged layers (mixed head
dims, e.g. Gemma-4 256/512, share one max-sized flat pair sliced
per layer) — and allocated under ``mem_pool`` (the engine's
``scratch.attn_codec`` named pool) when provided, mirroring how TKV
routes its ``TQBufferPool``.

Construct the per-layer tkv-bypass attn op.

``dispatch_owner`` scopes the op in the process-global dispatch
registry. The MAIN model uses the default ``"main"`` and registers
under its plain ``layer_idx``. An external draft model passes
``"drafter"`` so its ops register (and dispatch) under
``DRAFTER_DISPATCH_LAYER_BASE + layer_idx`` — a disjoint key range, so
the drafter never clobbers the main model's compiled-forward dispatch.

TKV per-layer (k_bits, v_bits) resolution from the calibration bundle.

This is the single source of truth for "what (k_bits, v_bits) does layer
N quantize K/V to".

Dispatch semantics (matches the project-wide ``TKV_BITS`` contract):

  * Float ``TKV_BITS`` (e.g. ``"4.0"``) → smart-mix. Reads the calibration
    bundle at ``TKV_CALIBRATION_FILE``, looks up
    ``byte_budget_table[<bpe>].layers[<idx>] = {"k_bits", "v_bits"}`` and
    returns a per-layer ``{layer_idx: (k_bits, v_bits)}`` map. Each
    paged-attention layer is sized + dispatched at ITS OWN bit widths.
  * Int ``TKV_BITS`` (e.g. ``"4"``) → uniform. No bundle lookup; every
    layer uses ``K=V=N``. Kept for unit tests / CI smoke; production
    should use the float form + a bundle.

Per the layout helper :func:`tkv.runtime.layout.slot_bytes`, k_bits /
v_bits may be any integer in ``[2, 8]`` — the codec, compress-store and
decode kernels handle arbitrary per-layer-varying 2..8.

Fail-loud contract: when smart-mix is active, every misconfiguration is
a hard error (missing bundle, wrong schema, missing BPE entry, bits out
of range) — never a silent fallback to a default width.

Parse ``TKV_BITS`` → ``(value, is_smart)``.

``value`` is the parsed float (or ``None`` when the env var is
unset). ``is_smart`` is True when the operator wrote a decimal /
fractional value (``"4.0"``, ``"4.5"``) selecting per-layer smart-mix
via the byte_budget_table; pure ints (``"4"``) select uniform K=V=N.

Fail loud on an unfinalized (stage-1) bundle unless explicitly allowed.

A bundle with ``finalized != True`` skipped the drift-scale/solver finalize
step and can silently degrade long-form output (drift into repetition
loops). By default smart-mix refuses to load it. The override env
``ARBI_ALLOW_UNFINALIZED_CALIBRATION=1`` permits it behind a loud banner for
dev / bring-up only.

Return the per-layer ``{layer_idx: (k_bits, v_bits)}`` smart-mix map.

Returns ``None`` when smart-mix is not active (``TKV_BITS`` unset or
integer form) — the caller falls through to uniform bits. When
smart-mix is active every misconfiguration raises (fail-loud).

Pull ``byte_budget_table[<bpe>].layers`` into a per-layer bits map.

Requires the current schema and a finalized bundle (arbi-serve-only
gates), then delegates the ``byte_budget_table`` walk to
:func:`tkv.runtime.calibration.resolve_layer_bits_for_bpe` — the same
resolver the vLLM integration calls directly — so the two can't drift
on the table shape. Every (k, v) is additionally range-validated.
Returns ``{int(layer_idx): (k_bits, v_bits)}``.

Return ``(max_k, max_v)`` across the BPE entry's per-layer table.

Used to size the backend spec string / advertised cache slot when
callers still need a single scalar pair (e.g. the ``tkv-k{K}v{V}``
spec name). The actual cache slabs are sized per-layer.

The single source of truth for tkv ``(k_bits, v_bits)``: ``TKV_BITS``.

A tkv config is declared exactly one way — one number + a mode:

  * Integer ``TKV_BITS`` (e.g. ``"4"``) → **uniform**: ``K = V = N``.
  * Float ``TKV_BITS`` (e.g. ``"4.5"``) → **smart-mix** at that
    bytes-per-element, resolved per layer from the calibration bundle at
    ``TKV_CALIBRATION_FILE``. The pair returned here is the
    ``(max_k, max_v)`` across that bundle's table — used only to size
    the cache slab / name the derived spec; the genuine per-layer bits
    come from :func:`load_layer_bits_from_env` (same env + bundle).

There is no other way to declare tkv bit-width: not the ``tkv-k{K}v{V}``
spec name, not ``TKV_K_BITS`` / ``TKV_V_BITS``. Any caller that carries
explicit bits in a spec string (an HTTP ``attention_backend`` request, a
directly-built ``ServerConfig``) is validated against this resolver and
rejected on mismatch (see :func:`arbi_serve.backends._parse_tkv`). Because
the declaration is a single scalar, an asymmetric ``k != v`` can never be a
*user* choice — only a calibrator-produced per-layer detail of a smart-mix
bundle.

Defaults to uniform ``4`` when ``TKV_BITS`` is unset (CI / smoke). Raises
``ValueError`` on any misconfiguration (fail loud, never silent).

Implementation: this delegates to the canonical
:func:`tkv.runtime.calibration.resolve_declared_tkv_bits` so the
declaration semantics live in one place. We only adapt its
``RuntimeError`` to arbi-serve's ``ValueError`` contract (the CLI
+ :func:`arbi_serve.backends._parse_tkv` handle ``ValueError``).

Per-step metadata builder for the turbo-attn (tkv) attention backend.

Holds :class:`TkvMetadataBuilder` and the four arbi-serve-side flag-truth
``path_counter`` fire sites its :meth:`~TkvMetadataBuilder._finalize`
increments (counters are created here, in the module that fires them).

Re-exported from :mod:`arbi_serve.backends.tkv_backend` so
``arbi_serve.backends.tkv_backend.TkvMetadataBuilder`` keeps resolving.
The row-split resolver / preflight helpers stay in ``tkv_backend`` (a
handful of tests poke them there); this module imports them lazily inside
the methods that use them to avoid an import cycle.

Per-step metadata builder for the TKV backend.

Inherits the shared ``(CuSeqlens, BlockTable)`` sub-builder pair and
:class:`AttnPagedKVMeta` construction from
:class:`PagedKvMetadataBuilder`. Wraps the persistent
:class:`TQBufferPool` and runs the per-step
``compute_page_metadata_cuda`` + ``compute_bypass_safe`` calls.
Returns a populated :class:`AttnPagedKVMeta` (the kind-specific
metadata dataclass) carrying the :class:`TQRunState` in its
``tq_run_state`` field.

The inherited sub-builders cover the standard paged-KV addressing
(``cu_seqlens_q/k``, ``query_start_loc`` alias, ``block_table``,
``slot_mapping``, ``seq_lens``, ``max_*``); :meth:`_finalize`
handles every TKV-specific bit:

  * :class:`TQBufferPool.ensure` — re-shape per-step scratch.
  * :meth:`cudagraph_capture_step` — re-populate
    ``indptr/indices/lpl`` from the persistent ``seq_lens`` /
    ``block_table`` (called once at build, again on every replay
    via the cudagraph capture path).
  * ``compute_bypass_safe`` — CPU-side bypass eligibility flag.
  * Turbo prefill scratch pre-warm (``prefill_cu_seqlens_k_i32``,
    ``prefill_block_table_i32``).
  * :class:`TQRunState` construction + the seven ``_tq_*`` meta
    attribute mirrors that the AttendKernel wrappers introspect.

Re-populate the page-metadata kernel scratch from the
(captured / persistent) batch buffers.

Called once during build and again on every replay (via the
cudagraph capture path's ``getattr(builder,
"cudagraph_capture_step", None)`` hook). Reads
``batch.seq_lens`` and ``batch.block_table`` (persistent
kernel-input tensors that ``CapturedGraph.replay`` ``copy_()``s
each step) and writes:

  * ``self.tq_pool.{indptr, indices, lpl}`` — the split-K
    decode/verify page metadata (the shared base's
    :meth:`_fill_page_csr`, via ``super()``).
  * ``self.tq_pool.{prefill_cu_seqlens_k_i32, prefill_block_table_i32}``
    — the Turbo prefill verify route's per-step KV extents (the int32
    cumsum of ``seq_lens`` + int32 block-table cast). A high-K
    verify chunk (block_m > the split-K register ceiling) routes
    through the Turbo prefill verify path, whose kernel reads these
    two buffers (attached to the per-step meta as
    ``_tq_cu_seqlens_k_i32`` / ``_tq_block_table_i32`` by
    ``decorate_metadata_for_tkv`` at build time). Without this
    re-derivation inside the captured region, the captured Turbo prefill
    kernel would read the capture-time synthetic seq_lens (the
    full max-context page coverage the capture batch uses) on
    every replay — its per-row K-loop trip count baked at that
    huge value → a catastrophic KV over-read that pins the GPU
    at 100% for the whole captured forward (a hang, not a wrong
    answer; on TP2 it reads as a lockstep "deadlock" because the
    rank-0 host loop blocks on the never-finishing replay). The
    split-K page metadata above is already refreshed here; the
    Turbo prefill extents must be too, for the same reason and from the
    same live buffers.

The in-place writes mirror ``decorate_metadata_for_tkv``'s
Turbo-prefill hoist (``tkv.runtime.page_metadata``) exactly — same
slices, same ``torch.cumsum(..., out=...)`` / int32 ``copy_()``,
no fresh allocation (capture forbids cudaMalloc).

Pre-condition: ``self.tq_pool.ensure(B, max_pages, N)`` and
``ensure_prefill(B, max_pages)`` were called once before
capture (the first :meth:`_finalize` does both; the capture path
re-asserts ``ensure`` on the synthetic batch before invoking this
hook), so the Turbo prefill buffers are already sized for this (B, pages).

Apply the full TKV per-step wiring to a meta built outside the
normal :meth:`build` path (the MTP drafter chain).

The drafter's :func:`build_draft_step_metas_from_lists` produces
bare :class:`AttnPagedKVMeta` objects (one query per row) that
never route through :meth:`build` / :meth:`_finalize`, so the
``_tq_*`` page buffers stay ``None`` and the decode_attend decode
kernel crashes on ``compute_page_metadata(... None ...)``. This
runs the same sequence :meth:`_finalize` runs — ``ensure`` →
``compute_page_metadata_cuda`` → :class:`TQRunState` →
``decorate_metadata_for_tkv`` — but reads the live extents off
the meta's own ``seq_lens`` / ``block_table`` / ``cu_seqlens_q``
(a ``ScheduledBatch`` is not in hand on the drafter path).

One query per row: ``N == B`` (the drafter decodes one token per
request per chain step). Mutates + returns ``meta``.

``bypass_safe`` is host-known from ``pre_lens`` at the caller
(:func:`decorate_draft_metas`) — passed in so this path
runs zero D2H sync.

Construct the per-step :class:`TQRunState` from the persistent pool.

Slices every :class:`TQBufferPool` buffer to the live
``(num_seqs, num_tokens)`` extents. The only per-call variable
is ``bypass_safe`` (host-computed on the build path, host-known
on the drafter path); every other field is a fixed view onto the
pool. Shared verbatim by :meth:`_finalize` and
:meth:`decorate_external_meta`.

Attach the run-state + persistent ``_tq_*`` mirrors to ``meta``.

``query_start_loc`` is the same ``cu_seqlens_q`` tensor — needed
by the AttendKernel wrappers' getattr-based introspection (see the
MetadataView protocol). ``decorate_metadata_for_tkv`` ensures the
decode pool covers this step, attaches the persistent
``_tq_<name>`` buffers (cu_seqlens_k_i32, block_table_i32, indptr,
indices, last_page_len) for the decode_attend attend wrapper, and
hoists the Turbo-prefill cumsum + int32 block_table cast off the
per-layer hot path into a single per-batch in-place write.
arbi-serve always runs the standalone
CUDA decode engine, so ``engine_is_cuda_standalone=True``. Shared
by :meth:`_finalize` and :meth:`decorate_external_meta`.

``arbi-serve bench`` subcommand — clean per-request perf benching.

Boots an in-process :class:`Engine`, sends N synthetic requests with
per-step timing capture enabled, and prints summary distributions for:

  - TTFT (raw and first-call-cost-subtracted)
  - TPOT (raw and warm-only)
  - Queue wait
  - Per-label first-call cost breakdown

The ``--warmup N`` flag drops the FIRST N requests entirely from the
histograms — separate from first-call cost subtraction; this is
request-level steady-state filtering useful when even the second /
third request might have a long tail (autotune table caches usually
take a few hits to settle).

The harness is sequential by default (one request at a time, deterministic
TTFT) and supports a ``--concurrency C`` flag to fan out C concurrent
requests at once for steady-state throughput measurement.

Output is plain text (one block per metric); pass ``--json`` for a
machine-readable dump.

Speculative decoding (MTP) bench: pass ``--enable-mtp --mtp-n-draft K``
to attach the model's bundled draft head (or a separate draft model via
``--mtp-draft-model-path``) at speculation depth K. The summary then includes
per-request ``mtp_proposed`` / ``mtp_accepted`` distributions plus an
``accept_rate`` histogram so an operator can compare TPOT-with-MTP-on
against TPOT-with-MTP-off in two adjacent runs.

The bench implementation is split across sibling modules by concern —
:mod:`arbi_serve.bench_run` (prompt/config building + the async engine
driver) and :mod:`arbi_serve.bench_report` (stats, summary aggregation,
text rendering) — while every public name stays importable from here.

BENCH-GUARD — wrap any bench invocation so its number carries its own evidence.

    python -m arbi_serve.bench_guard --gpus 0,1 --stamp stamp.json -- <bench cmd...>

A bench that drives a server ALREADY RESIDENT on those GPUs -- a flag A/B
against a live engine -- adds ``--standing-target --endpoint <url>``, which
lifts the GPU-residency refusal and nothing else.

Wraps an ARBITRARY bench command (turbo-attn's ``bench_mtp_decode.py``, ``vllm
bench serve``, anything) and around it:

  1. PREFLIGHT the host. A run queue already deeper than half the cores, a
     foreign process resident on the GPU, a hardware throttle already asserted --
     any of these and we REFUSE TO BENCHMARK. We do not emit a number and warn;
     we emit no number. A warning next to a plausible-looking figure gets
     screenshotted without the warning.
  2. SAMPLE the machine for the whole invocation, out-of-process, at 1 Hz.
  3. SCRAPE ``/v1/admin/flag_truth`` at both boundaries and diff it. Reading
     them over HTTP outside the timed window costs the engine nothing. The diff
     is the CODE-PATH truth: it says whether the engine took the same route
     through itself on a slow invocation as on a fast one.
  4. AUDIT the window afterwards and FAIL LOUD if the machine degraded mid-run.
  5. STAMP everything onto a JSON envelope beside the bench's own output.

WHY INVOCATION GRANULARITY
--------------------------
Because that is the anomaly's granularity. The slow mode struck a whole bench
invocation -- both of its concurrencies -- while a sibling invocation against the
SAME STANDING SERVER was clean. So the unit that must carry evidence is the
invocation, and the fields that must be captured are the ones that can differ
between two invocations of one server: host contention, clocks, link state,
thermals, and the engine's own path counters.

THE DISCRIMINATOR
-----------------
``env_gpu_busy_frac_nvml`` is the field to read first when diffing a fast against
a slow invocation:

  * host contention starves the engine's launch thread -> the GPU waits ->
    busy fraction FALLS while wall time rises;
  * a clock/link/thermal degradation slows the GPU's own work -> busy fraction
    HOLDS or RISES while wall time rises.

Two different diseases, opposite signs, one field. Then read the flag-truth diff:
if the captured-vs-eager step mix moved, it was never an environment problem at
all and the code took a different route.

Per-counter fired/refused delta across the invocation.

This is the field that says "the environment changed" vs "the code changed".
If a slow invocation has the same step count but a different captured-vs-eager
mix, it is not an environment problem and no amount of nvidia-smi will find it.

Hash the arbi_serve package so a row can prove WHICH CODE produced it.

The rig bind-mounts the worktree into a container. If the mount silently does
not take, the container runs the image's baked-in code and the A/B compares a
branch against itself. Hashing the tree the harness *believes* it is testing
is half the check; ``--container`` hashes what the server actually imported.

Summary statistics, aggregation, and text rendering for ``arbi-serve bench``.

Holds the latency :class:`_Distribution` helper, the structured-summary
builder :func:`_collect_bench_stats`, and the human-readable
:func:`_print_text_summary` / :func:`_to_samples` renderers. The public
names remain importable from :mod:`arbi_serve.bench`.

Compute the structured bench summary from the collected payloads.

Applies ``--warmup`` discard to timing-bearing completions only,
builds the TTFT / TPOT / queue-wait distributions, aggregate
throughput, per-label first-call costs, and MTP accept stats, then
returns the summary dict (the ``--json`` form).

Pull the raw sample list out of a dist dict back from raw_timings.

The dist dict only carries {n, mean, p50, p90, p99, max}; for
formatting we reconstruct the underlying list from raw_timings,
walking the ``components`` block.

Prompt/config construction and the engine driver for ``arbi-serve bench``.

Holds the prompt builders, MTP / parallel / server config assembly, and
the async :func:`_drive_bench` harness that boots an in-process engine
and fires the synthetic request loop. The public names remain
importable from :mod:`arbi_serve.bench`.

Build the per-request synthetic prompt token list.

Three modes:

  * ``--prompt-token-id N`` — every position is ``N`` (operator-
    forced; use this to reproduce a specific failure shape).
  * ``--seed N`` set, ``--prompt-token-id`` unset — deterministic
    random ids drawn from the safe mid-vocab range
    ``[256, vocab_size - 256)``. Avoids low-id control tokens
    (BOS/EOS/PAD live there in many tokenizers) and high-id
    special tokens (chat-template markers).
  * Neither set — legacy ``[1] * --prompt-tokens`` shape.

The legacy shape is the documented default in tests/test_timing_e2e
and the operator runbook; rotating it would silently break
cross-PR comparison runs. New behaviour is opt-in.

Read a bundled prompt fixture from ``arbi_serve/bench_prompts/``.

Uses :mod:`importlib.resources` so the read works regardless of
whether arbi-serve is installed as a wheel or run from the source
tree. Raises :class:`FileNotFoundError` with a clear message when
the fixture is missing — that's a packaging bug, not an operator
error, so the failure should be loud.

Repeat-then-truncate ``ids`` to exactly ``target`` length.

Single-token fixtures degenerate-but-don't-crash (the doubling loop
still terminates when len(ids) >= 1). Empty ``ids`` is a tokenizer
bug — surface it loudly rather than infinite-looping.

Build the per-request prompt token list.

Resolution order:

  1. ``--prompt-text TEXT`` — tokenize the literal string.
  2. ``--prompt-file PATH`` — tokenize file contents.
  3. ``--prompt-preset {wiki,code,chat,instruct}`` — tokenize the
     bundled fixture.
  4. ``--prompt-preset random`` (default) — fall through to the
     legacy synthetic random-token path :func:`_build_synthetic_prompt`.

All real-text paths repeat-then-truncate to ``--prompt-tokens``
so MTP accept-rate measurements use a fixed prefill shape
independent of the source text length.

Resolve the TP parallel config from ``--tp-size`` and the env.

world_size>1 requires a torchrun launch so RANK / WORLD_SIZE /
LOCAL_RANK are populated; the bench then builds a
DistributedEngineDriver and only rank 0 runs the request loop.
Raises :class:`SystemExit` on an inconsistent launch.

Build the per-request prompt ids and optional ``--prompt-set`` list.

Returns ``(prompt_ids, prompt_set_ids)``. ``prompt_set_ids`` is None
unless ``--prompt-set`` was given, in which case each line is wrapped
in the model's chat template and ``prompt_ids`` is its first entry.
Falls through to :func:`_build_prompt_ids` for the single-prompt path.

Drive ``one(idx)`` for every request, sequentially or fanned out.

Sequential (concurrency 1) is the default because TTFT cleanliness
is the canonical use case; otherwise fan out ``concurrency`` at a
time, drain, repeat (simple semantics — not a sliding window).

Read and print the MtpDriver acc-len / overall-accept counters.

Returns ``(acc_len, accept_overall)``; both NaN when no drafts were
issued. Counters live on ``eng.mtp_driver`` (the source of truth),
not the thin ``eng.spec_decode`` strategy — the per-request handle
counters can read 0 on this path.

Dump the verify-forward captured-graph lookup histogram and keys.

Reveals why (and how often) the MTP verify forward misses the
captured graph at concurrency. Gated by the runner's debug attrs
(``ARBI_DEBUG_CAPTURE_LOOKUP``) so steady runs pay nothing.

Live boot progress: which phase is open, how long the boot has run.

The startup-phase METRIC records a phase once it closes, so a phase that
is still running is invisible to it — and the first phase does not close
for most of a cold boot. This registry carries the other half: the
currently-open phase and its elapsed seconds, published from the moment
the process starts so a probe or a scrape during the boot reports
something true instead of nothing.

Process-global (one boot per process) and torch-free — it sits at the
package root, not under ``runtime``, so the health probes and the serving
gate can read it without pulling the engine/torch chain. The HTTP thread
reads it while the boot thread writes it. Boot-only bookkeeping — a
handful of dict writes per phase, never touched on the serving path.

At TP>1 only rank 0 serves HTTP, so only rank 0's registry is read. That
is enough: every rank walks the SAME collective build, so rank 0's open
phase is the fleet's phase — and a phase whose elapsed keeps climbing on
rank 0 means rank 0 is blocked in a collective waiting for a peer that
has not arrived. The topology (:func:`set_topology`) rides in the
snapshot so a reader knows it is looking at rank 0 of N.

Two kinds of window are published, because the boot has two kinds:

* the BUILD phases, which the boot procedure opens and closes in
  sequence (:func:`phase_opened` / :func:`phase_closed`);
* NESTED steps (:func:`nested_phase_opened` / :func:`nested_phase_closed`)
  — work that is triggered lazily from inside a dependency and can land
  anywhere, in a build phase or between two of them. A JIT cpp/CUDA
  extension build is the case that forced this: it runs for minutes, it
  is reached from a model import, and it has no place in the build's
  phase sequence. It gets its own window rather than borrowing the
  enclosing phase's name, and closing it RESTORES the enclosing phase
  instead of clearing it.

:func:`active_phase` reports the innermost open window — the most
specific true answer to "what is this process doing right now".

:func:`phase_timeline` is the same history with its STRUCTURE kept: one
record per window that closed, carrying where it started, how long it
ran, and which window contained it. Without the containment a nested
step renders as the next item in a sequence and a reader summing the
list double-counts it — a JIT extension build's seconds are ALSO inside
the phase that triggered it. With it, a nested step draws inside its
parent, and two records whose intervals genuinely overlap (an async DMA
still running under the next phase) are distinguishable from that by
their intervals rather than by guesswork.

Nothing here guesses. When no window is open, the snapshot says which
phase closed last and how long ago, and :func:`describe_position` renders
exactly that; it never names a cause it did not observe.

States: ``not_started`` → ``booting`` → ``ready``, or ``failed``.

One closed boot window, with the structure a timeline needs.

``start_s`` is the window's OPEN when one was published, and
``closed_at - seconds`` otherwise. The fallback is what the boot's own
telescoping phases give — several close without ever announcing an
open, and their duration is measured from the previous close, so
subtracting it lands on the true start. Where a phase charges seconds
off onto another (``_phase_done(charge=…)``), the bar sits where the
REPORTED duration puts it, so the drawing and the phase accounting
cannot tell different stories.

``parent`` is the window that was open when this one opened. With two
genuinely concurrent steps (extension builds on the boot-overlap
worker and the boot thread at once) that is the innermost window at
that instant, which is an observation rather than a derivation; the
intervals are what settle whether two records overlap.

Record which rank this process is, and how many there are.

Reported in the snapshot so a reader can tell a single-rank boot from
rank 0 of a collective one — the difference decides how to read a
phase that has stopped advancing.

Record ``name``'s measured duration and clear the open phase.

Phases accumulate: a phase the boot charges more than once (a hybrid
model's second compile warmup) reports the sum, matching the metric.

Open a nested step INSIDE whatever is already running.

For work reached lazily from a dependency — a JIT extension build —
which can land in any build phase or between two of them. The
enclosing phase keeps running and is restored on close.

The containment is stamped HERE, because the instant of opening is
the only moment it is observable: by the time the step closes, the
stack has moved on.

Close the nested step ``name`` and record its measured duration.

Removes the innermost entry with that name, so concurrent steps
(extension builds on the boot-overlap worker and the boot thread at
once) each close their own window rather than the top of the stack.

``did_work`` / ``detail`` are what the STEP measured about itself —
whether the bracketed work actually ran, and in what words. Only the
step can know that, so it is passed in rather than inferred here from
a duration, which would be a threshold pretending to be a fact.

Monotonic zero for ``start_s``. Caller holds ``_LOCK``.

The boot clock when one has started, the process start otherwise — so
a window that closed before :func:`boot_started` reports a negative
offset rather than being clamped into a position it never occupied.

``(name, seconds_open)`` for the innermost open window, or ``None``.

A nested step wins over the build phase containing it: it is the more
specific answer, and it is the one that is actually consuming the
seconds a reader is watching climb.

Closed windows in the order they closed, name → seconds.

NOT a partition of the boot. The build phases telescope, but a nested
step's seconds are also inside the phase that contained it — that is
what nested means, and this flat mapping cannot say so.
``arbi_serve_startup_phase_seconds`` is the surface that sums, and it
records build phases only; :func:`phase_timeline` is the one that
keeps the containment.

Every closed window as a :class:`PhaseRecord`, earliest start first.

One record per OCCURRENCE (a phase the boot charges twice draws two
bars), each with its start offset, duration, depth and parent — the
four things a timeline needs and the flat ``name -> seconds`` mapping
cannot express. Rendering these as a flat sequence double-counts every
nested step; rendering a nested record inside its parent does not.

``(name, seconds, age_s)`` for the phase that closed most recently.

``age_s`` is how long ago it closed — with no window open that is the
length of the stretch nothing has named, which is the number a reader
staring at a silent boot actually wants.

One line saying where the boot is — observed only, never inferred.

Renders from a snapshot so the prose and the JSON beside it cannot
disagree. Three cases, and each says exactly as much as the registry
knows:

* a window is open → name it and how long it has been open;
* none open but phases have closed → name the last one and the size
  of the silence since. NOT a cause: something is running that no
  bracket covers, and this line reports the gap rather than
  inventing what is in it;
* nothing has closed either → the boot has not reached its first
  phase. That is a fact about the registry, not a claim about the
  work, and it is the only case in which the pre-phase window may be
  named at all.

The fold-state emit: where a mid-step savepoint's state lands, and the step plan.

A savepoint boundary sits INSIDE a prefill step (arbicity/arbi-serve#2238):
the state after exactly ``boundary`` prompt tokens is swept through by one
kernel and held nowhere the commit path can observe. The recurrent kernel
writes it as it passes (:mod:`arbi_serve.kernels.fla_vendored.chunk_delta_h_emit`),
and this module owns everything that write needs which the forward itself
must not carry.

THE LAUNCH IS THE SAME WHETHER OR NOT A STEP ARMS
=================================================
Every served GDN prefill runs as ONE Inductor graph per decoder layer, and a
full prefill capture replays a CUDA graph of the whole forward. Anything the
forward reads to decide "emit here" — a row, an offset, a plan object — is a
Python value Dynamo guards on and a captured graph bakes. So nothing is read.
Every prefill launch carries the emit table (``-1`` everywhere when nothing
is armed) and gathers its conv window at a device index table; the tables
are persistent buffers the metadata builder writes host-side before the
forward, one H2D each. The traced graph, and the captured graph, are the
same with and without a plan: a step that arms differs from one that does
not only in what two small buffers hold.

ONE KERNEL, so the emitting launch and the plain launch are one Triton
function with one autotune: the state a savepoint stores is produced by the
very launch a step without a savepoint would issue, and its bit-identity
against the uncut fold (``tools/gdn_row_class/check_emit_state.py``) is a
property of the kernel, not of which config two kernels happened to pick.

NO VRAM FOR RETENTION, AND A CORE OTHER RECURRENT KINDS PLUG INTO
=================================================================
Retention lives in the pinned-host ring (and the store's disk tier beneath
it): a savepoint costs the card nothing, and the retained count scales with
host RAM, never with KV. The device cost of the whole mechanism is the one
staging slot below. An engine that caches recurrent state in DEVICE slots
pays a sequence's whole state per retained entry out of servable context;
this one pays a PCIe copy per snapshot instead, off the compute stream.

The staging, the plan, the publish, the spill and the capture landing are
kind-agnostic: geometry is read off the pool's own views, and the ops hand
in ``(emit table, slot)`` and gather a conv window at an index table. What
a recurrent KIND supplies is (a) a prefill kernel that can store its fp32
running state at a chunk index into ``emit_out`` and (b) which stream its
conv window is a slice of. GDN supplies both here. ShortConv's state IS the
conv window, so it needs only (b). Mamba2's boundary state is the SSD
scan's state-passing output, which upstream materialises per chunk only as
the next kernel's matmul operand (``out_dtype=C.dtype``); it supplies (a) as
one store of the scan's own fp32 register
(:mod:`arbi_serve.kernels.mamba_vendored.ssd_state_passing_emit`), the same
shape as the GDN emit.

WHAT AN EMIT COSTS THE STATE
============================
Nothing, and that is the property to hold onto rather than "it is fp32". The
emit is the fp32 register the kernel already carries, rounded once on store
into the staging slot, which is the SLAB's dtype whatever that is — GDN's
recurrent slab is fp32 under ``ARBI_GDN_RECURRENT_FP32``, Mamba's SSM slab is
the activation dtype (``_state_view_registry`` promotes GDN's alone). The
slab scatter rounds ``final_states`` once, the same way, so a savepoint holds
exactly what the slab would hold had the prefill been cut at that boundary:
not a nearby state, and not one carrying a rounding the model's own path does
not. The per-kind checkers measure that equality
(``tools/gdn_row_class/check_emit_state.py``,
``tools/mamba_fold_emit/check_emit_state.py``).

ONE DEVICE SLOT, ON THE LEDGER
==============================
The emit lands in a slot holding one recurrent row in the slab's own layout
and one conv window, sized off the pool's geometry. Layer ``L+1`` writes it
after layer ``L``'s copy has read it, and that wait is MEASURED
(``arbi_serve.savepoint.fold_emit_slot_wait_ms``, device time the compute
stream sat on the fence) rather than assumed away: a layer's out-projection
and MLP separate the copy from the reuse, and the copy is a few MiB. The
slot count is a parameter (``n_slots``) so that histogram, not an analogy
with another engine's ping-pong buffer, decides whether a second slot is
worth its bytes. A slot per layer would be a whole ring slot of VRAM, which
the boot ledger prices at tens of KV pages; a slot per row would be
``max_batch`` copies of a state only one row ever writes.

THE COPY IS OFF THE COMPUTE STREAM, AND CAPTURABLE
==================================================
A layer's slot is copied to host on this module's own stream, forked from
the compute stream behind the emitting kernel and joined where the slot is
next written. On an eager step the copy goes straight into the ring slot the
scheduler took for this snapshot. A captured graph cannot bake a per-snapshot
ring address, so inside a capture the copies go to a fixed pinned LANDING
zone — one ring slot's worth of host RAM, booked on the pinned-host ledger —
and after a replay the landing zone is moved into the ring slot by a host
copy on a worker thread, fenced for the store's consumers exactly as a device
event would be. The capture's side-stream work is joined before the capture
ends (:meth:`FoldEmitStaging.join_capture`), and the events it recorded live
as long as the graph does.

A fence over a host-side copy, duck-typed like a ``torch.cuda.Event``.

:class:`~arbi_serve.cache.recurrent_savepoint.SavepointEntry` carries one
``pending_event`` and calls ``synchronize()`` on it before the host bytes
are read or freed; a replayed captured step's bytes arrive by a worker
thread's copy, so its fence is this.

The emit's device slots, the two step tables, and the copy machinery.

Constructed once per pool geometry at boot
(:meth:`RecurrentLifecycleMixin.attach_fold_emit_staging`) and kept for
the process: a captured graph bakes the slots' and tables' addresses, so
a pool rebuilt after capture must find the same objects. ``landing`` is
allocated only when a capture mode can replay a prefill; an eager-only
boot pays no host RAM for it.

This step's fold-split plan, resolved to what the ops check and use.

Published by :func:`publish_step_plan` and read by the recurrent op's
real-impl through :func:`fold_emit_step_for`. ``row`` and ``cu_host``
are LAUNCH-relative: on a row-class-mixed step the chunk kernel sees the
prefill class alone, rebased to start at 0, and the plan is keyed on
that tensor.

Turn the scheduler's ``(req_id, offset, boundary)`` into the step's tables and plan.

The ONE derivation for every path a prefill can take — the metadata
builder (eager and compiled forwards) and the captured-prefill replay,
which builds no metadata — so a step is armed for the ops in exactly one
way. ``refuse`` names why not on the caller's counter. Returns the plan,
published and armed, or ``None`` with the tables unarmed.

The registered layer block that supplies the emitting chunk entry, or ``None``.

Sharing the GDN slab is not the same as supplying the emit, and only the
dispatch registry can tell the two apart: ``attach_fold_emit_staging``
sizes the slot off the pool's GDN views, so a kind that borrows
:class:`~arbi_serve.cache.recurrent_views.GdnLayerView` and registers
under the ``gdn_attention`` op gets a slot sized for it. KDA
(``bailing_hybrid``) is such a kind: its chunk entry is upstream FLA's,
which has no emit parameter, so it carries no
:meth:`_fold_emit_delta_rule`.

``None`` means this boot's recurrent kind cannot emit.

Launch the emitting chunk kernel once at boot, at the served geometry.

The kernel every prefill launch takes is warmed by the boot's own prefill
forwards once the staging exists; forwards before it (the KV profile)
take the plain upstream kernel. Cold, the first served launch runs the
Triton autotune on a live request after the KV grow, whose benchmark
allocates an L2-flush buffer larger than the headroom the grow leaves:
measured as a step OOM on the first armed step of a boot
(arbicity/arbi-serve#2238). One launch here, while the card is still
free, compiles and tunes it exactly as serving will call it.

Page-lock the capture landing zone, laid out exactly as a ring slot.

``layout`` is :func:`~arbi_serve.cache._recurrent_lifecycle._snapshot_arena_layout`'s
rows over the model's snapshot walk, so the landing zone and a ring
slot are the same bytes in the same order and the post-replay move is
one flat copy. Booked on the pinned-host ledger BEFORE it is locked.

Write the step's tables host-side and hand them to the device.

``emit_chunk[row]`` names the chunk at whose start the state is
emitted; ``conv_idx`` names the ``kernel`` flat tokens before the
boundary ``t`` — the conv state after ``t`` tokens is that window
(``gdn_prefill_conv_flat`` defines its own ``final_buf`` as exactly
the last ``kernel`` inputs). Two small async H2D copies on the current
stream, ahead of the forward; nothing enters a trace.

``(emit_chunk view, emit_out)`` for a launch of the ONE batch row ``row``.

A kernel whose prefill is issued per row rather than varlen
(:meth:`~arbi_serve.models.mamba2_block.Mamba2Block._forward_kernel_prefill`)
sees a batch of one, so its emit table is this row's single entry of
the step's table — the same buffer the metadata builder armed, sliced,
never a value read out of it.

Gather this launch's conv window into the layer's slot.

Unconditional and index-driven: the same kernels whether or not a
plan is armed, so the compiled and the captured graph are the same
either way. The slot is first fenced behind the copy that last read
it.

Put an already-gathered ``(K, C)`` conv window in the layer's slot.

The window a kind's conv state IS differs — GDN's slab row is the last
``kernel`` inputs, Mamba's rolling buffer keeps slot 0 out of the carry
— so the gather belongs to the kind and the slot's fencing belongs
here.

Copy the layer's slot to host on the staging stream; return the fence.

Eager: straight into ``snapshot``'s ring-slot views. Capturing:
into the fixed landing zone (``snapshot`` may be ``None``; a captured
graph copies every layer on every replay, and the post-replay move
decides whether the bytes go anywhere). The fence is what the next
writer of this slot waits on, and — eager — what the snapshot hands
the store as the entry's pending event.

Order the compute stream behind the last copy that read ``s``.

Eager: bracketed by a timing pair on the compute stream, read lazily,
so ``fold_emit_slot_wait_ms`` is the device time the stream actually
sat on the fence — zero when the copy had retired. Timing events
cannot be recorded inside a capture, so a captured wait is a bare
dependency the graph replays.

Join the staging stream's captured copies into ``stream``; verify coverage.

Called by the prefill capture inside its ``torch.cuda.graph`` region,
after the forward: a capture with unjoined side-stream work is
refused by CUDA, and a captured graph that copied fewer layers than
the model stages would replay a partial snapshot — refused HERE, at
boot, rather than at every commit.

After a replayed prefill: move the landing zone into the step's ring slot.

No plan, nothing to move — the graph copied into the landing zone
and the bytes are simply overwritten next time. With a plan, the move
is a host copy on a worker thread that first waits for the replay's
device event, and the snapshot's fence is that copy. Copies are
serialised through one worker so a slot's bytes are never written by
two moves at once; a ring slot is reused only after every other slot
has been handed out, which is many steps later than one move takes.

Does the plan's host copy of the boundaries describe a launch of ``n_tokens``?

The boundary tensor is a persistent buffer: its address and length
repeat across steps, so identity says "the same buffer", not "the
same contents". The launch's flat token count is the one fact about
its contents that is free to read (a shape).

GrowableRegion-backed PAGED_KV grow / sleep / wake lifecycle.

The :class:`GrowableKvMixin` wraps the primary ``PAGED_KV`` pool
view's growable-slab surface: query mapped footprint, grow to a
page target or to fill free VRAM, and sleep / wake / close the physical
mapping at a stable VA. Every wrapper funnels through
:meth:`_primary_paged_delegate` so the getattr-then-forward shape lives
in one place; the pool facade composes this mixin.

Forward ``name`` to the primary ``PAGED_KV`` pool view.

Looks up the PAGED_KV view, fetches its ``name`` attribute, and —
if that attribute is callable — invokes it with ``*args`` /
``**kwargs`` and returns the result; otherwise returns the
attribute value directly. Returns ``default`` when the view is
absent or does not define ``name`` (non-growable / non-PagedKV
pools). Backs the thin growable-KV wrappers below so the
``getattr``-then-forward shape lives in one place.

Physical bytes currently mapped across the GrowableRegion-backed
PAGED_KV slab(s); 0 when the KV pool is not growable. Surfaces the
growable KV slab's mapped footprint to ``kv_pool`` metrics + the B3
freeze invariant (the slab lives in its OWN VA reservation, NOT under
the ``kv_pool`` NamedMemPool tag, so it isn't counted by the
pluggable-allocator tag accounting).

Physical bytes a grow to ``target_pages`` leaves mapped, or ``None``
when the KV pool cannot answer (not growable / not GrowableRegion-backed).

THE seam every sizing decision that divides up a driver free reading
prices its target with. See
:meth:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool.growable_mapped_bytes_for_pages`
for why this is not ``pages × per_page_bytes``. ``None`` — never a
neutral 0 — so a caller that cannot get the real cost falls back to the
capacity arithmetic knowingly rather than sizing against a zero.

Grow the growable PAGED_KV pool to FILL the VRAM currently free,
keeping ``floor_bytes`` unmapped. Returns the new mapped-pages count.

For REBUILD paths (backend swap / sleep-wake / reload) that rebuild
the pool WITHOUT a cudagraph capture sweep. Unlike the cold-boot
prefix-then-grow (which leaves the capture transient room in the
unmapped tail), a rebuild does no capture — so the pool must end
FULLY sized, else its free list stays at the tiny capture prefix and
admission immediately rejects every request with
``kv_watermark_exceeded``.

``floor_bytes`` is the VRAM this grow must LEAVE FREE. It is not an
allocator margin: one in-flight serving step allocates real transients
(the forward arena's regrow, the speculative verify tail, the drafter's
gather), and a grow that does not hold them back leaves the card with
nothing for the first request. Callers pass
:func:`~arbi_serve.engine.inprocess_capture.serving_floor_for_grow` —
the same reserve the post-capture boot grow honours.

``max_pages`` caps the target for a caller RESTORING a known page count
rather than filling the card.

Per-page bytes are derived from the pool's OWN mapped footprint
(``growable_mapped_bytes_total / mapped_pages``) rather than the
engine's cached estimate, so a cross-backend swap (tkv↔bf16, which
changes bytes/page) sizes correctly. That average is what the pool can
STORE per page; what a target COSTS to map is
:meth:`kv_mapped_bytes_for_pages`, and the two differ by the
granularity round-up each per-layer region applies. So the average
proposes and the price list disposes: a proposal whose real map cost
exceeds the headroom is walked down until it fits, which is the only
direction that cannot leave the floor overdrawn. ``grow_kv_to_pages``
caps the target at the slab's reserved ``num_pages``. No-op (returns
current mapped pages) when the KV pool is not growable.

Wake (remap physical at the same VA) the GrowableRegion-backed
PAGED_KV slab(s). Mirror of :meth:`sleep_growable_kv`. No-op (0) when
not growable. ``discarded=True`` re-provisions a discard-parked
member's slab from scratch (``budget_bytes`` sizes it) — see
:meth:`arbi_serve.cache.paged_kv_pool.PagedKVStatePool.wake_growable_kv`.

Shared LRU + byte-budget host-RAM cache primitives.

Two host-tier KV stores — :class:`SuspendedJobStore` (parked batch
jobs) and :class:`PrefixTierStore` (demoted prefix pages) — keep an
LRU-ordered host-RAM map bounded by a byte budget, spill cold entries
to an optional disk tier (LRU + TTL bounded), and fault them back in.
``RecurrentSavepointStore`` shares the same host-tensor entry contract
(byte accounting + lazy-confirm D2H barrier). This module hoists the
duplicated pieces:

  * :class:`_LruByteBudgetEntry` — the per-entry host-tensor byte
    accounting (``compute_nbytes``) and async-D2H ``synchronize``
    barrier shared by every store's entry dataclass.
  * :class:`_LruByteBudgetStore` — a host-RAM LRU byte-budget store
    with an integrated optional disk-spill tier. The host map evicts
    cold entries to disk (or drops them when no disk tier is set) until
    the host byte budget fits; the disk tier evicts by LRU and a TTL
    sweep. Per-store specifics (the key/entry types, blob encode/decode,
    disk index format, spill gating) live in subclass hooks.

The entry contract
==================
An entry carries ``{... : host_tensor}`` payloads. The store never
inspects tensor contents — only ``element_size() * numel()`` for byte
accounting. When the snapshot was issued async (non-blocking D2H), the
recorded CUDA event in ``pending_event`` must fire before the host
tensors are read or freed; :meth:`_LruByteBudgetEntry.synchronize`
waits on it once, then clears it.

Host-tensor byte accounting + async-D2H barrier for store entries.

Subclasses are dataclasses that declare the ``nbytes: int``,
``_nbytes_cached: bool`` and ``pending_event`` fields and implement
:meth:`_byte_sources` to yield the host tensors whose byte cost
constitutes the entry. This base supplies the shared
:meth:`compute_nbytes` (O(1) cached) and :meth:`synchronize`.

Host-RAM LRU byte-budget store with an optional disk-spill tier.

Host tier: ``{key: entry}`` LRU-ordered (oldest = left), bounded by
``host_max_bytes``. Cold host entries spill to disk (or drop, when no
disk tier is configured). Disk tier: ``{key: rec}`` LRU-ordered,
bounded by ``disk_max_bytes`` AND a ``ttl_seconds`` age sweep.

``stage`` admits to host; ``fault_in`` returns an entry (host
fast-path or disk read, promoting back to host); ``drop`` removes
from both tiers; ``sweep_ttl`` drops expired disk blobs.

Per-store specifics live in subclass hooks: :meth:`_init_offload_dir`
(dir setup + index reload / crypto), :meth:`_should_spill_to_disk`
(spill gate), :meth:`_spill_write` (encode a host entry to a disk
blob + build its index record), :meth:`_unlink_blob` (remove a blob),
:meth:`_read_blob` (decode a blob back to an entry),
:meth:`_write_index_locked` (persist the disk index), and the
``_BLOB_READ_EXCEPTIONS`` / :meth:`_warn_blob_read_failed` /
:meth:`_log_stage` observability seams.

Thread-safe under a single re-entrant lock; blob I/O happens outside
the lock in subclass hooks (only bookkeeping is under it).

Sum host-tensor byte cost; cache O(1). Safe under pending D2H.

Reads only tensor metadata (``element_size`` / ``numel``), never
the underlying host bytes, so it is safe to call before a
pending async D2H event has fired.

Wait for the pending async D2H event, if any (idempotent).

Cheap no-op when ``pending_event`` is ``None`` (synchronous
snapshot path or a CPU-only test pool). When set, blocks the
host until the recorded event retires — at which point the host
tensors carry the snapshot and are safe to read or release — then
clears ``pending_event`` so subsequent calls fast-path.

Admit ``entry`` into the HOST tier under ``key``.

Evicts host-LRU to disk (or drops, when no disk tier) until the
post-admit host footprint fits ``host_max_bytes``. If the single
entry exceeds the host budget it is staged then immediately
spilled — disk is the floor. Re-staging the same key replaces the
prior entry (refresh recency + bytes) and drops any stale disk
blob for that key.

Return the entry for ``key`` (host fast-path or disk read).

Host hit: bump recency and return the entry (synchronised against
any pending async D2H). Disk hit: read the blob back into host
tensors, promote to host (subject to the host budget), and return
it. ``None`` on miss (including an expired or unreadable blob).

Spill host-LRU entries to disk until the host budget fits.

``protect`` is the just-touched key; it is not spilled first (it
is the MRU and the caller wants it resident). Caller holds lock.

Move one host entry → disk (or DROP when no disk tier).

When no disk tier is configured, or the subclass spill gate
declines (:meth:`_should_spill_to_disk`), or the blob write fails
(:meth:`_spill_write` returns ``None``), the host entry is simply
dropped — host overflow is a hard drop in those cases. Otherwise
the entry is written to a blob and recorded in the disk index.
Caller holds lock.

Create the offload dir + restore/init its disk-tier state.

Called from ``__init__`` when ``offload_dir`` is set. Subclasses
make the directory, reload any persisted index, and set up any
per-process state (e.g. blob encryption).

Encode ``entry`` to a disk blob and return its index record.

Returns the ``rec`` dict to store under ``key`` in the disk
index, or ``None`` when the blob could not be written (the entry
is then dropped). Caller holds lock; the budget has already been
enforced for ``entry.nbytes``.

Metrics shim used by the prefix-sharing surface.

Two emit channels in lockstep:

  1. An in-process registry (counters / histograms keyed by
     name + label dict) that backs ``snapshot()`` for tests and the
     admin endpoint. This works without OTEL configured, which is why
     the unit tests in :mod:`tests.test_radix_pagetable` can call into
     it directly with a fake pool.
  2. A bound :class:`arbi_serve.server.metrics.Metrics` instance, when
     present, into which the same emits are forwarded so OTEL counters
     light up under live engines. The bind is a one-shot call from
     :class:`arbi_serve.engine.engine.Engine.__init__`; tests / scripts
     that never construct an :class:`Engine` get the in-process channel
     only and don't pay the OTEL import cost.

Metric names follow this contract:

  - ``arbi_serve.prefix_cache.hit_total{kind}``
  - ``arbi_serve.prefix_cache.tokens_saved{kind}``
  - ``arbi_serve.prefix_cache.eviction_total{reason}``
  - ``arbi_serve.batch.in_batch_shared_tokens``  (histogram)

The OTEL forwarder uses the canonical
``arbi_serve.prefix_cache.tokens_saved_total`` name; the in-process
registry keeps the shorter ``tokens_saved`` for backward compatibility
with existing radix-cache tests that snapshot it directly.

Paged-pool allocation, accounting, page-table factory, and draft slots.

The :class:`PagedAllocationMixin` groups the engine/scheduler-facing
page-allocation surface: byte accounting across pools, the page-table
factory + tenant-namespaced prefix lookup, page alloc/free, suspended-
job page snapshot/restore, and the speculative draft-slot lifecycle.
Every method routes through the shared ``self._views`` registry and the
primary-paged-kind resolution the pool facade also composes.

Free pages for ``kind`` (default PAGED_KV).

MLA shared-KV (``StateKind.MLA_SHARED``) is also paged. Mamba /
GDN return 0 — recurrent state is per-request not per-page.

For MLA-only models (DeepSeek V2/V3, Kimi K2) the engine /
scheduler keep asking for ``PAGED_KV`` (the default argument);
transparently route those queries to the primary paged kind so
the scheduler does the right thing without per-arch branching.

Construct the page-table impl registered for ``kind``.

Resolution:
  1. If ``kind`` is ``None``, route to :meth:`primary_paged_kind`
     (PAGED_KV → MLA_SHARED) so MLA-only models get the right
     slab without caller-side conditionals.
  2. If a factory was registered for ``kind`` in
     ``page_table_factories``, call it with ``self``.
  3. Otherwise default to :class:`FlatPageTable` (no prefix
     cache).

Each call constructs a fresh page table — the engine is the
single owner and binds the result via :meth:`bind_page_table`
once at startup. Re-calling for the same kind in tests is safe
but unusual.

Tenant-namespaced :meth:`PageTable.lookup_prefix` convenience.

Engine call sites that hold a :class:`Request` or a
:class:`TenantContext` route through here so the radix-tree
namespace is always derived from
:attr:`TenantContext.tenant_id` (not e.g. the bearer token,
not :attr:`Request.tenant_id`). Anonymous traffic
uses the empty string ⇒ global namespace ⇒ cross-request
sharing semantics.

Returns the matched page IDs (longest page-aligned prefix in
the tenant's radix tree) or ``None`` on miss / on the flat
impl. The bound page table must satisfy the
:class:`PageTable` Protocol. When no page table is
bound, returns ``None`` — the same outcome as a flat impl,
which is what tests expect.

Resolution order for the tenant id:
  1. ``tenant.tenant_id`` if ``tenant`` is given.
  2. ``request.tenant.tenant_id`` if ``request`` is given AND
     carries a non-anon tenant.
  3. legacy ``request.tenant_id`` field for source-compat.
  4. ``""`` (global namespace).

Allocate ``n`` pages from the named kind's pool.

Returns the page ids or ``None`` if insufficient. ``kind=None``
routes to the primary paged kind (PAGED_KV for paged-KV models,
MLA_SHARED for MLA models). Recurrent pools allocate
per-request, not per-page, and raise.

Gather ``page_ids``' KV content into host tensors.

Returns ``{tag: host_tensor}`` (see
:meth:`PagedKVStatePool.snapshot_pages`). When ``async_copy`` is
True, also call :meth:`snapshot_pages_event` immediately after to
fence the in-flight D2H before reading the host buffers.

Record + return a CUDA event after an async page snapshot.

Mirrors :meth:`snapshot_recurrent_event`: the
:class:`SuspendedJobStore` waits on this event before reading the
pinned-host gather buffers. Returns ``None`` when CUDA is
unavailable (CPU test pools) — the snapshot was synchronous then.

Register the engine's page table for draft-slot bookkeeping.

Called by :func:`Engine.build` after both the pool and the
page table have been constructed. Decoupling construction
from binding lets tests build a pool without a page table and
skip the MTP path.

Allocate ``k`` draft token slots tail-extending ``req_id``.

Returns the per-token slot indices (page_id * page_size +
offset) for the kernel's slot_mapping, or ``None`` if the
underlying paged pool is exhausted. Pool-exhaustion is not an
engine-fatal error — the caller falls back to the K=1 path
for that step.

Composes with :meth:`free_pages_list` semantics: every new
page taken during this call starts in the same allocated
state as a normal page. On rollback the pages are immediately
returned via the same ``free_pages_list`` path.

Copy whole KV token slots within the paged slab.

The pool half of a TREE accept's compaction: a tree's accepted
path is scattered across the step's draft slots, and
:meth:`free_draft_slots` trims a SUFFIX — so the accepted rows
must be moved down into the kept prefix before the trim, or the
trim keeps the wrong tokens' KV with nothing raising.

Thin pass-through, matching the rest of this section: the page
table owns which slots move (:meth:`FlatPageTable.draft_path_moves`)
and the paged view owns the bytes.

Trim rejected draft slots after the verify pass.

Accepted draft slots stay (they become first-class KV the
next step's attention can read as already-cached prefix).
Trailing pages that become fully empty are returned to the
pool so steady-state high-K speculation does not pin idle
pages.

``k_accepted + k_rejected`` must equal the ``k`` from the
prior :meth:`allocate_draft_slots` call.

Shared allocator bases for the cache state pools.

Two small mixins hold the bookkeeping shared across the pool classes:

  * :class:`PerRequestRowPool` — the per-request slab-row allocator
    shared by :class:`~arbi_serve.cache.recurrent_pool.RecurrentStatePool`
    and :class:`~arbi_serve.cache.short_conv_pool.ShortConvStatePool`.
  * :class:`PagedFreeListPool` — the free-page allocator shared by
    :class:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool` and
    :class:`~arbi_serve.cache.mla_pool.MLAStatePool`.

Neither base defines ``__init__`` — the subclasses keep their own
(very different) construction signatures and call the small
``_init_*`` helper once they have set the attributes the base reads.

Set ``rows`` of ``buf`` (along dim 0) to ``fill``, with NO device sync.

The obvious spelling, ``buf[rows] = fill``, costs a blocking
device-to-host sync per call. Measured, not assumed: with
``torch.cuda.set_sync_debug_mode("warn")``, ``buf[rows] = 0`` reports 1-2
implicit syncs while ``buf.index_fill_(0, rows, 0)`` reports none, and the
index's device is irrelevant -- ``buf[rows_on_cpu] = 0`` syncs just the
same. It is the SCALAR right-hand side: the ``index_put_`` lowering has to
materialise and broadcast it, and that path reaches the host.
``index_fill_`` takes the scalar natively and never does.

This matters because every pool zeroes recycled rows one buffer per LAYER,
inside a Python loop. On the served 48-GDN-layer hybrid that is 96 indexed
writes -- and so was 96 syncs -- in the single step that drains a request's
admission zero-clears. It is paid per ADMISSION rather than per chunk, so
it does not amortise over a long prefill and it DOMINATES a short one.

Use this for every recycled-row reset. Assigning a scalar through
``__setitem__`` puts the sync straight back.

Per-request slab-row allocator (one row per live request).

Shared by :class:`RecurrentStatePool` and :class:`ShortConvStatePool`.
Both hand out a slab ROW per request (no paging) and reclaim it on
finish.

Zero-sentinel slab row
----------------------

Slab row 0 is reserved as a permanently-zero "sentinel" row and is
NEVER handed out by :meth:`alloc_for_request`. The free list is
initialized with rows ``1..max_num_seqs`` (skipping 0); slabs are
sized ``max_num_seqs + 1`` so row 0 stays at the boot-time zero
allocation for the engine's lifetime. See
:mod:`arbi_serve.cache.recurrent_pool` for the full rationale.

The pool keeps a LIFO ``_free_rows`` stack (pop from tail for O(1)
alloc, append on free), a ``_row_for_req`` dict, and a
``_pending_zero_rows`` list. Rows are zeroed eagerly on alloc
(deferred to :meth:`flush_pending_zero_clears`) so admission of a
new request always sees a clean slate without depending on a prior
tenant's free-time correctness.

Pending-zero-clear buckets and thread contract
----------------------------------------------

Admission (``Scheduler.add``) runs on the engine loop thread and may
land WHILE a step is in flight on the forward-executor thread, i.e.
concurrently with :meth:`flush_pending_zero_clears`. Two buckets:

  * ``_pending_zero_rows`` — the step-visible bucket. The GDN
    metadata fail-loud guard (``has_pending_zero_clears``) probes
    ONLY this bucket.
  * ``_pending_zero_rows_deferred`` — rows queued by an admission
    that landed while a forward was in flight
    (``alloc_for_request(..., forward_in_flight=True)``). Such a
    row comes off the free list, so NO in-flight batch can
    reference it — keeping it out of the step-visible bucket is
    what makes the in-flight step's metadata-build guard exact
    instead of racy. Both buckets drain together at the next flush,
    strictly before the row's owner can enter a slate.

Every bucket mutation (append and flush-swap) holds
``_pending_zero_lock`` — a plain list swap is not atomic against a
concurrent append, and a lost append is a never-zeroed slab row
(silent state corruption). The empty check in the flush stays
lock-free: a miss leaves the row queued for the next flush, which
is exactly the queue's semantics.

Subclasses must set ``max_num_seqs`` and ``device``, call
:meth:`_init_per_request_rows` once, and implement
:attr:`_row_pool_label`, :meth:`_zero_pending_rows` and
:meth:`zero_all_rows`.

Free-page allocator shared by :class:`PagedKVStatePool` and
:class:`MLAStatePool`.

Page 0 is reserved as the "null" page (so int32 block tables can use
0 to mean "unused"); usable pages are ``[1, num_pages)``. The free
list is a LIFO stack — ``pop()`` from the tail is amortized O(1).

Subclasses must set ``page_size``, ``_bytes_per_token``, call
:meth:`_init_free_pages` once, and implement :meth:`bytes_total`.

Clear every row this pool owns, the reserved sentinel included.

The boot cudagraph capture sweep drives synthetic forwards through
the slot-state blocks, so any row the sweep addressed carries writes
no request made — including sentinel row 0, which is never handed
out and so is never reached by the admission zero-clear.

Every subclass answers. A pool whose reads are bounded by the owning
request's own position count has nothing to clear and says so with
an explicit no-op body; the contract is that the answer is stated,
because a pool that merely omitted the method is indistinguishable
from one that was never wired for the sweep.

Initialize the row free-list + bookkeeping (call from ``__init__``).

``_free_rows`` is a LIFO stack initialized most-recent-first so
the first alloc returns row 1 (row 0 is the permanent zero-
sentinel and is intentionally absent from the free list).

Reserve a slab row for ``req_id`` and queue it for zero-clear.

Returns the row index. Idempotent: re-alloc for the same
``req_id`` returns the existing row without re-queuing. Raises
:class:`RuntimeError` when the slab is exhausted.

The zero-clear is DEFERRED to a pending list flushed by
:meth:`flush_pending_zero_clears` (called at start-of-step), so
the N admissions in one step coalesce into a SINGLE indexed
:func:`fill_rows_` write per slab instead of one tiny ``.zero_()``
kernel launch per admission.

``forward_in_flight=True`` (an admission landing on the engine
loop while a step runs on the forward-executor thread) routes
the row to the deferred bucket — invisible to the in-flight
step's metadata-build guard, drained by the next flush. See the
class docstring's bucket contract.

Apply every queued admission's zero-clear in one indexed write.

Drains BOTH buckets (step-visible + deferred). Builds a single
row-index tensor and delegates the per-slab zeroing to
:meth:`_zero_pending_rows`. One indexed-zero per slab total
(regardless of how many admissions are pending). Safe to call
when nothing is pending — early-returns, so the caller (engine
start-of-step) does not need to gate. The empty fast path is
lock-free; the swap holds ``_pending_zero_lock`` so a
concurrent admission's append is never lost.

Row indices as a device tensor, without a synchronous host copy.

``torch.tensor(pending, device=cuda)`` allocates PAGEABLE host memory,
and a pageable H2D ``cudaMemcpy`` is synchronous with respect to the
host — one blocking sync per flush, which the sync census sees. Staging
through a PINNED buffer makes the same copy asynchronous.

``_zero_stage_event``
guards that. The wait is on the PREVIOUS flush's copy, issued a whole
engine step earlier, so in steady state it is already complete and the
wait returns immediately — it is a correctness interlock, not a stall.

Falls back to the plain constructor on CPU, where there is nothing to
pin and nothing to synchronise.

Radix prefix-tree node + shared type aliases.

Lives outside :mod:`arbi_serve.cache.radix_pagetable` so the page-table
mixins can construct / annotate nodes without importing the main module
(which imports them), avoiding an import cycle.

One node in the radix prefix tree.

Each node represents exactly ONE page (``block_size`` tokens). The
edge from ``parent`` to ``self`` is labelled with the page's
full-tuple of token IDs (``key``).

Recurrent (Mamba / GDN / ShortConv) per-request state lifecycle.

The :class:`RecurrentLifecycleMixin` owns everything that touches the
per-request recurrent slab rows: admission alloc / free + zero-clear,
the savepoint-resume queue, host↔device savepoint snapshot / restore,
and the MTP per-token verify snapshot + partial-accept rollback. It
reads the shared ``self._views`` registry and the
``self._pending_savepoint_resumes`` queue the pool facade owns.

Apply every queued recurrent slab write, in the one order that is safe.

Zero-clears first, savepoint resumes second: a resume overlays state
onto a freshly-zeroed row, so the reverse order would wipe the
restore (see :meth:`RecurrentLifecycleMixin.flush_pending_savepoint_resumes`).

Duck-typed and unconditional. A pure-attention pool exposes neither
hook and a not-yet-wired engine has no pool at all; both are silent
no-ops here rather than a condition every caller has to restate.

The boot-ledger line relating the recurrent pool's SLABS to a savepoint ROW.

Two counts share this pool and invite the wrong division. A slab is one
``(layer, attr)`` tensor holding ``max_num_seqs + 1`` rows; a row is one
sequence's state within it. A savepoint — and so a ring slot — is one row
across EVERY slab, which is why its bytes do not move with ``max_num_seqs``
and why a pool total divided by the slab count is a slab, not a row. The
row bytes are read off :meth:`RecurrentLifecycleMixin.savepoint_row_nbytes`,
the function that sizes the ring; nothing here derives them from the slab
count or the pool total.

``None`` for a pool with no recurrent state. Duck-typed, so a pure-attention
pool reads as absent rather than as a zero-slab line.

``(rows_with_offsets, total_bytes)`` for one snapshot.

One layout function serves the sizing call (at boot, over row 0) and the
snapshot call (per boundary, over a request's row), so the ring's slot size
and the bytes a snapshot writes into it cannot disagree.

One savepoint's host buffers, filled MID-step by the forward.

WHY THIS EXISTS AT ALL
======================
Every other snapshot in this file reads a slab row AFTER the step, when
the row holds the state at the step's end. A savepoint boundary is not
at the step's end: the step runs its prefill row to a width the budget
chose, and the position a resume can be spliced at is the last KV page
boundary inside it. The recurrent slab never holds that state at any
moment the commit path can observe — the forward sweeps through it
inside one kernel — so the forward splits its fold there and hands the
intermediate state over WHILE it exists.

The GDN ops fill it once per layer from the emit staging slot
(:mod:`arbi_serve.cache._fold_emit_staging`): the recurrent kernel writes
the boundary state into the slot as it passes the boundary's chunk, and
the copies to these buffers ride the pool's spill stream behind that
kernel. :attr:`fence` is the last such copy's event, which the store's
entry carries and its consumers wait on. :meth:`stage` remains the
current-stream spelling for a caller that holds the state itself.

WHICH ROW THE BYTES MAY COME FROM
=================================
:attr:`sources` maps every ``(layer, attribute)`` to the device address
of the slab row it was laid out for, and :attr:`slab_rows` carries that
row per :class:`~arbi_serve.models.layer_spec.StateKind`. A launch proves
it is scattering into its own kind's row before it emits
(``GDNBlock._fold_emit_delta_rule``): the state it would stage is
otherwise some other request's, labelled with this one's boundary.

The row is per KIND because a request holds one row in EACH pool it has
layers in, and those rows are allocated independently — a hybrid whose
GDN row is 3 can hold ShortConv row 5. One row number would name at most
one pool's, so every other kind's identity check would compare an address
against the wrong slab and refuse by name forever.

NOTHING PARTIAL IS EVER STORED
==============================
:attr:`expected` is every ``(layer, attribute)`` a snapshot of this
model consists of, taken from the same walk
:meth:`RecurrentLifecycleMixin.snapshot_recurrent_row` uses, so the two
cannot disagree about what a complete snapshot is. A snapshot that
reached the store one attribute short would be a restore that silently
keeps a stale tensor for that layer, which is the failure mode the whole
savepoint machinery exists to prevent — so :attr:`complete` is false
until every one of them has been staged, and the commit path refuses
(and counts) rather than storing what it has. A layer kind the GDN
forward does not stage (PLE context, a Mamba row sharing the step)
therefore reads as an honest refusal, not as a snapshot.

This snapshot's slab row in ``kind``'s pool, or ``None``.

``None`` is the honest answer for a kind the request holds no row in,
and every emit treats it as a refusal: without a row there is nothing
to compare the launch's scatter target against, and an emit that
cannot prove the row is its own may not run.

Copy one ``(layer, attribute)`` of the boundary state to host.

Silently ignores a key this model's snapshot does not contain: the
forward offers what it holds, and what a snapshot consists of is
decided by the walk that built :attr:`tensors`, in one place. An
offer that matches nothing leaves :attr:`complete` false, so it
cannot pass for a snapshot.

Allocate per-request slot state across every slot-state pool.

Iterates over every :data:`SLOT_STATE_KINDS` pool view present and calls
``alloc_for_request(req_id)`` on it. The pool queues a
zero-clear for the freshly-handed-out row; the engine flushes
it in :meth:`flush_pending_zero_clears` before the next
forward reads the slab. Idempotent per pool. No-op if the
model has no recurrent / short-conv layers.

``forward_in_flight=True`` MUST be passed when the caller runs
on the engine loop while a step is in flight on the forward-
executor thread (``Scheduler.add`` mid-forward): it routes the
zero-clear to the pool's deferred bucket, invisible to the
in-flight step's GDN metadata-build guard. See
:class:`~arbi_serve.cache._pool_base.PerRequestRowPool`'s
bucket contract.

Declare the context an admitted request can reach, per slot pool.

A pool that backs its stores on demand
(:class:`~arbi_serve.cache.dsv4_pool.DSv4StatePool`) needs the
request's REACH — its prompt plus its output budget — not its
current length, so the mapping extends once per admission instead of
once per token. The reach is recorded against the request's own slab
row, so it extends that row's backing and no other. ``None`` means
the request carries no output budget and runs to the pool's own
reservation.

Runs AFTER :meth:`alloc_recurrent_state`, which is what gives the
request the row this declaration is attributed to.

Host-only by contract: the pool records the target and applies it at
the next start-of-step flush, because admission runs on the engine
loop while a forward may be in flight. Pools that back their stores
eagerly expose no hook and are skipped.

Free per-request slot state across every slot-state pool.

Returns the slab row to the pool's free list so it can be
re-used by a future admission. Idempotent — unknown
``req_id`` is a no-op on each pool. Must be paired with
:meth:`alloc_recurrent_state` from the request lifecycle hooks
(scheduler.add → finished/remove/preempt). Without this pair,
the slab row carries the prior request's end-of-decode state
forward into the next request — a state-leak bug.

Also cancels any pending savepoint resume for ``req_id``: once
the row is freed it can be handed to another request before the
resume flushes, so a stale resume would crash the flush
(``row_for`` on a gone req) or clobber the new owner's row.

Flush every per-pool pending zero-clear list in one shot.

The recurrent (Mamba / GDN) and short-conv pools defer the
admission-time zero-clear of a freshly-allocated slab row to
a pending list (see
:meth:`RecurrentStatePool.flush_pending_zero_clears`). The
engine calls this aggregator method at start-of-step, before
the model forward reads any state; it then issues one
indexed-zero per layer attribute per pool kind that has work
pending — replacing the pre-fix per-admission per-layer per-
attribute ``.zero_()`` calls (kernel launches under high
request churn).

Pure-attention models (no recurrent / short-conv view) make
this a trivial no-op.

Clear every row of every slot-state pool, sentinel rows included.

The boot cudagraph capture sweep drives synthetic forwards through
the slot-state blocks, leaving writes no request made in the rows it
addressed — sentinel row 0 included, which no admission zero-clear
ever reaches. A kind that derives "carries prior state" from the
sequence length alone would otherwise read the sweep's writes as a
live request's initial state.

Every :data:`SLOT_STATE_KINDS` view present is dispatched to, and
each answers :meth:`PerRequestRowPool.zero_all_rows` — a kind with
nothing to clear does so with an explicit no-op. Skipping a view
that does not implement it would make "this kind is exempt" and
"this kind was never wired for the sweep" the same event.

Pure-attention models (no slot-state view) make this a no-op.

True when any recurrent pool still holds queued admission
zero-clears.

The recurrent-state correctness invariant: a slab row handed out
at admission is ZEROED before any forward reads it (the GDN
prefill recurrent path gathers h0 UNMASKED — see
``_gdn_fla._forward_prefill_fla`` — so a stale row corrupts
silently, not loudly). This probe is the cheap host-side check
(one small dict walk + list-truthiness per view) behind the
fail-loud guards at :meth:`flush_pending_savepoint_resumes` and
the recurrent metadata builders.

Probes ONLY the step-visible bucket. Rows in the deferred
bucket (queued by an admission that landed while a forward was
in flight) are intentionally excluded: a deferred row comes off
the free list at queue time, so no in-flight batch can
reference it — probing it would make the metadata-build guard
trip on a provably-safe concurrent admission.

True when any recurrent view still holds a queued zero-clear
(either bucket) for ``req_id``'s slab row.

Used by :meth:`flush_pending_savepoint_resumes` to keep the
per-row ordering invariant (zero-clear BEFORE restore) across
the deferred bucket: a restore applied while the row's
zero-clear is still queued would be wiped by the later flush.

Queue a savepoint restore for ``req_id``.

``tensors`` is the per-(layer_idx, attribute) host-tensor dict
produced by :meth:`snapshot_recurrent_row` (or its CPU clone
from :class:`SavepointEntry.tensors`). The actual
``restore_recurrent_row`` call is deferred to
:meth:`flush_pending_savepoint_resumes`, called at start-of-
step before the model forward reads any state.

Idempotent at the per-(req_id) level: the LATEST queued tensors
win (dict overwrite). In practice each request enters the queue
at most once per admission.

Pure-attention models (no recurrent / short-conv view) MUST NOT
reach this method — :meth:`Scheduler.add` only queues resumes
for hybrid models. Calling it on a pure-attention pool would
be a no-op at flush time but a wasteful queue alloc; we don't
defend against it here.

Queued on the engine loop thread; the flush may run on the
forward-executor thread — every queue mutation holds
``_savepoint_resume_lock`` (a lost entry is a request decoding
from zeroed instead of restored state, silently).

Resolve any completed restore-timing events into per-request ms.

Lazy on purpose: ``Event.elapsed_time`` requires the events to have
completed, and forcing that on the step path would sync the engine on a
copy the step does not otherwise wait for. Called from
:meth:`pop_savepoint_resume_ms` at request FINISH, by which point the
restore has long since retired, so ``query()`` is true and no wait
happens. Events not yet complete stay queued for the next drain.

Device milliseconds spent restoring ``req_id``'s recurrent state.

``0.0`` when the request had no savepoint resume (a pure-attention
prefix hit reuses pages by pointer and restores nothing) or when the
timing could not be resolved. Pops, so the entry cannot leak.

Drop any queued savepoint resume for ``req_id``. No-op if absent.

Called from :meth:`free_recurrent_state` so a request whose slab
row is freed (remove / preempt / suspend) before the resume
flushes never restores into a row that has since been handed to
another request.

Drain the queued savepoint resumes into device-side slab rows.

Call site: engine start-of-step, IMMEDIATELY AFTER
:meth:`flush_pending_zero_clears`. Order is critical — the
zero-clear runs first (clearing the freshly-allocated slab
row), then this method overlays the savepoint state on top of
the zeroed row. Reversing the order would zero-out the
restore.

Each queued entry triggers one :meth:`restore_recurrent_row`
call — a synchronous per-(layer_idx, attribute) ``copy_(non_blocking=True)``
from the host clone into the on-device slab row. The host
tensors come from :class:`SavepointEntry.tensors` which are
independent CPU clones (see
:meth:`snapshot_recurrent_row`'s ``copy=True`` invariant), so
the copy is safe — no aliasing with future engine work.

Cudagraph-safety note: the captured prefill graph has not yet
been replayed when this method runs (we are at start-of-step,
pre-forward). The slab rows the captured kernels reference are
the same slabs we write into here. Live-GPU validation under
captured prefill replay needs to confirm the restore lands
before any captured kernel reads the row — the engine's call
site placement (mirroring zero-clear) gives that ordering by
construction.

Safe to call when nothing is pending — early-returns. The
caller (engine start-of-step) does not need to gate.

Raises if a queued ``req_id`` is unknown to the recurrent pool.
A free between queue and flush is NOT a bug — it cancels the
resume via :meth:`free_recurrent_state` →
:meth:`cancel_pending_savepoint_resume`, so the only way to reach
the flush with a gone req is an alloc that never happened (a real
scheduler bug); hard-fail.

An entry whose row still has a zero-clear queued in the pool's
DEFERRED bucket (its admission landed while a forward was in
flight, after this step's zero-clear drain) is NOT restored now
— restoring it would be wiped by the queued zero-clear. It
stays queued and flushes at the next drain pair, strictly
before its owner can enter a slate. The step-visible-bucket
raise above is unchanged: pending clears there mean the caller
skipped the paired ``flush_pending_zero_clears()``.

Yield ``(layer_idx, attr, slab_row_tensor, flavor)`` for ``req_id``.

Walks every recurrent / short-conv pool view, resolves the
request's slab row via ``row_for`` (raises :class:`KeyError`
if unknown), and yields one entry per ``(layer, attribute)``
state slab the savepoint snapshot/restore touches:

  * GDN: ``"recurrent_state"`` then ``"conv_state"``
  * Mamba1 / Mamba2: ``"ssm_state"`` then ``"conv_state"``
  * ShortConv: ``"conv_state"``

``slab_row_tensor`` is the per-row slice (``slab[row]`` /
``getattr(lv, attr)[row]``) and ``flavor`` is the pool-kind
label (``"GDN"`` / ``"Mamba"`` / ``"ShortConv"``) used to
phrase restore error messages. Shared by
:meth:`snapshot_recurrent_row` and
:meth:`restore_recurrent_row` so the per-view-type dispatch
lives in exactly one place; iteration order is the contract
both directions rely on.

Serialize ``req_id``'s slab row into a dict of host tensors.

Walks every recurrent / short-conv pool view, looks up the
request's slab row, and clones each per-attribute tensor
slice (one per layer × attribute) to the CPU. The returned
dict is the ``tensors`` payload of :class:`SavepointEntry`.

Per-pool attributes:

  * GDN: ``"recurrent_state"`` + ``"conv_state"``
  * Mamba1 / Mamba2: ``"ssm_state"`` + ``"conv_state"``
  * ShortConv: ``"conv_state"``

Synchronous mode (``async_copy=False``, default)
------------------------------------------------
Synchronous device→host copy. Caller (the prefill-commit hook)
is responsible for ensuring the slab row's content is the
post-chunk-boundary state — typically by inserting a stream
sync after the prefill chunk's forward and before this call.

Async mode (``async_copy=True``)
--------------------------------
Allocates ONE pinned host arena for the whole snapshot (per-
tensor views into it) and issues ``non_blocking=True`` D2H
copies on the current CUDA stream WITHOUT a host-side sync.
The returned dict carries the (still-empty until the copies
retire) host tensors. The caller MUST also invoke
:meth:`snapshot_recurrent_event` immediately after to obtain a
``torch.cuda.Event`` recorded on the same stream — the
consumer-side ``SavepointEntry.synchronize()`` waits on that
event before reading the host buffers (and before freeing
them on eviction).

Lazy-confirm contract: chunk N+1's forward shouldn't need
chunk N's snapshot to be complete on host until the entry is
actually consumed (admission-time restore) or evicted (host
page free). The async path moves the D2H off the commit
critical path; the cost surfaces only at admission, when it
is overlapped with admission's other work and is tens-to-
hundreds of µs at worst.

Falls back to the synchronous path automatically when CUDA is
not available or the slab tensors are not on a CUDA device
(CPU-only test pools, ``device='cpu'``). Under fall-back the
caller's :meth:`snapshot_recurrent_event` returns ``None`` and
the entry's ``pending_event`` stays ``None``.

Pure-attention models (no recurrent / short-conv pool view)
return an empty dict.

Idempotent + non-mutating: snapshot doesn't change device-side
state. Safe to call multiple times for the same row at the
same chunk boundary; the returned host tensors are independent
clones.

Raises :class:`KeyError` if ``req_id`` has no slab row
allocated — a bug in the caller (snapshot before
``alloc_recurrent_state`` was called or after
``free_recurrent_state``).

Lay out the host buffers a MID-step savepoint will be staged into.

Called at metadata-build time, before the forward runs, for the one
prefill row whose step contains a resume-grid boundary. Returns
``None`` for a model with no recurrent state (nothing to snapshot)
and for a request with no slab row (nothing to read).

The layout is :func:`_snapshot_arena_layout` over the request's own
row — the same function that sizes a ring slot at boot and the same
walk :meth:`snapshot_recurrent_row` takes — so a boundary snapshot
and an end-of-step snapshot are the same bytes in the same order,
and a restore cannot tell them apart.

Buffers come from the ring when one is attached, exactly as they do
for an end-of-step snapshot, and the hand-out's stamp travels with
the snapshot so the store can refuse it once the slot is reused. The
copies themselves are NOT routed to the spill stream: a spilled copy
is fenced against the NEXT step's slab write, and this one has to
land between two kernels of the CURRENT step, which is exactly what
the current stream already guarantees.

``offset`` is where in the row's own segment of this step the fold
splits; it travels with the snapshot so the metadata builder can hand
the layers a plan without re-deriving anything the scheduler decided.

A SLOT IS SPENT WHETHER OR NOT THE SNAPSHOT COMPLETES
=====================================================
The ring hand-out happens here, before the forward runs, because the
forward stages into these buffers and they must exist first. So a
step that arms and then stages nothing — a compiled/captured GDN
prefill, which cannot split (see
``GDNBlock._forward_prefill_fla``) — still consumes a slot, and the
ring wraps that much sooner, costing RETENTION of older snapshots
rather than correctness. That is what a nonzero
``snapshot_skipped_total{reason="fold_split_not_staged"}`` is telling
an operator, and it is the number to read before concluding the ring
is too small.

Only one boundary snapshot is open at a time. A second open replaces
the first, which is the honest outcome: the slate advances one
prefill row depth-first (``Scheduler.schedule``), so a second armed
row in one step means the caller armed something it should not have,
and carrying both would put two labels on one ring slot.

Each recurrent pool's slab row for ``req_id`` — host-side, no device read.

A kind whose pool has no row for the request is absent from the map
rather than present as ``None``: the emit's identity check reads a
missing row as "this snapshot names no row for my kind" and refuses,
which is the same answer without a sentinel to mistake for row 0.

Hand the step's boundary snapshot over and clear it.

Cleared unconditionally, complete or not: a snapshot left behind
would be staged again by the next step's forward and then labelled
with THIS step's boundary — a coverage label naming a state it does
not hold, which is the one thing this machinery may never do.

Bytes ONE snapshot occupies, laid out exactly as a snapshot lays it.

Sized off row 0 of every recurrent slab — the geometry is a property of
the model, not of a request — so the ring can be allocated at boot,
before any request exists, and priced on the boot ledger.

Slabs one savepoint row spans: one per recurrent ``(layer, attr)`` tensor.

Counted over the same walk :meth:`savepoint_row_nbytes` sizes from, so
a ledger line printing both the count and the row bytes describes one
set of tensors.

Page-lock the snapshot ring on the HOST; return its total bytes.

Called once at boot, never lazily: page-locking stalls the device, so a
route that allocates per snapshot pays that on the first long prompt of
every boot. Allocating the whole bound up front is the only reason the
request path can be allocation-free.

The bytes are HOST RAM. Nothing is taken from the KV budget, so the slot
count is bounded by host memory rather than by servable context — and
that bound is enforced, not assumed: the whole reservation is booked
through :func:`~arbi_serve.runtime.pinned_host_budget.reserve_pinned_host`
BEFORE a page is locked, so a host that cannot afford it refuses here
rather than handing the decision to the kernel's OOM killer later.
``slots`` is expected to have come from
:func:`~arbi_serve.cache._savepoint_ring.plan_savepoint_ring`, which
asked the same budget; the booking is what makes the ring a row on the
host ledger and what makes a later reservation see these bytes.

The ring itself is the PROCESS's, not this pool's
(:func:`~arbi_serve.cache._savepoint_ring.acquire_shared_ring`): a boot
builds the state pool up to four times and every one of them lands
here, so a ring per pool page-locked the same bytes four times over.

Returns 0 without allocating when the pool holds no recurrent state.

``(kind, rec_slab, rec_attr, conv_slab, n_layers)`` the emit slot is sized off.

One boot emits for ONE kind: the slot is a single row of a single
pool's layout, and a model mixing two recurrent kinds would need a
slot per kind before either could stage. The first kind present in
``RECURRENT_KINDS`` order that carries a conv window wins, and a
second recurrent kind on the same boot reads as an incomplete
snapshot at the commit seam rather than as a wrong one.

``rec_slab`` is ``None`` for a kind whose entire state is the conv
window, which is what tells the staging it carries one attribute;
``rec_attr`` is what THAT kind's snapshot walk calls the matrix, which
is the name the spill has to land under.

``(the layer view's recurrent matrix, what its kind calls it)``.

The NAME travels with the tensor because the snapshot walk keys its
host buffers by it (GDN's ``recurrent_state``, Mamba's ``ssm_state``),
and a slot spilled under the other kind's name lands nowhere.

Tested against ``None`` explicitly: a tensor has no truth value, so
an ``or`` chain over these raises rather than falling through.

Attach the fold-state emit's device slot(s) and tables; return their bytes.

Sized off the GDN pool's own geometry — one recurrent row in the
slab's layout and dtype, one conv window, the emit table and the conv
index table — so the bytes are the model's, not a guess, and the KV
grow that follows prices them by measuring free VRAM after they
exist. The staging is the PROCESS's: a captured graph bakes its
addresses, and the pool is rebuilt after the deferred KV resize, so a
rebuild with the same geometry finds the same object
(:func:`~arbi_serve.cache._fold_emit_staging.set_staging`).

``landing`` page-locks the capture landing zone — one ring slot's
worth of pinned host RAM laid out as a snapshot — for a boot whose
capture mode can replay a prefill. Returns 0 without allocating when
the pool holds no GDN view.

Give the emit's device slot back — this model supplies no emitting kernel.

:meth:`attach_fold_emit_staging` sizes the slot off the pool's GDN
views, which is geometry; whether any layer can EMIT is a fact only
the dispatch registry holds
(:func:`~arbi_serve.cache._fold_emit_staging.emitting_block`). A kind
that borrows the GDN slab without the vendored emitting chunk entry
gets a slot sized for it, and holding it would leave device bytes and
an ``expected_layers`` count standing for a path this boot cannot
take, while ``gdn_prefill_conv_flat`` — which such a kind does call —
staged conv windows into a slot nothing spills.

Detaches the process global too: the ops read the staging from there,
not from the pool.

Time one slot's write and restore; return ``(write_us, restore_us)``.

The host ring trades VRAM for a PCIe transfer, so the transfer is the
number an operator needs and it belongs to this boot on this card —
bus width, pinned-ness and slot size all move it. Measured once at
attach, against the ring that was just allocated, so the boot log can
state the cost in the units the slot now costs.

``None`` off CUDA, or with no ring attached.

Record the start of a DIRECT-route snapshot's copies on the current stream.

Called only where those copies are issued on the current stream;
:meth:`snapshot_recurrent_event` closes the pair with the fence it
records right after them, so the pair brackets exactly the copies. A
spilled snapshot never opens one — its copies are on the spill stream,
which times them itself. Never fatal: without CUDA there is no event
and no observation. A second open before a close drops the earlier
start rather than leaking it.

Resolve retired direct-route write pairs into ``arbi_serve.savepoint.write_ms``.

Events not yet complete stay queued for the next drain, so this never
waits on the device. The spill route's pairs drain inside the spill
sequencer, at its fence, into the same histogram.

Record a CUDA event on the current stream.

Paired with ``snapshot_recurrent_row(async_copy=True)``. The
caller issues the per-layer non-blocking D2H copies, then
calls this method to fence them with a single event. The
event is then attached to the :class:`SavepointEntry` and
consumed at ``get`` / eviction time (host-side wait).

Returns ``None`` when CUDA is unavailable — the snapshot
path fell through to synchronous ``to('cpu', copy=True)``
which already finished the D2H before returning, so no event
is needed.

Why a method on the pool (not on the snapshot dict): the
event must be recorded on the CUDA stream the snapshot's
``copy_`` calls were issued on. Today every recurrent slab
lives on the engine's default stream (the engine doesn't
side-stream the recurrent kernels); recording on
``current_stream`` of the pool's device matches the same
stream the copies were issued on.

Copy host-side ``tensors`` into ``req_id``'s slab row.

Inverse of :meth:`snapshot_recurrent_row`. The dict shape MUST
match what ``snapshot_recurrent_row`` would produce for the
same model — same layer indices, same attribute names, same
per-layer tensor shapes. Mismatches RAISE ("no silent
fallbacks").

Cudagraph-safety constraint
===========================

Direct ``slab[row].copy_(host_tensor)`` from a fresh host
tensor is UNSAFE under captured prefill graphs — the captured
kernel holds the slab's device pointer, but the host source
of an ad-hoc ``.copy_()`` from a freshly-allocated host
tensor uses a per-call staging buffer whose device pointer
becomes stale after the next sync.

This method's synchronous copy (pre-step, pre-forward) is
correct under eager mode and cudagraph capture-only paths. The
caller must issue this method only from the start-of-step
flush (mirroring ``flush_pending_zero_clears``'s placement).

Idempotent: repeated restore with the same ``tensors`` dict
produces the same end state. Free to retry.

Raises:
  * :class:`KeyError` if ``req_id`` has no slab row.
  * :class:`RuntimeError` if a layer×attribute key the pool
    expects is missing from ``tensors``, or if a tensor's
    shape/dtype/device doesn't match the slab's
    corresponding row.

Hard-fail on shape / dtype mismatch between savepoint and slab.

Configuration drift between snapshot and restore (different
TP slicing, different dtype, different model — all things
that should never co-exist in one engine but DO if a stale
store is loaded) corrupts the recurrent kernels silently.
Raise loud.

Min free per-request slab rows across recurrent / short-conv pools.

On hybrid models (GDN / Mamba / ShortConv) the per-request slab
row is a SCARCE resource distinct from paged-KV pages — a batch
backlog can exhaust recurrent rows while paged-KV pages stay
plentiful. The suspended-job offload policy watches this so it
parks (and frees the recurrent row of) resident batch jobs under
recurrent-pool pressure too. Returns ``None`` on pure-attention
models (no recurrent pool), so the caller skips the check.

Physical every parked state arena will map on the next wake.

The wake maps the arenas AFTER the growable-KV re-provision, so the
re-provision's budget has to hold these bytes out or the arenas run the
card dry behind it (:func:`arbi_serve.engine.sleep._growable_kv_wake_budget`).
Read from the arenas rather than predicted: 0 while they are awake.

Yield ``(kind, view)`` for every per-request-slot pool view — all
of :data:`SLOT_STATE_KINDS`, not one favoured kind.

Every slot-state kind advances its state on EVERY token the verify
forward consumes, so every one of them must be reached by the
snapshot-attach and partial-accept rollback paths below. A kind
filtered out here silently keeps the rejected drafts' state and
speculative decoding stops being output-preserving. That is why the
predicate is :data:`SLOT_STATE_KINDS` itself and never a second
enumeration beside it.

The named snapshot / rollback entry point on ``view``, or RAISE.

A recurrent kind that does not implement it cannot roll its state
back to the accepted prefix; refusing here is the only correct
outcome, since serving on would emit tokens conditioned on rejected
drafts.

Allocate per-token verify snapshot buffers across recurrent pools.

Called from :func:`arbi_serve.engine.build.attach_mtp_driver`
once the driver's ``max_k`` is known. Pure-attention models
(no recurrent pool view) make this a no-op. Each view's
snapshot is sized ``(max_k_plus_1, max_num_seqs, *state_geom)``.

Every recurrent kind present is dispatched to. A kind whose pool
has no snapshot path fails LOUD here — at drafter attach, before
the engine serves a single token.

Roll partial-accept rows' recurrent state back to ``snap[n_accepted]``.

Per partial-accept row, look up its slab row index (in the
recurrent pool's slab), read ``snap_X[n_accepted, slab_row]``
from the verify-pass snapshot, and copy it into
``slab[slab_row]`` for every state attribute on every recurrent
pool. Issues one indexed copy per layer per row — no replay
forward.

``slab_rows`` and ``n_accepted_per_row`` are parallel lists
(same length); each row's rollback target is ``snap[n_acc,
slab_row]``. Pure-attention models make this a no-op (no
recurrent views).

Hot-path implementation: build a single ``(slab_rows,
n_accepted)`` tensor pair on the recurrent pool's device and
dispatch to :meth:`RecurrentStatePool.rollback_batch`, which
issues one indexed copy per layer attribute.

Device-tensor partial-accept rollback — NO host sync.

The GPU-resident analogue of :meth:`rollback_partial_accept`. The
async hybrid GDN/Mamba verify path calls this with the DEVICE
``n_accepted`` ``(B,)`` tensor straight off the greedy verify op —
before any ``.tolist()`` host pull — so the recurrent reconcile is
enqueued data-dependent on the device count and rides the compute
stream with zero host round-trip. This is the piece that lets arbi
run async MTP-verify for hybrid GDN/Mamba models.

``slab_rows`` is the per-row REAL pool slab row (``row_for``) on
this pool's device; ``n_accepted`` is the per-row accepted-draft
count ``(B,)``, same length and device. Both are passed through to
:meth:`RecurrentStatePool.rollback_batch` with
``no_host_sync=True``.

``k_uniform`` is the uniform draft depth ``K``, used only by the
RECOMPUTE rollback to skip fully-accepted rows whose committed
post-token-K state the forward already left in place. MASKED REPLAY
ignores it and commits every row it is handed: its verify forward
commits no state, so a row left uncommitted keeps the state it
entered the step with.

Ordering. The caller enqueues this on the COMPUTE stream
immediately after the verify forward and before parking the
deferred step, so it lands before step N+1's GDN forward AND before
step N+1's verify forward overwrites the replay buffers —
single-buffer-safe because the compute stream serializes them.

Bounded ring of pinned-host buffers for recurrent savepoint snapshots.

A savepoint is one sequence's recurrent state at a prefill chunk boundary. Two
costs decide where it can live, and they pull in opposite directions.

Page-locking is one. Allocating a pinned arena per chunk boundary stalls the
DEVICE, not merely the calling thread, so a route that page-locks per snapshot
pays that on the first long prompt after every boot. That is what a BOUNDED ring
fixes: ``slots`` buffers, page-locked ONCE at boot, handed out round-robin. The
allocation leaves the request path entirely, and the bound is what makes
allocating them all up front possible.

Residency is the other, and it is why the ring is on the host. A slot is one
sequence's whole recurrent state, and on a card with no headroom VRAM is KV — so
a device-resident slot costs servable CONTEXT, measured at ~35 KV pages each on
the boot ledger. Retention then trades directly against context, and the count
has to stay small enough to be useless: at eight concurrent conversations, a
one-slot ring serves ZERO resumes, because each conversation's snapshot is
overwritten by its neighbours before its own next turn. Pinned host memory has
no such exchange rate against KV: the slot count can cover real concurrency and
the KV pool is untouched.

It has a DIFFERENT ceiling, and it is not softer. Page-locked pages cannot be
swapped or reclaimed, so a ring the host cannot afford does not degrade — the
kernel's global OOM killer picks a victim, and on a shared box the engine is
the fattest one. So the slot count resolves against two inputs, not one:
concurrency says what retention needs, and
:mod:`arbi_serve.runtime.pinned_host_budget` says what the host can page-lock
beside the engine's own measured working set. :func:`plan_savepoint_ring`
takes the smaller and reports both.

What the host costs instead is the transfer. A snapshot is a device→host copy
and a resume a host→device one, over PCIe rather than within the card. Both are
bounded by the slot size and both are issued on the SPILL STREAM
(:mod:`arbi_serve.cache._savepoint_spill`), fenced where the slab row is next
written, so neither sits on the decode critical path. The alternative to a
transfer is re-prefilling the chunk it stands in for, which is the comparison
that matters.

**Reuse is the correctness core.** When the ring wraps, the slot's previous
occupant is overwritten, so any store entry still pointing at it must stop being
servable. Each hand-out therefore carries a ``(slot, generation)`` stamp and
:meth:`is_live` answers whether that stamp still owns the slot; the store checks
it before serving an entry and drops it when it does not. A stale entry is a
MISS — the request re-prefills, exactly as it would have if the snapshot had
never been taken — never a wrong restore.

The allocator is injected so the sequencing is exercisable without a GPU.

A hand-out was given back that cannot safely be given back.

Raised only for the two shapes that would corrupt the ring rather than
merely fail to recover a slot: a slot index the ring does not have, and a
stamp that has already been released. A double release would rewind
``_next`` twice, so a later hand-out would return a slot the ring believes
is two positions ahead — the round-robin order and the generations would
stop describing the same thing, which is the same class of defect as the
stale slab row ``_recurrent_lifecycle`` documents, arriving from the other
side.

A release that is merely too LATE (another hand-out has happened since) is
not this: it is answered ``False``, because nothing is wrong with the ring
and there is simply nothing left to recover.

Raised when a hand-out is requested with no slot free to overwrite.

Cannot happen with the round-robin policy (every slot is overwritable);
exists so a future policy that pins a slot has somewhere to fail loudly
rather than silently serving a slot that is still in use.

Slots to retain: the operator's value, or AUTO from concurrency.

``configured > 0`` pins the count. ``0`` means AUTO, and auto is a function
of ``max_num_seqs`` rather than a constant, because the quantity the ring
has to survive is CONCURRENCY: the ring is global, so a sequence's snapshot
must outlive every snapshot its peers lay down between two of its own
turns. A number that does not move with ``max_num_seqs`` is wrong on every
deployment but the one it was measured on.

``max_num_seqs`` is the only concurrency estimate available at boot — no
count of active conversations exists before any request has arrived — and
it is the right conservative one, since sequences beyond it are queued and
cannot be interleaving snapshots.

``entry_floor`` is the store's own retention figure — the entry count its
byte cap admits (``savepoint_min_entries`` when the cap is defaulted, the
cap divided by a slot when it is stated). The round-robin bound above
assumes ONE write per sequence per round, which held while only a
page-aligned step end could write; a prefill now writes at every
``chunk_prefill`` crossing and at its last page (arbicity/arbi-serve#2238),
so a sequence's turn holds several entries and the store's floor is the
figure that already prices that on the host. Two halves of one machinery
must not disagree about how many snapshots are retained: the ring keeps
at least what the store budgets, and :func:`plan_savepoint_ring` still
fits the result to the host.

Never returns less than 1: a ring must have somewhere to put a snapshot.

The slot count this boot may actually have, and everything that set it.

Two inputs, and the ring needs both. CONCURRENCY says how many slots
retention requires (:func:`resolve_savepoint_ring_slots`); HOST HEADROOM
says how many the machine can page-lock without putting the kernel's OOM
killer in charge of which process dies. The count is the smaller, and the
plan carries the other one so the boot line can state what was asked for
beside what was granted — a ring that silently shrank is a retention
figure nobody can reproduce, and a ring that silently did not shrink is
how the engine got killed mid-request.

Resolve the ring against BOTH of its inputs: concurrency and host headroom.

``slot_bytes`` comes from the pool's recurrent geometry and is known before
anything is allocated, which is what lets the whole reservation be priced
before a single page is locked. ``host`` / ``already_pinned_bytes`` are
injectable so a starved host and an ample one are both reachable in a test
without being on one.

A grant of 0 turns the ring OFF rather than building a ring that cannot be
afforded. The savepoint path then falls back to its per-boundary host arena
— slower per snapshot, but bounded by one snapshot rather than by the whole
retention window.

The boot line for a ring that was built: count, WHERE it came from, cost.

A slot count with no parameter beside it reads as a tuned number and the
next reader cannot tell whether their own ``max_batch`` was accounted for.
When the HOST is what set it, that is stated too, with the arithmetic — a
ring that quietly shrank is a retention figure nobody can reproduce, and
the alternative to shrinking loudly is not "16 slots", it is a SIGKILL.

``slab_count`` is the pool's own walk. A slot is one ROW across every
recurrent slab, and a reader dividing the pool total by the slab count
gets a slab (``max_batch + 1`` rows), not a slot; naming the unit on the
line is what stops that division from looking right.

Pure string assembly so both branches are readable off a test rather than
off a boot.

The boot line for a ring that was NOT built, and why.

Zero slots is a real outcome of the host budget, not an error, and it must
not be silent: the savepoint path falls back to a per-boundary pinned arena
that is bounded by ONE snapshot instead of by the retention window, which
is a different performance profile the operator has to be able to see.

The process's ring at ``(slots, nbytes)``, built once and reused.

**A boot builds the state pool up to FOUR times** — the throwaway KV warmup
pool, the capture-sizing measurement pool, its restore, and finally the
serving pool — and every one of them runs ``build_active`` and asks for a
ring. The geometry is identical each time, because a slot's size is the
model's recurrent geometry and the slot count comes from ``max_batch``: so
building a new ring per pool page-locked the same bytes over and over,
each ``cudaHostAlloc`` synchronizing the DEVICE, and left three rings' worth
of host RAM to the garbage collector to give back at a time nobody chose.

Reusing it is safe because nothing survives a pool rebuild that could still
be pointing into a slot: ``build_active`` rebuilds the savepoint STORE too,
so every entry stamped against the old hand-outs is gone with it. The
``(slot, generation)`` reuse contract is untouched — it is the same ring,
with the same counters.

A DIFFERENT geometry (a model swap under residency) replaces the ring and
re-books it; the old one is released from the host ledger first, so the
replacement is priced against a host that is not counting its predecessor.

``slots`` pinned-host buffers of ``nbytes``, handed out round-robin.

``alloc_fn(nbytes)`` returns one buffer; injected so the ordering is
testable off-GPU. ``slots`` is the retention bound: the ring holds the
newest ``slots`` snapshots and nothing older.

The backing buffer for ``slot``. For the boot transfer probe only.

Reading a slot outside a hand-out would race a live snapshot; at boot
there are none, which is the only time this is called.

Hand-outs, overwrites, and the geometry.

``overwrites`` is how often the ring wrapped onto a slot that had
already been used — i.e. how often retention was exceeded. A
deployment reading zero has never lost a snapshot to the bound.

``releases`` is how often a hand-out was given back unused
(:meth:`release`). It is not a subset of ``overwrites`` and does not
cancel one: the hand-out's generation bump already invalidated the
slot's previous occupant, so a release recovers the ring's position
and not that snapshot. Read the pair together — ``releases``
approaching ``handouts`` says most hand-outs are being taken by steps
that store nothing, which is retention paid for and not received, and
the reason belongs to whoever is taking them.

Claim the next slot; return ``(buffer, slot, generation)``.

``nbytes`` must fit the slot — the snapshot size is a property of the
model, so a mismatch is a wiring bug (a different resident, or a pool
whose geometry changed under a live ring) and fails loudly rather than
truncating a snapshot.

Give back a hand-out whose bytes were never used. True if recovered.

WHY THIS EXISTS
===============
A hand-out is taken BEFORE the writer knows whether it will produce a
snapshot — the mid-step boundary snapshot
(:meth:`~arbi_serve.cache._recurrent_lifecycle.RecurrentLifecycleMixin.
open_boundary_snapshot`) has to own its buffers before the forward can
stage into them. A step that then stages nothing would otherwise
advance the ring by one, and a run of such steps walks the whole ring:
after ``slots`` of them EVERY retained snapshot has been invalidated
by hand-outs that stored nothing. Rewinding puts the next hand-out
back on the same slot, so one slot absorbs the whole run and the rest
keep their snapshots.

WHAT IT CANNOT RECOVER, AND WHY THAT IS THE RIGHT TRADE
======================================================
The slot's PREVIOUS occupant is already gone and does not come back.
:meth:`take` bumps the generation at hand-out, which is what makes the
buffer safe to write into: from that moment the store reads the old
entry as stale and stops serving it, so nothing can read bytes that
are being overwritten. Deferring the bump until the snapshot completed
would recover that one entry and open a window in which the old entry
is still servable while the new one is being copied on top of it —
trading a certain, bounded retention loss for a possible wrong
restore. This ring does not make that trade, so ``release`` recovers
the ring's POSITION and never a generation.

Refuses (``False``) when another hand-out has happened since — the
rewind would then hand out a slot that is currently being written.
Raises :class:`SavepointRingReleaseRefused` for a double release or an
unknown slot; see that class for why those two are different.

Recurrent-savepoint device→host copy on a spill stream, fenced where the row is reused.

The savepoint snapshot (``MultiStatePool.snapshot_recurrent_row``) serializes one
slab row — every recurrent layer's state plus its conv window — into pinned host
RAM at each prefill chunk boundary. Issued the direct way, that is one
``non_blocking=True`` device→host copy per (layer × attribute) on the CURRENT
stream. ``non_blocking`` there is a statement about the HOST: it returns before
the transfer retires. On the DEVICE the copy still occupies the compute stream.

This module issues the same copies on a spill stream instead, and holds the one
piece of ordering that makes them safe.

The hazard, and where the fence belongs
=======================================

The slab row is written IN PLACE by the next step's recurrent kernels. A spill
stream reading it while the compute stream overwrites it tears the snapshot
across a chunk boundary, and the entry would be labelled with a
``num_tokens_covered`` the bytes do not hold. So the compute stream must wait for
the spill — the only question is WHERE.

Fencing at the snapshot's own call site (``Scheduler.commit``, right after the
step's forward) re-serializes the transfer: nothing is queued behind it there, so
the wait lands immediately and the compute stream sits out the whole copy. The
fence belongs instead at the point the row is next WRITTEN, which is the start of
the following step — :func:`drain_pending_recurrent_flushes`, on the recurrent
backend's ``build``. That seam already exists for the same class of invariant
(admission-queued zero-clears must land before the forward gathers ``h0``), and
it is chosen for the same reason: draining where the slab is touched makes the
ordering a property of the access rather than of the caller.

Whether that leaves the copy free is a question about a particular deployment,
not a property of the design: it depends on how much host work separates the
snapshot from the next step's first slab write. The fence therefore reports
whether the spill had already completed when it ran
(``arbi_serve.savepoint.spill_fence_total`` and
``arbi_serve.savepoint.spill_fence_not_ready_total``), so the answer is read off
a counter rather than assumed.

Two durations sit behind those counters and are timed here as well, both with
CUDA events and both read LAZILY — the elapsed time is taken only once the
events are known to have retired, never by synchronizing on the step path:

* ``arbi_serve.savepoint.write_ms`` — device time one snapshot's copies occupy
  the SPILL stream, from just after it has waited for the row's writes to just
  after its last copy. This is what a ring slot costs to fill. The histogram is
  shared with the direct route, where the copies ride the current stream and
  the pool records the pair there
  (``RecurrentLifecycleMixin._open_savepoint_write_span``); one name, and each
  route times the stream its copies are actually on.
* ``arbi_serve.savepoint.spill_fence_residual_ms`` — device time the COMPUTE
  stream sat on the fence when the spill had not retired by the time the fence
  ran. Recorded only on that ``not_ready`` branch, so its count equals that
  counter and its sum is the stall the counter only hints at. It is measured
  with events on the compute stream because the wait is a stream wait
  (``wait_event``), not a host wait: the host never blocks at the fence, so a
  host-side timer around it would measure nothing.

Contract
========

The fence event this returns takes the place of the direct path's event on the
:class:`~arbi_serve.cache.recurrent_savepoint.SavepointEntry`. Consumers
(``get``, eviction) call ``entry.synchronize()`` before reading or freeing the
host buffers, so the lazy-confirm contract is unchanged: the host bytes are valid
after the event, and only after it.

No device buffer is allocated. The copies read the slab row directly and land in
the entry's own pinned host arena, so this costs no VRAM and takes nothing from
KV.

The stream/event constructors are injected rather than reached for, so the
ordering above is exercisable without a CUDA device.

Sequencer for the spill-stream savepoint copy. One instance per pool.

``event_factory`` builds the ordering events (fences); ``timing_event_factory``
builds the timing-enabled ones. They are separate because a fence event is
waited on by a stream and a timing-enabled event is slower there, while a
timing event is never waited on — only queried and read.

The pool's sequencer, built on first spill and cached on it.

``None`` off CUDA: a CPU pool's snapshot path is synchronous already, so
there is no stream to route the copies to.

Route one snapshot's copies to the spill stream; return the fence, or ``None``.

``None`` means the spill route did not run — the flag is off, or the pool is
not on CUDA — and the caller issues its own device→host copies. On the spill
route ``out`` is filled with the same pinned host views the direct route
would produce, so the entry's payload is identical either way; only the
stream the bytes travel on differs.

Route one snapshot's copies into a RING SLOT on the spill stream.

The same sequencing as :func:`spill_snapshot_copies`, over a buffer the ring
page-locked at boot rather than one allocated for this snapshot. It is NOT
gated on ``savepoint_async_spill``: that flag chooses whether the per-boundary
host route is worth a second stream, while for the ring the side stream is
the reason the route is affordable at all — the copy crosses PCIe, and on the
compute stream it would sit in front of the next decode step.

``None`` when there is no spill stream to route to (a CPU pool, or no CUDA);
the caller then issues the same copies itself.

Fence a pool's in-flight savepoint spill against this step's slab writes.

Called from :func:`~arbi_serve.cache._recurrent_lifecycle.drain_pending_recurrent_flushes`,
which every recurrent forward reaches before it touches the slab. Duck-typed
and unconditional: a pool that never spilled has nothing outstanding and
returns ``None``.

Issue one snapshot's copies on the spill stream; return its fence.

``plan`` is a sequence of ``(host_view, src)``: ``host_view`` slices the
entry's pinned host arena and ``src`` is the slab row's tensor for one
(layer, attribute).

Order, and why each step is where it is:

1. an event on the compute stream marks the step's writes to the row
   complete — the spill stream may not read the row before them;
2. the spill stream waits on it;
3. the copies are issued there, bracketed by a timing pair on the same
   stream, so the compute stream carries none of the transfer and the
   pair measures only the copies (the wait in 2 precedes its start);
4. the returned event fences those copies, both for the entry's
   consumers and for :meth:`fence_before_slab_write`.

Order the compute stream behind the in-flight spill. Call before any slab write.

Returns ``None`` when no spill is outstanding, ``True`` when the spill
had already completed (the wait is free), and ``False`` when it was
still in flight (the compute stream may stall on it). The return value
is the measurement: it says whether the host work between the snapshot
and the next slab write covered the transfer.

The wait is enqueued either way, so correctness never depends on the
query — one code path, and the query is purely observational. On the
``False`` branch the wait is bracketed by a timing pair on the compute
stream: the elapsed time between them is exactly how long that stream
sat on the fence, since nothing else is enqueued between the two. That
is the residual the ``not_ready`` counter can only count.

Read every retired timing pair into its histogram; keep the rest.

``elapsed_time`` requires both events complete, and forcing that would
synchronize the step on a copy it does not otherwise wait for. So a
pair is read only once its end event queries complete, and a pair that
has not stays for the next drain. Nothing here waits.

Shared primitives for the multi-state-kind pool and its concern mixins.

Holds the pieces every :class:`MultiStatePool` concern mixin needs in
common — the pool key alias, the per-state-kind view Protocol, the
recurrent-kind tuple, and the lazy ``@torch._dynamo.disable`` wrapper —
so the mixins import them from one place instead of from each other.

``@torch._dynamo.disable`` without the ~1.5 s import at module load.

Reading ``torch._dynamo.disable`` at *decoration* time forces the whole
``torch._dynamo`` package (~1.5 s) to import on the serving-prep critical
path — paid on every boot, including warm ones, just to obtain a
defensive decorator. The decorated method is only ever called eagerly
from ``__init__`` and never from a compiled region, so we defer wrapping
to first call: by then ``torch.compile`` has pulled ``torch._dynamo`` in
if it is needed at all. The graph-break guarantee is preserved — the real
``disable`` wraps the body before it runs.

True when ``pool`` owns per-request recurrent slab rows.

The predicate behind every "is this a hybrid model" gate: whether a
prefix credited without a matching recurrent savepoint would decode
against zero state. It lives here, once, because the sites that ask
are on opposite sides of the engine — the rank-0 scheduler's
admission (``Scheduler._has_recurrent_state``) and the SPMD worker's
delta replay (``SpmdRankLoop``) — and a gate that holds on one of
them and not the other is the whole shape of the defect it guards.

``has_recurrent_pools`` reports MAMBA / GDN; a ShortConv-only model
(LFM2) answers False there and still carries per-request rows, so the
view registry is the second source. Answers False for a pool exposing
neither (mock pools in CPU tests) and errs toward True when the two
sources disagree — a spurious True costs a re-prefill, a spurious
False is a wrong answer.

What every per-state-kind pool exposes for the engine.

``layer_view(i)`` returns whatever opaque per-layer object the
matching backend's :class:`AttnOp` expects (uint8 slab for tkv,
bf16 pair for FA, dense state pair for Mamba, MLA slot for MLA).

Per-state-kind view construction + per-layer resolution + kind routing.

The :class:`StateViewRegistryMixin` owns how a :class:`MultiStatePool`
turns its ``layer_specs`` into one :class:`StatePoolView` per
``(StateKind, backend_hint)`` key, how a global ``layer_idx`` resolves
to the right per-layer slab view, and how paged-kind queries route to
the dominant paged pool. The pool facade composes this mixin; the
mixin reads the shared ``self._views`` registry and the build params
set by the facade's ``__init__``.

Eagerly resolve the per-layer slab view for ``layer_idx``.

Pulled out of :meth:`layer_view` so the resolution happens once
at construction time (populating :attr:`_per_layer_views`) and
the steady-state :meth:`layer_view` is a flat list lookup.
Honours :class:`LayerSpec`'s ``kv_source_layer`` aliasing — see
:meth:`layer_view` for the contract.

Layers with state_kind NONE return ``None`` (no slab to view —
MLP-only blocks).

Decorated with ``@torch._dynamo.disable`` belt-and-braces: this
method is only called from :meth:`__init__` (eager,
construction time) and never reaches a compiled hot path, but
the decorator hard-pins that contract so a future caller doing
``pool._resolve_layer_view(...)`` from inside a Dynamo trace
graph-breaks loudly instead of silently re-introducing the
per-``layer_idx`` specialization the
:attr:`_per_layer_views` cache eliminates. The two
``self.layer_specs[layer_idx]`` reads inside this body are
therefore safe — no compiled code path can ever observe them.

Construct the per-kind pool view for ``backend``.

``backend`` may be ``None`` for state kinds that do not run
through the :class:`AttentionBackend` / :class:`AttnOp` seam
(e.g. ``StateKind.SHORT_CONV`` — the LFM2 ShortConv kernel is
invoked directly inside the model block).

Resolve per-PAGED_KV-layer geometry + global indices.

Returns ``(geometries, indices)`` where:

* ``geometries[p]`` is ``(num_kv_heads_per_rank, head_dim)`` for
  the p-th PAGED_KV layer.
* ``indices[p]`` is the GLOBAL position of that layer in
  ``layer_specs`` (i.e. its ``layer_idx``).

Both lists have length equal to the number of PAGED_KV layers —
NOT the model's total layer count. Hybrid models (Qwen 3.5/3.6
with GDN linear-attention layers, NemotronH with Mamba and
MLP-only blocks) interleave non-paged layers;
the paged-KV slab is sized to the paged subset only, and
:class:`PagedKVStatePool` maps a global ``layer_idx`` back to
the slab's paged position via ``_global_to_paged``.

Per-rank ``num_kv_heads`` is sliced by ``parallel_cfg.tp_size``
before returning. Hetero geometry across PAGED_KV layers is
supported (Gemma 4 sliding vs full attention) — the caller
decides whether to fold to a single stacked slab (homogeneous
fast path) or keep one slab per layer.

Return the per-layer view for whatever state kind layer ``layer_idx`` has.

Honours :class:`LayerSpec`'s ``kv_source_layer`` aliasing —
when a layer's spec declares an earlier source layer (Gemma 4
``num_kv_shared_layers`` mechanism), the returned view is the
SOURCE layer's view, not this layer's own. The shared layer
thus reads/writes the source layer's slab.

Implementation note. The aliasing / spec-walk lives in
:meth:`_resolve_layer_view`, called once per layer at
construction time and cached on :attr:`_per_layer_views`. This
method is then a flat list lookup, so a Dynamo trace that
reaches it (e.g. an eager capture-time call) emits no
per-``layer_idx`` guard on ``self.layer_specs[layer_idx]``. The
production hot path does not route through this method: the
orchestrator (:meth:`LayerStack.forward_positional`) reads
:attr:`_per_layer_views` and passes the per-layer view as the
block's ``state_view`` argument directly, so the block's
compiled forward sees an opaque per-layer view (one Dynamo
cache slot for every layer instance instead of N).

Return the eagerly-resolved per-layer view list.

``len(out) == len(self.layer_specs)``; ``out[i]`` is whatever
opaque object the matching backend's :class:`AttnOp` expects
for layer ``i`` (uint8 slab for tkv, bf16 pair for FA, dense
state pair for Mamba, MLA slot for MLA), or ``None`` for
:class:`StateKind.NONE` layers (NemotronH MLP-only blocks).

Cross-layer aliasing (Gemma 4 shared-KV) is already resolved —
a shared layer's entry IS the source layer's view, not its
own.

Used by :class:`LayerStack` to thread per-layer views through
``block_args[2]`` so the block's compiled forward sees an
opaque per-layer slab handle instead of the
:class:`MultiStatePool` itself (which would force a Dynamo
recompile per layer index — see the :meth:`layer_view`
implementation note).

Page-count mirror of a PEER attention-DP rank's paged pool.

Under ``attn_dp_size > 1`` each attention-DP set owns whole requests and
its own KV, so admission has to be decided per set: rank 0 runs one
:class:`~arbi_serve.scheduler.scheduler.Scheduler` per attention-DP rank
and each needs a page table whose free count tracks the pool that set
actually holds. Only rank 0's OWN set is backed by real device memory;
the peers' pools live on other processes.

This pool stands in for one of those. It holds no tensors and no device
memory — only a free list of page ids of the same length the peer's real
pool has — because rank 0 never reads or writes a peer's KV: the slot
mapping and block table a peer forwards on are derived by that peer,
against its own real table, from the same broadcast delta.

WHAT MUST MATCH IS THE COUNT, NOT THE IDS. The peer's table and this one
are fed the identical admit / free / ``allocate_slots`` sequence (one
:class:`~arbi_serve.distributed.spmd.SpmdRankLoop` each, driven off the
same :class:`~arbi_serve.distributed.spmd.SlateDelta`), so their free
counts stay equal step for step, which is exactly what
``can_allocate`` / ``available_pages`` read. Handing out different ids
than the peer's real allocator is harmless and never crosses the wire.

Exposes only the paged surface :class:`~arbi_serve.cache.pagetable.FlatPageTable`
and the scheduler's page accounting touch. Everything a recurrent /
offload / savepoint path reaches for is deliberately ABSENT so those
paths refuse rather than silently operating on a pool with no memory
behind it; the attention-DP boot gate refuses those features up front.

DeepSeek-V4 per-layer state: VA-reserved streams that grow in place.

A DSv4 attention layer holds three things per request slot:

  * a **window ring** of the last ``window`` latents, addressed by
    ``position % window`` — fixed size, always resident;
  * a **compressed stream**, one latent per ``compress_ratio`` tokens —
    grows with the context at 1/ratio the token rate;
  * on the fine family, the **indexer's own stream** at the same rate but a
    narrower latent.

Three growth rates in one layer is what forces a paged engine into
per-cache-group page accounting. This pool takes the other route the cuMem
allocator makes available: reserve the VA for the whole 1M-token span up
front and physically back only what live requests have reached
(:class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion`). What that
buys, concretely:

  * **Addressing is arithmetic.** Entry ``e`` of slot ``s`` lives at flat
    slot ``(s // g) * rows * g + e * g + (s % g)``
    (:class:`~arbi_serve.models._deepseek_v4_kv.DSv4SlabAddress`). There is
    no block table and no page indirection, so the id lists the model
    builds are already physical addresses.
  * **Rates stop mattering.** Each store maps at its own pace; nothing has
    to reconcile them.
  * **Growth keeps the address.** The slab's ``data_ptr`` is fixed at
    reservation, so growing mid-run cannot invalidate a captured graph.

The layout is GROUP-MAJOR over slots, ENTRY-MAJOR inside a group, because a
cuMem mapping backs a contiguous byte range. A group of ``g`` slots is one
such range, so mapping ``E`` entries of a group backs row ``E`` for that
group's slots and nothing else: the mapped footprint tracks the longest
live request IN EACH GROUP rather than the longest in the batch.

``g`` is derived, not chosen: the smallest group whose one mapping granule
fits inside what one slot of the compressed stream holds at
:data:`MIN_GROUPED_CONTEXT`, so a group's first granule is fully used. One
slot per group would make that granule span thousands of entries a short
request cannot use; all slots in one group is the whole-batch prefix. Both
ends are reachable — ``g == slots`` when the derivation says the store is
too narrow to split, ``g == 1`` when its rows are wide enough.

The partition is ONE per pool (:func:`group_slots_for`): a request holds
one slot across every layer and every store, so the stores cannot disagree
about which group that slot is in.

Every gather clamps its entry index to the slot's backed row count
(:attr:`DSv4LayerState.kv_backed`), so an id past what a group has mapped
reads a backed row instead of faulting on reserved-but-unbacked VA. The
clamped ids are ones the caller's own mask discards.

Slots cut into groups, each group's entry prefix backed on demand.

``growable=False`` allocates the whole thing eagerly — the same
``growable`` switch :class:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool`
takes, and what a CPU test or a context small enough not to need
on-demand backing wants. The addressing is the same either way.

One layer's addressable state, plus the accumulators that feed it.

The model addresses the store by VIRTUAL id: ``[0, window)`` is the ring
and ``[window, window + stream_capacity)`` the compressed stream. An id
plus a slot is a flat slot index through :attr:`kv_addr` — the whole
point of the VA-reserved layout — so a gather needs no translation
table.

The ``StateKind.DSV4_SPARSE`` pool: one store set per DSv4 layer.

Rows are per-request slots, handed out by the shared row allocator (row
0 is its reserved sentinel), and every layer's stores share the row
index — so a request's window, streams and accumulators all live at the
same slot across the stack.

Recycling a row must reset the ACCUMULATORS. The ring and the streams
are safe to leave: an id list only names entries the request's own
position count makes visible, so a fresh request cannot address the
previous tenant's rows. The accumulators have no such guard — they are
read as the partial group of whatever request holds the row, so a stale
one would pool a previous request's tokens into this request's first
cached entry.

Growth contract
---------------

A growable pool backs only the prefix its requests have reached, so
something must extend that prefix before a forward addresses past it.
Two calls, in this order:

  * :meth:`request_context` at ADMISSION — records, on the host alone,
    the furthest context the newly admitted request can reach, against
    the request's own slot. It maps nothing: the mapping and the zero of
    what it backs are DEVICE work, and admission runs on the engine loop
    while a forward may be in flight on the executor thread.
  * :meth:`flush_pending_zero_clears` at START-OF-STEP — applies the
    recorded targets, on the forward-executor thread, before the step's
    metadata build and forward. This is the same drain the recurrent
    pools' admission zero-clears ride, and it sits outside every
    captured region, which is what makes the mapping legal.

A slot's target extends the GROUP that holds it, so a batch of one long
request and many short ones backs the long one's group deeply and leaves
the rest at their own depth.

:meth:`assert_context_mapped` is the backstop: the metadata builder
refuses a step whose sequences have run past their slots' mapping
instead of letting the forward fault on unbacked memory.

The slot partition every store of every layer is laid out on.

ONE partition, because a request's slot is one slot pool-wide: splitting
it per store would put a request in two different groups. It is derived
from the store whose slot holds the MOST at
:data:`MIN_GROUPED_CONTEXT` — the compressed stream of the fine family,
which is also the store that dominates the footprint — so the
granularity is free where it is worth the most. A narrower store on the
same partition pays at most one extra granule per group, against a
length saving of a whole store's depth.

Back ``group`` to ``rows``, zeroing whatever became mapped.

Newly mapped physical is undefined, and the model reads a row only
after writing it — except the ``-1`` pad id, which clamps to row 0.
Zeroing on map keeps a padded read deterministic instead of
returning whatever the driver handed back. A granule that straddles
two groups is mapped by whichever reaches it first and zeroed only
then, so the group already holding rows in it keeps them.

Back the rows a request at ``context_len`` addresses, every slot.

Mapping only: :meth:`publish_backed` is what makes the new depth
visible to the gathers, so a caller extending several slots pays one
publish for all of them.

Allocate the pre-block frames and the retained per-token inputs.

Every layer keeps a copy of the ring rows a verify block
overwrites: the ring is addressed MODULO the window, so a drafted
position ``y`` takes the row of the still-live position
``y - window``, and no later step ever rewrites it. Undoing the
rejected tail needs what was there before.

The accumulator frames and the hidden states that advance them are
allocated only where a compressor exists.

Idempotent at a same-or-smaller depth; a larger depth reallocates.

No-op: no DSv4 read reaches a row its owner did not write.

Window ids, compressed-stream ids and indexer entries are all
derived from the reading request's own position count, and the
``-1`` pad clamps to a backed zero row that the kernel masks off
before it scores. Boot capture-sweep writes are therefore the same
class as a previous tenant's rows — the class the accumulator reset
in :meth:`_zero_pending_rows` already covers — and slot 0 is a
discard sink this pool writes on purpose (see
:meth:`rollback_batch`), not a zero sentinel anything reads.

A DSv4 read that ever becomes position-independent invalidates this
and must clear the rows here instead.

Record the context an admitted request can reach. Maps nothing.

The reach is recorded against the request's own SLOT, so it extends
that slot's group and no other. ``None`` means the request carries
no output budget and so runs to the reservation. The target is
applied by the next :meth:`flush_pending_zero_clears` — see the
class docstring's growth contract for why the map cannot happen
here.

A request that could run past the reservation is refused: the stores
reserve exactly ``max_context`` entries, so there is no row for it to
address and serving it would fault the device.

Back the rows EVERY slot needs for a request of ``context_len``.

Monotone and idempotent: the mapping only ever extends, and a target
at or below what is already backed does no device work. Issues cuMem
mappings, so it MUST run outside any captured region.

Back the rows ``slot`` needs for a request of ``context_len``.

Extends the group that holds ``slot``, so the slots sharing that
group are backed to the same depth and every other group is left
where it is. The record is per slot and holds what THAT slot
declared, which is a lower bound on what its group backs.

Map ``slot``'s group up to ``context_len``. True iff it moved.

Mapping only — the caller publishes, so a flush that extends several
slots pays one publish per layer rather than one per slot.

Extend the mapping to what admission asked for, then clear rows.

Called at start-of-step, before the step's metadata build and
forward, which is what puts every mapping outside the captured
region and on the thread that runs the forward.

Refuse a step whose sequences run past their slots' mapping.

With per-row ``rows`` and ``seq_lens`` — host ints the scheduler
already holds, so no device read — each row is checked against ITS
slot. Without them the batch's longest sequence is checked against
the deepest mapping any slot has, which is the weakest statement
that never refuses a step it cannot prove wrong.

The step's own tokens extend past a committed length by at most one
row per ``compress_ratio`` of them, so this trails the true reach by
one step; that is the point at which an admission target that failed
to cover a request becomes a loud error instead of a fault on
unbacked memory.

Replay each row's accepted prefix over its pre-block accumulators.

Per layer: restore the accumulators from the pre-block frame, then
pool the retained inputs again with every token past ``n_accepted``
masked off — the same :meth:`DSv4Compressor.advance_batched` the
forward runs, so there is one pooling implementation and the replay
cannot drift from it.

``k_uniform`` routes a fully-accepted row to the zero sentinel
row: it has no rejected tail to put back and its replay would land
the values the forward already left, so its work is a provable
no-op. ``no_host_sync`` skips the host-side row range check, for
the device-resident caller that must not sync.

MLA shared-KV pool — DeepSeek V2 / V3, Kimi K2.

Paged shared-KV slots. Each token in each layer occupies one slot of
``slot_bytes = packed_nope + 2 (norm) + qk_rope_head_dim*2 (bf16 RoPE)``;
the packed NoPE bytes come from
:func:`tkv.runtime.mla.spec.mla_slot_bytes`. One ``(num_mla_layers,
num_pages, page_size, slot_bytes)`` uint8 slab is allocated up front; the
per-layer view is a ``(num_pages, page_size, slot_bytes)`` slice.

Slab spans the MLA layers only; ``layer_view`` takes the GLOBAL layer
index and raises on a non-MLA one.

Same free-list machinery as :class:`PagedKVStatePool` — both share the
:class:`~arbi_serve.cache._pool_base.PagedFreeListPool` base, so only
the slot byte layout + slab construction differ here. Page 0 is
reserved as the "null" page so int32 block tables can use 0 to mean
"unused".

The :class:`MlaBackend` declares the slot-bytes size; the pool calls
``backend.kv_cache_shape(...)`` to get the per-layer tensor shape and
``backend.bytes_per_token(...)`` for per-token accounting.

TP is irrelevant for MLA shared-KV: the stored vector is shared across
heads (one per-token vector per layer), so there is no per-head split
to slice. ``parallel_cfg`` is accepted for API symmetry but every TP
rank holds the full MLA cache.

Paged MLA shared-KV pool. One uint8 slab covering every MLA layer.

Args:
    layer_specs: full model layer specs; the pool indexes by
        ``layer_idx`` via :meth:`layer_view`.
    backend: the active :class:`MlaBackend` (drives slot-byte size).
    model_dims: model-global facts.
    num_pages: total page count. Page 0 is reserved as a null sentinel.
    page_size: tokens per page.
    device: CUDA device.
    parallel_cfg: TP / EP topology (unused for MLA — the cache is
        shared across heads — but accepted for API symmetry).

Multi-state-kind pool — the engine's single handle to all per-layer state.

Indexed by :class:`StateKind`. Each key resolves to one
:class:`StatePoolView` impl:

  - :class:`PagedKVStatePool` for ``StateKind.PAGED_KV``.
  - :class:`MLAStatePool` for ``StateKind.MLA_SHARED``.
  - :class:`RecurrentStatePool` for ``StateKind.MAMBA`` /
    ``StateKind.GDN``.

The engine calls ``cache.layer_view(i)`` and gets the right opaque
view for whatever state kind layer ``i`` carries — the model is
unaware of which backend / slab it is hitting.

:class:`MultiStatePool` is a thin facade that constructs the per-kind
views and holds the shared ``_views`` registry + savepoint-resume
queue; the engine-facing behaviour lives in four concern mixins:

  - :class:`StateViewRegistryMixin` — view construction, per-layer
    resolution, paged-kind routing.
  - :class:`PagedAllocationMixin` — page alloc/free, byte accounting,
    page-table factory, suspended-job page snapshot/restore, draft
    slots.
  - :class:`GrowableKvMixin` — GrowableRegion PAGED_KV grow/sleep/wake.
  - :class:`RecurrentLifecycleMixin` — per-request recurrent state
    alloc/free, savepoints, MTP partial-accept rollback.

Engine-facing pool, keyed by :class:`StateKind`.

Constructed from ``layer_specs``, the active per-:class:`StateKind`
backend map, and ``parallel_cfg``. Internally builds one
:class:`StatePoolView` per :class:`StateKind` and indexes each
layer to the right one based on its :class:`LayerSpec`.

The construction / view-resolution, page-allocation, growable-KV,
and recurrent-lifecycle behaviour is provided by the composed
concern mixins; this class owns ``__init__`` and the shared state
those mixins read.

PageTable Protocol — engine-side page allocator contract.

Two impls implement the Protocol structurally:

  - :class:`arbi_serve.cache.pagetable.FlatPageTable` — flat (no prefix
    cache). :meth:`lookup_prefix` always returns ``None``.
  - :class:`arbi_serve.cache.radix_pagetable.RadixPageTable` — radix
    tree (prefix cache). :meth:`lookup_prefix` consults the tree, scoped
    by tenant.

One :class:`PageTable` Protocol;
:class:`MultiStatePool` accepts a page-table-factory per
:class:`StateKind` so MLA / Recurrent pools can opt into prefix cache.

The Protocol intentionally stays narrow — engine-side: "give me
pages, take pages back, tell me if you've seen this prefix for this
tenant, and tell me how many pages exist". The full per-request
bookkeeping (``add_request`` / ``allocate_slots`` / ``commit_full_pages``)
lives on the concrete impls; the engine drives those via the impl.

Per-:class:`StateKind` page allocator with optional prefix cache.

The flat impl (:class:`FlatPageTable`) returns ``None`` from
:meth:`lookup_prefix` always; the radix impl
(:class:`RadixPageTable`) returns the matched page list when the
prompt's prefix is cached for the given tenant.

The Protocol intentionally does NOT expose the radix internals
(commit / refcount / evict) — :class:`MultiStatePool` drives those
via the concrete impl. The Protocol is the engine-side contract:
"give me pages, take pages back, and tell me if you've seen this
prefix for this tenant".

Release ``page_ids`` back into the free pool.

For the radix impl: drops the radix-tree refcount on each
commit-cached page; the page returns to free only when the
refcount hits zero (so concurrent prefix-cache hits keep the
page live).

Look up a prefix-cache hit scoped to ``tenant_id``.

Returns the list of cached page IDs for the longest matching
prefix in the tenant's namespace, or ``None`` for a miss
(including the flat impl, which always returns ``None``).

``tenant_id`` is the radix tree's namespace key — tenant A's
tree never sees tenant B's pages (defeats the timing
side-channel attack documented in
``AGENTS.md::Prefix cache — privacy controls``).

Paged KV state pool — one allocation per PAGED_KV layer.

The :class:`StatePoolView` impl behind :class:`MultiStatePool` for the
``StateKind.PAGED_KV`` slot.

Slab dimensionality is sized to ``len(paged_layer_indices)`` — the
number of layers whose ``StateKind`` is ``PAGED_KV`` — NOT the model's
total layer count. Hybrid models (Qwen 3.5 / 3.6 with GDN linear-
attention layers, NemotronH with Mamba and MLP-only blocks, LFM2 with
ShortConv blocks) interleave non-paged layers; allocating a slab slot
for those would waste several GB on a 27 B model with 64 layers of which only ~17
are full-attention.

The shape and dtype come from the active :class:`AttentionBackend`.
Every paged backend returns a RANK-3 ``(num_pages, page_size,
slot_bytes)`` per-layer slab — a single packed plane:

  - TKV: ``slot_bytes`` is the compressed uint8 KV slot.
  - tkv-bypass: a ``2*num_kv_heads*head_dim`` bf16 slot packing
    ``[K | V]`` into one fused per-slot row.

The per-layer view is the raw ``(num_pages, page_size, slot_bytes)``
tensor, handed to the attn op as its ``state_view``. Pages on dim 0 of
each per-layer slab make a byte-prefix a page-prefix, so the slab is
GrowableRegion prefix-backable (the capture-then-grow boot path).

Two construction paths:

  - **Homogeneous** (default — every PAGED_KV layer shares
    ``num_kv_heads`` + ``head_dim``): build a single stacked slab over
    the paged layers and slice per-layer views off it. Standard Qwen3-
    style models hit this path.
  - **Heterogeneous** (``per_layer_geometry`` provided — one
    ``(num_kv_heads, head_dim)`` per PAGED_KV layer): build a list of
    per-layer slabs sized to each layer's geometry. Gemma 4 hits this
    path because sliding-attention vs full-attention layers carry
    different head_dim (and on the 31B fixture, different num_kv_heads
    too).

The free-page allocator is shared across both paths — only the slab
representation changes. Page 0 is reserved as the "null" page so int32
block tables can use 0 to mean "unused". A single page-id maps to one
slot per paged layer: ``slot = page_id * page_size + offset_in_page``,
then each per-layer slab indexes its flat slot dim by that ``slot``.
There is no cross-layer page sharing — every paged layer's slab has its
own ``num_pages × page_size`` slot space addressed by the same page-id.

Layer-view addressing: callers pass the GLOBAL layer index (the spec
position in the full ``layer_specs`` list). Internally the pool maps
this to the slab's paged-layer position via ``_global_to_paged``;
calling ``layer_view`` with a non-PAGED_KV global index raises.
:class:`MultiStatePool` already routes by ``state_kind`` before
delegating, so this never fires from the engine — only from buggy
direct callers.

Hot-swap: the engine destroys the old pool entirely (drops the slab,
empties the CUDA cache) and constructs a fresh one for the new backend.
Free-list is rebuilt; in-flight requests must already be drained.

TP awareness: the per-rank ``num_kv_heads`` is sliced from the
model-wide value by ``parallel_cfg.tp_size`` at construction time.
For hetero geometry the per-layer slice happens inside
:class:`MultiStatePool` before the per-layer list is handed in.

Bytes ONE page costs in the auxiliary index-key stream.

A sparse-attention layer (Qwen4-Exp QSA) carries an index-key stream
beside its KV: a ``(num_pages, page_size, width)`` slot per token, plus a
``(num_pages, width)`` pooled row per PAGE that the selection reads
instead of re-pooling the page every step. Both are allocated by
:class:`PagedKVStatePool` out of the same page budget as the KV slab, and
neither is part of the backend's ``bytes_per_token`` — the backend
describes the KV slot, not the layer's indexer.

THE page-cost seam. :func:`~arbi_serve.engine.memory_budget.kv_budget.
per_page_bytes_for_paged_kv` — the number every KV sizing path divides
free VRAM by to decide how many pages fit — calls this, and so does the
pool's own per-token accounting, so "what a page costs" cannot mean two
different things on the two sides of the allocation.

``width`` is NOT sliced by ``tp_size``: every rank materializes the whole
index stream (see ``MultiStatePool._make_view``), so its share of a page
GROWS as TP splits the KV heads.

Owns the per-layer paged KV slab(s) + the free-page bookkeeping.

Args:
    backend: the active :class:`AttentionBackend` (drives shape +
        dtype + bytes-per-token).
    model_dims: the model's :class:`ModelDims`. Used for backend
        calls that need the full model context; **NOT** for slab
        sizing (use ``paged_layer_indices`` for that — hybrid
        models have non-paged layers that must not be sized).
    num_pages: total page count. Page 0 is reserved as a null
        sentinel; usable pages are ``[1, num_pages)``.
    page_size: tokens per page.
    paged_layer_indices: GLOBAL layer indices (positions in
        ``layer_specs``) of every layer whose ``state_kind`` is
        ``PAGED_KV``. Determines slab leading dim;
        :meth:`layer_view` maps a global index back to its paged
        position via this list. Must be non-empty and strictly
        increasing.
    num_kv_heads: per-rank KV-head count (already divided by
        ``parallel_cfg.tp_size``). Required when
        ``per_layer_geometry`` is None (homogeneous fast path);
        ignored otherwise — kept for stats / display.
    head_dim: per-head dim. Same homogeneous-only semantics as
        ``num_kv_heads``.
    device: CUDA device.
    parallel_cfg: TP / EP topology (defaults to single-process).
    per_layer_geometry: optional ``list[(num_kv_heads, head_dim)]``,
        one entry per PAGED_KV layer (length must equal
        ``len(paged_layer_indices)``). When supplied, builds per-
        layer slabs sized to each layer's geometry instead of one
        stacked slab (heterogeneous path used by Gemma 4). Per-rank
        ``num_kv_heads`` slicing must be applied by the caller
        before passing the list in.

Physical bytes the slab holds once :meth:`grow_to_pages` has reached
``target_pages`` — the number to size against free VRAM.

A page target is TWO different byte counts wearing one name, and they
are not equal:

  * ``target_pages * per_page_bytes`` — the SERVABLE capacity. What a
    request can store. This is what the capacity log and the servable-
    context arithmetic report.
  * this — the RESIDENT physical. Each layer owns its own
    ``GrowableRegion`` (one region per layer, so a byte-prefix is a
    page-prefix), and :meth:`~arbi_serve.runtime.cumem_allocator.GrowableRegion.map_to`
    rounds its mapped prefix UP to the device allocation granularity.
    The cost is therefore ``Σ_layers ceil(N × page_bytes_layer / gran)``
    — a STEP function of ``N``, not a line through the origin, and it
    differs from the capacity by up to one granule per layer in either
    direction.

Anything sized against a driver free reading must use this one.
Pricing a grow at the capacity number is what made the post-capture
sizing and its ``kv_sizing_late_growth`` detector disagree about how
much physical a grow had just consumed.

Monotone non-decreasing in ``target_pages`` (regions never unmap on the
grow path), so a search over it may bisect. A target at or below the
mapped prefix costs what is already mapped. Returns 0 when the pool is
not growable.

Map physical for pages up to ``target_pages`` at the SAME stable VA
the captured graphs baked, and extend the free list to cover the newly
backed pages.

Captured graphs reference the slab by its stable base VA + page index;
growing physical at offsets beyond the captured prefix is invisible to
the graph's baked pointers but fully readable/writable once mapped — a
serving step that allocates a grown page and a captured decode graph
that gathers it both resolve to the now-backed VA. Returns the new
mapped-pages count. No-op (returns current) when not growable or when
``target_pages`` is at/below the current prefix.

Drop physical pages for every GrowableRegion KV slab behind their
stable VAs (sleep), returning bytes released.

The growable KV slab lives in its OWN cuMem VA reservation, NOT under
the ``kv_pool`` NamedMemPool tag, so the pool-tag sleep path
(:meth:`NamedPoolRegistry.sleep_all`) can't see it. This registers it
into the sleep/wake lifecycle explicitly: each region's physical is
offloaded D→pinned-host (``offload=True``) then unmapped, keeping the
VA so captured graphs' baked pointers stay valid and :meth:`wake_growable_kv`
can remap the IDENTICAL prefix. No-op (0) when not growable.

Remap physical for every slept GrowableRegion KV slab. Mirror of
:meth:`sleep_growable_kv`. No-op (0) when not growable.

``discarded=False`` — the content-preserving wake: every region remaps
its IDENTICAL chunk layout and restores offloaded contents.
``budget_bytes`` must be ``None`` (a content-bearing slab can never be
partially restored).

``discarded=True`` — the discarded-park re-activation: the slab's
content was discarded on park, so there is nothing to "wake" — the
physical backing is RE-PROVISIONED from scratch, sized to the page
count that fits ``budget_bytes`` (capped at the pre-park size;
``None`` = re-provision the full pre-park size, fail loud on a
driver refusal). The pool bookkeeping (``mapped_pages`` + free
list) is rebuilt FROM the re-provisioned size and every
re-provisioned page is zeroed (the boot-zero invariant: page 0
null sentinel + zeros-is-no-initial-state).

Fail-loud: a driver refusal raises out of ``map_range``; after either
path the per-region mapped bytes are VERIFIED against the bookkeeping
and a mismatch raises — this pool never again reports pages the
driver does not back.

Fail loud unless every layer region physically backs the pool's
claimed ``mapped_pages`` prefix.

The one invariant whose silent violation segfaults serving: the
free list hands out page ids up to ``_growable_mapped_pages``, so every
layer's region must back at least that many pages. Bookkeeping that
exceeds the regions' actual mapping is a lie the first touch pays for
with ``cudaErrorIllegalAddress`` — raise here instead.

No-op on a NON-growable pool: it owns no GrowableRegions and no
``_growable_mapped_pages`` claim (its slab is fully physical or
offload/restored via its named pool), so there is nothing to verify.

Return the per-layer KV view (backend-shape).

``layer_idx`` is the GLOBAL spec index (the position in
``layer_specs``). The pool maps it back to the slab's paged-
layer position; calling this with a non-PAGED_KV global index
raises :class:`KeyError` — :class:`MultiStatePool` already
routes by ``state_kind`` before delegating, so the engine never
triggers the raise.

The view is the raw rank-3 ``(num_pages, page_size, slot_bytes)``
per-layer slab — uint8 for TKV, bf16 for the fused tkv-bypass slot.

Return the per-page pooled index keys for one sparse layer.

``(num_pages, width)`` beside the layer's ``index_layer_view``. The
pool allocates and carries it; the sparse layer that writes the
index keys owns its contents and must re-pool a page whenever it
rewrites that page's keys.

Copy whole KV token slots within the slab, all layers at once.

``moves`` is ``[(src_slot, dst_slot), ...]`` in page-table slot
coordinates (``page_id * page_size + offset``). This is the KV
half of a TREE accept's compaction: the accepted path's rows are
scattered across the step's draft slots and the trim that follows
keeps a PREFIX, so the accepted rows have to be moved down into
that prefix first (see
:meth:`FlatPageTable.draft_path_moves`, which owns the
arithmetic and the ordering guarantee).

The moved bytes are finished KV — already RoPE'd at the position
the token actually occupies, since a tree's depth-``d`` node was
written at absolute position ``P + d``. Nothing is recomputed
here and nothing may be: re-applying RoPE would rotate a row that
is already correct.

One slice per move across the LAYER axis, not one per layer: the
stacked slab carries every layer on dim 0, so a move is a single
copy. A tree accepts at most ``depth`` nodes, so this is a
handful of small copies per row per step, and none at all when
the accepted path is already in place (the chain case, and the
common tree case where the drafter's top-1 won).

Return ``[(tag, backing_tensor, page_axis), ...]`` for the layout.

``page_axis`` is the dimension that indexes PAGES in each backing
tensor — this is NOT always dim 0:

  * ``list-uint8`` — one tensor PER LAYER, each shaped
    ``(num_pages, ...)``; pages are dim 0.
  * ``stacked-uint8`` — a single STACKED tensor shaped
    ``(num_paged_layers, num_pages, ...)``; pages are dim 1
    (dim 0 is the layer index).

The suspend/resume gather + scatter MUST index ``page_axis``, not
a hard-coded dim 0 — gathering dim 0 on a stacked slab indexes the
LAYER axis with page ids and trips a device-side index assert (a
page id ≥ ``num_paged_layers`` is out of range on dim 0). Tags are
stable across snapshot/restore for one engine lifetime (the layout
kind is fixed at construction).

Gather the KV content of ``page_ids`` into host tensors.

Returns ``{tag: host_tensor}`` where each host tensor is shape
``(len(page_ids),) + slab.shape[1:]`` — the request's pages, in
the order given. ``async_copy=True`` issues ``non_blocking=True``
D2H copies into pinned host buffers on the current stream (the
caller records an event); ``False`` is a synchronous gather.

Empty ``page_ids`` returns ``{}``.

Scatter host page content back into ``page_ids``.

Inverse of :meth:`snapshot_pages`. ``host_tensors`` is the dict
returned by a prior snapshot; ``page_ids`` is the fresh page
allocation the content lands in (typically a DIFFERENT page set
than the original — the pool re-allocates on resume). The host
rows are written in order, so ``page_ids[j]`` receives the page
that was at snapshot-time index ``j``.

Synchronous H2D ``index_copy_``. The copy is BLOCKING (default
``non_blocking=False``) so the page content is guaranteed
resident before the request rejoins a slate — the resume runs on
the engine loop thread while the model forward runs on a separate
worker-thread CUDA stream, so a blocking copy is the cheap,
correct cross-thread barrier (resume is bounded to a couple jobs
per step). Raises if a tag is missing (snapshot/restore layout
drift — a hard bug).

Total bytes backing the slab(s), summed over the layout kind.

For the per-layer ``list-uint8`` layout, sum every layer slab.
(Growable slabs report their FULL VA span — mapped-physical is via
``growable_mapped_bytes_total``.)

Per-request page table + per-step slot mapping.

Each request owns a list of page IDs; there is no prefix sharing. This is
the FLAT table, one of two: :class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`
implements the same interface with prefix sharing (refcounted nodes and
chain hashes) and is what the engine builds by default — see
``engine/active.py``. Selected by ``ARBI_USE_RADIX_PAGETABLE``.

The page table reaches into the PAGED_KV pool only via the
:class:`MultiStatePool` API (``alloc_pages``, ``free_pages_list``,
``free_pages``); it never holds a direct handle to
:class:`PagedKVStatePool`.

This module also satisfies the :class:`arbi_serve.cache.page_table_protocol.PageTable`
Protocol structurally — exposing ``alloc / free /
lookup_prefix / free_pages / total_pages`` so the engine can drive any
concrete page-table impl uniformly. The flat impl returns ``None`` from
:meth:`lookup_prefix` (no prefix matching).

``[(src_ordinal, dst_ordinal), ...]`` for a scattered accept.

The page-table-independent half of ``draft_path_moves``: which draft
ORDINALS have to move, and the three ways a caller can hand over a
path that is not a path through the tree it claims. Shared by both
page tables so the guards cannot be present on one and missing on the
other — a tree served through the table without them keeps the wrong
tokens' KV, and nothing raises downstream.

Emitted in ascending destination order. Every destination ``d`` is
``<= path[d]`` because a breadth-first tree puts all of depth ``d-1``
before any depth-``d`` node, so a destination can never overwrite a
source a later move still needs and the sequence applies in place.

Flat per-request page table.

Each request has:
  - ``page_ids``: list[int] of page IDs (length grows as the request
    lengthens; one page per ``block_size`` tokens).
  - ``length``: int — number of tokens currently allocated.

The engine asks for ``allocate_slots(req, k)`` to extend a request
by ``k`` tokens, which may grow ``page_ids`` if the last page is
full. Returns a list[int] of slot indices that map to
``(page_id * block_size + offset_within_page)``.

Register a request; return the cached-prefix length (always 0).

``prompt_token_ids`` is accepted for API parity with
:class:`RadixPageTable` (which uses it for prefix matching) and
ignored here — the flat table never aliases pages across
requests, so every prompt token is fresh and the credit is 0.
``tenant_id``, ``cache_enabled`` and ``max_match_len`` are also
accepted for signature parity and ignored — there's no cache to
scope, opt out of, or cap. Returning 0 keeps the call site
uniform across both page-table backends; the scheduler treats it
as ``req.prompt_consumed += 0`` which is a no-op. (It also means
the flat table cannot carry a hybrid savepoint/KV skew: with no
prefix reuse there is nothing to resume from.)

Free a preempted victim's pages and re-register it page-less.

Signature parity with
:meth:`arbi_serve.cache.radix_pagetable.RadixPageTable.readmit_preempted`
so the preempt paths stay page-table-agnostic. The flat table
carries no cache identity (no tenant namespace, no cache opt-out
— it never shares a page), so there is nothing to preserve here;
this is the plain free-and-re-register. ``prompt_token_ids`` /
``max_match_len`` are accepted and ignored for the same reason
:meth:`add_request` accepts them: there is no prefix tree to
re-query, so the returned match is always 0 and the caller's
``prompt_consumed`` credit is a no-op. An unregistered id stays
unregistered, matching the radix impl.

Unregister ``req_id`` WITHOUT freeing its pages; return them.

Suspended-job KV offload uses this on the suspend path: the
caller snapshots the page content (D2H), then calls
:meth:`release_detached_pages` AFTER the snapshot to return the
pages to the pool. The pages are stashed in ``_detached`` until
then so the snapshot reads valid (not-yet-reused) bytes. Returns
``(page_ids, length)``; the request id is removed from the table
so a later :meth:`import_request` can re-register it against a
fresh page allocation.

Re-register ``req_id`` against an existing page allocation.

Inverse of :meth:`detach_request`. Used by the suspended-job
resume path: the caller has ``alloc``-ed ``len(page_ids)`` fresh
pages and scattered the snapshotted KV content into them; this
binds them to the request with the recorded logical ``length``
so decode continues seamlessly. Raises if the id is already
registered (a resume-before-detach bug).

Page count for ``req_id`` — ``len(page_ids(req_id))`` without
materializing the list (the flat table stores a single ``pages``
list, so this is a direct ``len``). Used by the incremental
block-table path in ``_build_batch``.

Pages reachable for a fresh allocation right now.

For the flat table this is the pool's free list — there is no
prefix-cache eviction layer. RadixPageTable overrides to add
the count of evictable cached pages.

No prefix matching on the flat impl — always ``None``.

Accepted for Protocol parity; ``tokens`` and ``tenant_id`` are
intentionally unused. Production deployments that want prefix
cache wire :class:`RadixPageTable` via the
:class:`MultiStatePool` page-table-factory.

Latch the decode-KV-ring base at the request's current length.

Called by the batch builder on a request's DECODE rows (idempotent,
no-op unless :attr:`decode_kv_ring` is set): the first call runs
after the whole prompt is consumed, so the latched base equals the
prompt length — the same value the reference records as
``_prefill_length`` on its first decode step. From then on
:meth:`allocate_slots` appends normally until ``base + W`` and
wraps afterwards (see :attr:`decode_kv_ring`).

A preempt-restart (or suspend-offload resume) drops the latch with
the rest of the request row; re-arming latches the RE-prefilled
length (prompt + generated so far), so the window keeps bounding
decode KV — a benign drift from the reference's original base.

Extend ``req_id`` by ``k`` tokens, returning per-token slot indices.

Slot index = ``page_id * block_size + offset_within_page``.
Raises ``RuntimeError`` if the pool is out of pages — the
scheduler's preempt path uses ``can_allocate`` as the admission
gate so allocation here only fails on a programmer bug.

On a ring-armed request (:meth:`arm_decode_ring`), tokens past
``ring_base + W`` REUSE the ring slots round-robin instead of
growing: the returned slot overwrites the oldest retained decode
KV and ``length`` stays capped at ``ring_base + W`` (which caps
the ``seq_lens`` the attention kernels see). RoPE positions are
sourced from the REQUEST's logical length by the batch builder,
so they keep climbing — exactly the reference ring semantics.

New pages extending ``req_id`` by ``k`` tokens would allocate.

Mirrors :meth:`allocate_slots`'s arithmetic exactly — the
scheduler's per-step page-commit accounting calls this instead of
re-deriving the crossing predicate. Ring-armed requests stop
needing pages once ``length`` hits the ring cap.

New pages an MTP verify step extending ``req_id`` would allocate.

A verify step takes a tail slot AND ``step_k`` draft slots off the
SAME contiguous tail (``allocate_slots(req_id, 1)`` then
``allocate_draft_slots(req_id, step_k)``), and page arithmetic is
additive across a contiguous tail, so the step's demand is exactly
the ``1 + step_k``-token demand. The scheduler reserves THIS for an
MTP row: reserving the decode tail alone under-reserves by the
drafts, and a saturated slate then admits rows whose draft
allocation cannot be satisfied — the whole step collapses to
``K = 0`` and the speculation an operator is paying for silently
stops happening under load, which is the shape of a throughput bug
nothing reports.

``step_k <= 0`` degrades to the plain tail demand. Same contract as
the radix table's, so the scheduler reserves identically whichever
page table it was booted with.

Extend ``req_id`` by ``k`` draft token slots; rollback-on-fail.

Returns the ``k`` slot indices, or ``None`` if the underlying
pool ran out of pages. On ``None``, any pages taken from the
pool during this call are returned and the request's state is
unchanged — the caller can fall back to the K=1 path safely.

Slots are tail-extending: a follow-up ``allocate_slots`` for
the bonus token after verify writes immediately past the last
draft slot. ``finalize_draft_slots`` trims rejected ones.

``[(src_slot, dst_slot), ...]`` to compact a SCATTERED accept.

:meth:`finalize_draft_slots` is a pure suffix trim, which is
exactly right for a chain: a chain accepts draft slots
``0..n-1``, a prefix, so trimming the tail leaves the accepted KV
where it belongs. A TREE's accepted path is scattered — accepting
nodes ``1, 4, 11`` of 14 and trimming to length 3 keeps slots
``0, 1, 2``, which is three tokens the model never committed.
Nothing raises; the next step simply reads plausible KV for a
prefix that does not exist.

``path[d]`` is the draft-slot ORDINAL (0-based within this step's
``k_total`` allocated draft slots) accepted at depth ``d``. The
accepted node at depth ``d`` belongs at ordinal ``d``, so this
returns the moves that put it there, skipping the ordinals
already in place. Call BEFORE the trim — the ordinals are
resolved against the current (pre-trim) length.

RoPE must NOT be re-applied by the caller: siblings share a
position and the node accepted at depth ``d`` was written at
absolute position ``P + d``, which is where it belongs. This is a
move of finished KV, not a recompute.

Moves are emitted in ASCENDING destination order, and every
destination ordinal ``d`` is ``<= path[d]`` because a node at
depth ``d`` is never earlier than the ``d``-th draft slot
(breadth-first ids put all of depth ``d-1`` before it). So a
destination can never overwrite a source that a later move still
needs, and the sequence is safe to apply in place.

Trim trailing ``k_rejected`` draft slots after the verify pass.

Accepted draft slots stay (they become real KV for the next
decode step). Pages that became fully empty after the trim are
returned to the pool so a long-running tail does not pin idle
pages.

``k_accepted + k_rejected`` must equal the

Multi-tier prefix-KV store — VRAM(hot) → RAM(warm) → disk(cold).

The radix prefix cache (:class:`arbi_serve.cache.radix_pagetable.RadixPageTable`)
keeps prompt-prefix KV pages resident in VRAM (the HOT tier) and shares
them across requests of the SAME tenant. Under VRAM pressure the radix
LRU evicts a refcount-0 leaf and today simply DROPS its page — the
prefix is gone and a future request covering it must re-prefill.

This store catches that dropped prefix and demotes it down a tier
instead of losing it:

    VRAM page ──D2H snapshot──▶ RAM (warm)  ──encrypted blob──▶ disk (cold)

and faults it back on a later prefix-cache lookup:

    disk blob ──decrypt+read──▶ RAM ──H2D restore──▶ fresh VRAM page

so a cold prefix survives eviction at the cost of HOST RAM (then disk),
not VRAM. One entry == one radix node == one page of ``page_size``
tokens; integration entries are keyed by the radix node's incremental
tenant-seeded Merkle digest. The tenant scope is preserved verbatim from the radix
tree — tenant A's snapshot can NEVER be restored for tenant B (the hash
namespace is partitioned by ``tenant_id``).

Built on :class:`arbi_serve.cache._lru_byte_budget_store._LruByteBudgetStore`
(host-RAM LRU + byte budget with pinned-host D2H staging, lazy-confirm
event barrier, and an optional disk-spill tier). The prefix-tier
specifics over that base: a tenant-namespaced hash key, the break-even
gate (:meth:`should_tier`), and AES-256-GCM blob encryption with secure
unlink for the disk tier.

Tiers
=====
  * HOST (warm) — ``{tag: cpu_tensor}`` page snapshots, LRU-ordered,
    bounded by ``host_max_bytes``. Cold host entries spill to disk.
  * DISK (cold) — OPT-IN (``offload_dir`` set). LRU + TTL bounded.
    The on-disk KV is ~invertible to the prompt, so blobs are ENCRYPTED
    with an EPHEMERAL per-process key (see below), written ``0600``, and
    SECURELY unlinked (overwrite-then-remove) on evict.

Confidentiality (disk tier)
===========================
A KV snapshot is approximately invertible to the prompt that produced
it, so an on-disk blob is as sensitive as the prompt. The disk tier is
therefore:

  * OPT-IN — no blob ever touches disk unless ``offload_dir`` is set.
  * ENCRYPTED with AES-256-GCM under a per-process key generated from
    ``secrets.token_bytes`` and held ONLY in this process's RAM. The key
    is NEVER persisted. On process exit (or :meth:`crypto_shred`) the key
    is dropped; every blob on disk becomes undecryptable ciphertext —
    crypto-shredding, so a crash/forensic image after restart cannot
    recover the prompts. (A restart consequently CANNOT reuse old blobs;
    they are unlinked on construction.)
  * Written with ``0600`` perms (owner-only) and SECURELY unlinked
    (best-effort overwrite of the ciphertext before ``os.unlink``).

The break-even gate
===================
Restoring a prefix from a tier is NOT always faster than re-prefilling
it. A restore pays a fixed-ish transfer cost (RAM→VRAM H2D, or
disk-read+decrypt+H2D); re-prefill cost grows with prefix length. So
there is a crossover prefix length below which re-prefill WINS and
caching/​restoring is wasted work + wasted RAM/disk.

This store does not itself time prefill — it exposes configurable
``warm_break_even_tokens`` / ``cold_break_even_tokens`` thresholds and a
:meth:`should_tier` gate. The radix integration consults the gate before
demoting (don't snapshot a prefix shorter than the break-even — just let
it drop) and before restoring (don't restore a short prefix — re-prefill
it). The defaults are conservative unvalidated placeholders (see the
docstring on :data:`DEFAULT_WARM_BREAK_EVEN_TOKENS`).

Thread-safety
=============
A single re-entrant lock guards the in-RAM maps + the disk index. Blob
encryption/IO happens OUTSIDE the lock (only bookkeeping is under it).

Content-address a page-aligned prefix, namespaced by tenant.

Returns the hex SHA-256 of ``tenant_id`` plus the token-id sequence.
The tenant is folded into the hash so tenant A's entries occupy a
DISJOINT key space from tenant B's — a lookup for B can never collide
onto A's blob even if the prompts are identical.

Key a page prefix by its already-computed radix Merkle digest.

Radix nodes maintain this digest incrementally, so tier admission and
lookup can be O(number of pages) instead of re-hashing every growing
root-to-page token prefix from scratch.

One demoted prefix page's KV content.

Carries everything the radix integration needs to splice the page
back into the tree bit-identically:

  * ``page_tensors`` — ``{tag: cpu_tensor}`` from
    :meth:`MultiStatePool.snapshot_pages` for the ONE evicted page.
  * ``tenant_id`` / ``token_ids`` — the radix node's tenant scope and
    its full root→leaf token-id sequence (length a multiple of
    ``page_size``). ``token_ids`` lets the restorer re-derive the
    node's edge key + verify the lookup matches.
  * ``num_tokens`` — ``len(token_ids)``; the prefix length the
    break-even gate is measured against.

Warm-RAM + cold-disk LRU store of evicted prefix-KV pages.

Keyed by ``(tenant_id, prefix_hash)``. ``stage`` admits a demoted
page to the warm tier; ``fault_in`` returns it (warm fast-path or
cold decrypt+read); ``drop`` removes from both tiers.

Args:
    host_max_bytes: warm (RAM) tier byte budget; LRU spill to disk.
    disk_max_bytes: cold (disk) tier byte budget; LRU unlink.
    ttl_seconds: cold-tier age cap (0 disables the TTL sweep).
    offload_dir: enables the cold tier when set (OPT-IN). The dir is
        created ``0700``; blobs are ``0600``, encrypted, and securely
        unlinked. ``None`` ⇒ warm-only (host overflow is a hard drop).
    warm_break_even_tokens / cold_break_even_tokens: the crossover
        prefix lengths gating :meth:`should_tier`. The module defaults
        are UNVALIDATED placeholders — see them for what would replace
        them.

Generate the ephemeral per-process AES-256-GCM key (RAM-only).

Fails LOUD when ``cryptography`` is unavailable — the disk tier's
whole premise is that on-disk KV (≈ the prompt) is encrypted, so
we refuse to enable it rather than silently writing plaintext.

Drop the ephemeral key — every disk blob becomes unrecoverable.

After this, no cold-tier blob can be decrypted (fault-in of a
disk entry fails loud / is treated as a miss). Called at teardown
for defence-in-depth; process exit drops the key anyway.

True iff a ``num_tokens`` prefix is long enough to tier/​restore.

``tier`` is ``"warm"`` (RAM) or ``"cold"`` (disk). Returns False
when ``num_tokens`` is below the configured break-even for that
tier — meaning re-prefilling the prefix is faster than restoring
it, so the caller should NOT snapshot it to the tier (on evict)
nor restore it (on lookup); it just re-prefills. Cheap pure check.

Best-effort overwrite-then-remove of an encrypted blob.

The bytes are already AES-GCM ciphertext, but overwriting before
unlink removes the ciphertext from the inode's blocks so a later
forensic image can't even recover the (key-shredded) ciphertext.
Best-effort: a failed overwrite still proceeds to unlink.

Serialize + AES-256-GCM encrypt an entry's host tensors to disk.

Plaintext = ``torch.save`` of the page tensors + token ids + tenant.
Ciphertext = ``nonce || AESGCM(plaintext)``. Written ``0600`` via an
atomic tmp+replace. The tenant + token-count is bound as the GCM AAD
so a blob can't be silently swapped under a different key.

Drop every blob from a PRIOR process run (undecryptable now).

The encryption key is ephemeral per-process, so blobs written by a
previous run cannot be decrypted by this one — keeping them would
leak ciphertext + waste disk. Unlink them (securely) and start the
cold tier empty. Single-threaded (construction); no lock held.

Radix prefix-cache page table for paged-KV state.

Engine-facing surface: ``add_request``, ``remove_request``,
``allocate_slots``, ``page_ids``, ``length``, ``can_allocate``,
``block_table_row``.

Also satisfies the :class:`arbi_serve.cache.page_table_protocol.PageTable`
Protocol: ``alloc / free / lookup_prefix / free_pages /
total_pages`` so the engine can drive any concrete page-table impl
uniformly. ``lookup_prefix`` consults the radix tree (per-tenant
namespace) and returns matched page IDs or ``None`` on miss.

Behaviour:

  - ``add_request(req_id, prompt_token_ids)`` walks a token-keyed radix
    tree to find the longest matching prefix; matched pages are
    aliased into the request's block-table at no allocation cost.
  - ``commit_full_pages(req_id)`` registers the request's completed
    full pages (NOT the partial last page) into the tree so sibling
    requests can hit the cache.
  - ``release(req_id)`` (alias of ``remove_request``) decrements
    refcounts along the request's path; refcount-0 nodes stay in the
    tree as "evictable" until pool pressure reclaims them.
  - LRU eviction kicks in when ``alloc_pages`` returns ``None``: we
    walk the cache evictable-list (refcount==0, leaf-first) and free
    pages back to the pool until the request is satisfied.

Multi-tenant safety: the tree is shared. Since matching uses
TOKEN IDs (not raw strings) and tokens carry no tenant identity, two
tenants who share a system prompt deduplicate their KV pages — that's
the whole point of the cross-batch share. There is no information leak:
a tenant only learns it hit the cache (visible in metrics), not which
peer populated the entry.

Lookup keying — pages are keyed on TOKEN-ID TUPLES of length
``block_size``; matching only happens at page boundaries. A request
whose prompt is 100 tokens with ``block_size=16`` can match up to 6
pages (96 tokens); the trailing 4 tokens are always private.

OTEL — counters / histograms emitted via :mod:`arbi_serve.cache._metrics`:

  - ``arbi_serve.prefix_cache.hit_total{kind}``  — radix matches
    (``kind`` is the radix-cache kind label, currently always
    ``paged_kv``).
  - ``arbi_serve.prefix_cache.tokens_saved{kind}`` — total prefix
    tokens deduplicated by the cache.
  - ``arbi_serve.prefix_cache.eviction_total{reason}`` — evictions per
    cause (``lru_pressure``).

Walk the tree against ``token_ids``, page-aligned.

Returns ``(matched_pages, matched_nodes, match_len_tokens)``.

Page-alignment: only complete ``block_size``-token chunks can
match. A 100-token prompt with bs=16 matches at most 6 pages
(96 tokens); the trailing 4 tokens never participate in
matching even if a parent-of-the-matching-page exists with
them.

``tenant_id`` selects which per-tenant tree root to walk; the
empty string is the global namespace (single-tenant default).
Tenant A's matches never reach tenant B's pages.

``allow_tier_restore`` (single-process / rank-0 only) extends the
hot (in-VRAM) match by faulting subsequent prefix pages back from
the warm/cold tier and splicing them into the tree (skipping their
re-prefill). The SPMD worker path leaves this False — a tier
restore allocates a fresh page (non-derivable across ranks), so it
must stay rank-0-decides / broadcast like the LRU eviction.

Register a request and (if ``prompt_token_ids`` given) try a prefix match.

Callers that omit ``prompt_token_ids`` (e.g. the preempt path,
which re-adds with fresh token state) get plain per-request
allocation — no shared pages, no risk.

Returns ``match_len_tokens`` so the scheduler can credit the
request's prefill consumption in one shot.

``tenant_id`` namespaces the cache lookup. Tenant A's match
never sees tenant B's cached pages — closes the existence
side-channel for multi-tenant deployments. Empty string =
global single-tenant namespace (default).

``cache_enabled=False`` disables BOTH the prefix match (no
cache read) and the later commit (no cache write). The
request's prompt + completion never enter the radix tree —
the privacy-conscious opt-out matching OpenAI / Anthropic's
no-cache request controls.

``max_match_len`` CAPS the match (page-aligned floor) instead of
taking the longest matching prefix. This is the hybrid-model
correctness gate: on a model carrying recurrent state (GDN / Mamba /
ShortConv) the KV a prefix match aliases is only half the resumable
state — the recurrent state must come from a savepoint, and the
savepoint store is a separate bounded LRU that evicts on its own
schedule. FULL match skips the recurrent layers over
``[coverage, match_len)`` — those tokens never pass through GDN at
all. The scheduler therefore caps the match at the savepoint's
coverage so the two caches resume from the SAME token. Pure-
attention models pass ``None`` and are unaffected.

Free ``req_id``'s pages and hand back its cache IDENTITY.

The half of the readmit seam that both re-registration flavours
share (:meth:`readmit_preempted` on rank 0,
:meth:`readmit_preempted_spmd` on an SPMD worker): read the
namespace + cache-write flag off the entry being dropped, then drop
it. ``None`` when the table never registered the id — the caller
must then leave it unregistered rather than resurrect a dead id.

Free a preempted victim's pages and re-register it.

The one seam every preempt/resume path goes through — rank 0's
:meth:`arbi_serve.scheduler.preemption.Preemptor.preempt_for_space`,
its admission-time re-bind, and (via
:meth:`readmit_preempted_spmd`, which shares the same drop) the
SPMD worker's ``SlateDelta.preempt`` replay. It exists so none of
them has to restate the victim's cache IDENTITY: the namespace
it matches/commits under (``tenant_id``) and its radix-cache write
flag (``cache_enabled``) are read back off the entry being
dropped and carried onto the new one. Both are privacy controls —
a re-admit that let them fall back to :meth:`add_request`'s
defaults would commit the victim's pages under the GLOBAL root
(visible to every other tenant's prefix match) and would re-enable
the radix write for a request admitted with ``cache_enabled=False``.
Taking no identity arguments is the point: there is nothing here a
caller can forget to pass.

``prompt_token_ids`` is the victim's REBUILT sequence (its whole
prompt plus what it has generated). Passing it re-queries the radix
tree, so the victim re-aliases the pages its own prefill already
produced instead of recomputing them: the preempt only
``_ref_down``s those pages (they stay in the tree at ref 0,
LRU-evictable), so an admission pass reaching the victim before
anything reclaims them finds them all. Omitting it is the page-less
re-admit the PREEMPT itself takes — the claim belongs to admission,
because a matched page stops being evictable and a request holding
one from ``waiting`` starves the pool
(``Scheduler._rebind_preempted_prefix``).

``max_match_len`` CAPS the re-match. It is where the caller's
page-headroom bound and the hybrid recurrent-savepoint coverage cap
both land — see that method.

Returns the matched prefix length in tokens (0 when no prompt is
passed, on a miss, or when the id was never registered) so the
caller can credit ``prompt_consumed`` in one shot — the same
contract :meth:`add_request` has.

A request the table never registered is left unregistered: it
holds no pages and no identity to carry, and the one way to reach
that state is a same-tick finish that already removed it — see
``tests/test_spmd_preempt_readmit_same_tick.py`` — where
re-registering would resurrect a dead id.

Replay rank 0's preempt re-admit from the BROADCAST match (worker).

The preempt-path twin of :meth:`add_request_spmd`, and it exists for
the same reason: the radix match is the one non-derivable decision
(rank 0's cap depends on its step's page headroom and its savepoint
store, neither of which a worker has), so it rides the
:class:`~arbi_serve.distributed.spmd_delta.PreemptRow` and the worker
binds exactly those pages instead of re-walking. Identity is carried
off the dropped entry by the shared seam, so — as on rank 0 — there
is nothing here a caller can forget to pass.

``cache_enabled`` OVERRIDES that carried flag, and ``None`` (keep it)
is the default for the reason above: the carried value is already
right for every preempt, and a required argument is one a caller can
get wrong. Rank 0 sets it only where its re-bind RESOLVED a different
value — the hybrid force-disable — which the dropped entry cannot
know about.

The request's RESOLVED radix-cache write flag (table entry).

This is the flag the request was (re-)admitted with — it reflects
the savepoint admission's hybrid force-disable (a re-admit with
``cache_enabled=False`` when the radix matched but no savepoint
covers the prefix), which the Request object itself does NOT. The
SPMD :class:`AdmitRow` builder reads it so the worker's commit
behaviour matches rank 0's exactly. Unknown ids report True (the
default admit flag).

The request's cache NAMESPACE (table entry).

Companion to :meth:`cache_enabled_for`: the table — not the
Request — is the authority on the namespace a registered request
matches and commits under, because the scheduler folds the LoRA
adapter into it (``Scheduler._tenant_id_for``) and a re-admit can
change it. The SPMD :class:`AdmitRow` builder reads it so the
worker registers the request under the SAME root rank 0 used.
Unknown ids report the global namespace.

Conservative check: can SOME request extend by ``k`` tokens?

We count both pool-resident free pages AND evictable
radix-cache nodes (refcount==0). The evictor only fires on the
actual ``alloc_pages`` path, but the budget question must
include them or admission decisions become overly pessimistic.

Free pages plus evictable radix-cache nodes.

Used by both ``can_allocate`` (scheduler) and the admission
gate. Counting evictables matters: a request that finishes
leaves its shared cached pages in the radix tree at
``ref_count=0`` — those pages aren't on the pool's free list,
but the LRU evictor reclaims them on the next ``alloc_pages``.
Treating them as "used" makes admission falsely deny under
steady-state cache occupancy.

Walk the radix tree for ``tenant_id``; return page IDs or ``None``.

Returns the list of cached page IDs covering the longest
page-aligned prefix in ``tenant_id``'s tree, or ``None`` on
miss (no matching pages). Updates each matched node's
``last_access`` so LRU tracking reflects the lookup as a hit.

Tenant scoping is hard: tenant A's lookup never reaches tenant
B's tree (the per-tenant root partition).

Decrement ``node.ref_count``, maintaining ``_evictable_n``.

A ``1 -> 0`` transition on an in-tree node adds it to the
evictable set. The caller keeps the site-specific underflow
guard (``ref_count < 0``); on that bug path ``ref_count`` is
negative so the evictable bump correctly does not fire.

Count of radix nodes whose pages can be reclaimed — O(1).

Returns the maintained :attr:`_evictable_n`. A node is evictable
when ``ref_count == 0`` (not leaf-only — an evicted non-leaf's
children are already ``ref==0`` too; refs propagate up via
``shared_nodes``). See :meth:`_evictable_count_slow` for the
authoritative DFS this caches.

Yield every non-root node across all tenant trees (or one).

``tenant_id=None`` walks every tenant; pass a specific tenant
to restrict iteration (used by per-tenant flush). The global
LRU evictor uses the all-tenant view so cache pressure can
reclaim from whichever tenant has the coldest pages.

Extend ``req_id`` by ``k`` tokens; return per-token slot indices.

Timed under ``ARBI_STEP_PHASE_PROBE`` (``radix_allocate_slots_us``),
because "the radix page-table bookkeeping" is one of the named
suspects for the prefill chunk boundary's host time and the pair
marks around the batch build cannot separate it from the flat-vector
walk it sits inside. Disabled, the probe costs one module-global bool
read per call.

Slot index = ``page_id * block_size + offset``. The first
``shared_prefix_len(req_id)`` token positions reuse cached
pages; positions beyond that are filled from private pages,
allocating new ones at page boundaries.

Note: when the request's currently-claimed-shared length plus
``k`` would step into already-shared pages, this method does
NOT re-extend into shared territory — the engine should never
call ``allocate_slots`` for tokens already covered by the
shared prefix. The scheduler/engine is responsible for crediting
those tokens on admission via the value returned by
``add_request``.

Implementation note: instead of a per-token Python loop, all
required private
pages are allocated in a single pass and slot indices are
either built via a tight list-comprehension (k ≤
``_NP_VECTORIZE_K``, the decode case) or a numpy ufunc chain
(k large, the chunk-prefill case). This avoids ~k Python
iterations on the engine main thread.

Drop cached pages — admin / privacy hammer.

``tenant_id=None`` flushes EVERY tenant's tree. Pass a specific
tenant to scope (default-tenant flush is ``tenant_id=""``).
Pages currently held by in-flight requests (``ref_count > 0``)
are NOT freed — they stay live until the request completes;
only the tree's reachability is broken so future matches miss.
Pages at ``ref_count == 0`` (evictable) ARE returned to the
pool's free list immediately.

Returns ``{"freed_pages": N, "tenants_flushed": [...]}`` for
the admin response. Emits
``arbi_serve.prefix_cache.eviction_total{reason="admin_flush"}``
per freed page so the dashboard reflects the action.

Post-step commit: promote completed private pages into the radix tree.

The cross-batch sharing seed — once a request's full pages are known they
enter the tree so later requests can hit them (with commit-time dedup of a
private page onto an existing shared node). Mixed into
:class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`.

Promote any newly-completed private pages into the radix tree.

Called at end-of-prefill (or at any safe point — repeated calls
are idempotent). A "completed" private page is one whose
``block_size`` tokens are all known, i.e. it isn't the
currently-being-filled tail page.

This is the cross-batch sharing seed: once request A's prefill
completes, its prompt's full pages enter the tree so requests
B/C/D arriving later can hit them. At request finish the
scheduler calls this again with ``token_ids = prompt + output``
so decode-completed full pages also land in the tree — that is
the chat-continuation hit path (a follow-up turn whose new
prompt covers the previous prompt + previous reply matches all
of it).

``token_ids`` defaults to the request's ``prompt_token_ids``
(the prefill-only behaviour). Pass an extended list (typically
``prompt_token_ids + output_token_ids``) to also commit
decode-completed pages. The list MUST start with the same
prefix as ``prompt_token_ids`` and only differ by appended
tokens — it's used as the page-key source.

LRU eviction, eviction listeners, path-hash helpers, and admin flush.

The single-process LRU evictor (heap-backed), the savepoint-store
eviction-listener dispatch, the root->node token/phash path helpers, and
the admin/privacy tree flush. Mixed into
:class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`.

The evicted node's root→node token path, concatenated on demand.

Every consumer on the eviction hot path reads only ``len(...)`` — the
savepoint listener to derive the resume depth and its page size, the
tier gate to ask whether the prefix is long enough to keep — while
building the sequence copies one token per token of CONTEXT, per page
evicted. A chunk-prefill step evicts one page per KV page it allocates,
so that copy lands inside the prefill chunk boundary once per page and
its cost grows with the prompt already cached.

This carries the per-page key references from the single parent walk
that produces the digests, so the length is exact (summed from the keys
rather than assumed uniform) and no token is copied unless something
actually reads one. It is a ``Sequence[int]``, equal to the tuple it
would have been, so a consumer that DOES read the tokens is unchanged.

Register ``fn`` to be called on every node eviction.

The callback receives ``(tenant_id, path_token_ids, page_phashes)``
where
``path_token_ids`` is the concatenated token-id tuple from the
tenant's root to the evicted node, length always a multiple of
``block_size``. Used by :class:`RecurrentSavepointStore` to
drop checkpoints under prefixes whose KV pages just went away
— without this hook, savepoints would linger past the radix
node they depend on, and a future tenant inheriting the freed
page could "match" the dead savepoint and decode against
another tenant's recurrent state — a cross-tenant leak.

Listeners are notified in registration order. Exceptions are
swallowed (counted via
``arbi_serve.prefix_cache.eviction_listener_error_total``) so
a buggy listener can't break the eviction critical path.

Idempotent: registering the same callable twice fires it twice.
Use a single registration for the savepoint store; the
scheduler wires this in at construction.

Concatenate the edge-label token-id chunks from root to ``node``.

Walks parent pointers. Length is always
``len(parents_in_path) * block_size`` (every node holds exactly
one ``block_size``-token chunk as its edge key). Root's
``key=()`` is skipped.

Per-page Merkle digests from root to ``node`` (root→leaf order).

Reads the STORED ``phash`` off each node along the parent chain
— no hashing. ``len(result) == len(path_token_ids) //
block_size``; index ``i`` is the digest of the first
``(i + 1) * block_size`` path tokens.

Return token and digest paths with one parent-chain traversal.

Eviction listeners require both sequences.  Constructing them
independently doubled the already depth-dependent work on that
uncommon path.

The walk itself is O(depth) — one key reference and one digest per
page — and the token path is handed back as
:class:`_EvictedTokenPath`, which concatenates only if a consumer
reads a token. See that class for why the difference is the whole
cost at long context.

Per-page prefix digests for ``token_ids``, read off the tree.

Walks ``tenant_id``'s tree page-by-page (read-only: no
``last_access`` bump, no ref changes, no tier restore) and
collects each in-tree node's stored Merkle ``phash`` (truncated
to the savepoint key length). Element ``i`` is the digest of
the ``(i + 1) * block_size``-token prefix — the exact
``SavepointKey.prefix_hash`` the snapshot side stored via
:class:`~arbi_serve.cache.recurrent_savepoint.ChainPrefixHasher`
(same pure function). ZERO hashing happens here — this is what
makes the admission-side savepoint lookup free.

The walk stops at the first uncached page, so the list may be
shorter than the candidate's page count. That is semantically
exact: a savepoint whose radix prefix pages are gone cannot be
used (the admission KV-coverage cross-check requires
``match_len >= num_tokens_covered``), and the radix eviction
listener has already dropped it from the store in lockstep.

Evict the LRU refcount-0 leaf; return freed page_id (or None).

We only evict LEAVES — evicting an internal node would orphan
its children. The chosen leaf is the one with smallest
``last_access`` among evictable leaves.

Yield ``(node, tenant_id)`` for every non-root node.

Walks per-tenant subtrees so the listener can be told which
tenant's eviction it observed. Root nodes (``parent is None``)
are skipped — they have no edge label and are never evicted.

Drop cached pages — admin / privacy hammer.

``tenant_id=None`` flushes EVERY tenant's tree. Pass a specific
tenant to scope (default-tenant flush is ``tenant_id=""``).
Pages currently held by in-flight requests (``ref_count > 0``)
are NOT freed — they stay live until the request completes;
only the tree's reachability is broken so future matches miss.
Pages at ``ref_count == 0`` (evictable) ARE returned to the
pool's free list immediately.

Returns ``{"freed_pages": N, "tenants_flushed": [...]}`` for
the admin response. Emits
``arbi_serve.prefix_cache.eviction_total{reason="admin_flush"}``
per freed page so the dashboard reflects the action.

Suspended-job offload: detach / deferred-release / re-import.

The suspend path unregisters a request but PINS its shared refs + private
pages until the scheduler has taken the page-content snapshot, then
finalises the release; resume re-imports the request against fresh pages.
Mixed into :class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`.

Unregister ``req_id`` for suspended-job offload; return its pages.

Returns ``(all_page_ids, length)`` — every page backing the
request (shared cache pages first, then private), in logical
order, plus the token length.

Contract: this method does NOT release any pages or drop any
refs. It only unregisters the request from
``_reqs`` and stashes the shared-node refs + private pages in
``_detached`` so they stay PINNED while the scheduler takes the
page-content snapshot (D2H). The scheduler MUST call
:meth:`release_detached_pages` once the snapshot is safely on the
host — that drops the shared refs (each refcount-0 node then
stays LRU-evictable) and frees the private pages back to the pool.

Keeping the shared refs across the snapshot is load-bearing: a
shared prefix page that is also read by another RUNNING job must
not be evictable / re-allocatable in the window between detach and
the snapshot read, or the concurrent job's bytes drift. The
request snapshotted the shared pages' content too, so resume
restores the full sequence into fresh private pages regardless of
whether the shared nodes survive afterwards.

Mirrors :meth:`FlatPageTable.detach_request` (which also defers the
page free to the caller) so the scheduler's offload path is
page-table-agnostic.

Finalise a :meth:`detach_request` once its snapshot is on host.

Drops the shared-node refs (a refcount-0 node stays evictable in
the tree) and frees the private pages back to the pool. Idempotent
/ no-op for an unknown ``req_id`` (a FlatPageTable detach, or a
double-release). Called by the scheduler's suspend path AFTER
``snapshot_pages`` has returned the page content — never before,
or a concurrent job sharing the prefix can read reused bytes.

Re-register ``req_id`` against ``page_ids`` as ALL-PRIVATE.

Inverse of :meth:`detach_request` for the offload resume path.
The caller has allocated ``len(page_ids)`` fresh pages and
scattered the snapshotted KV content into them; we bind them as
the request's private pages (no re-sharing — a resumed job
forgoes prefix aliasing, which is correct and simple). The
request can still grow + commit full pages back into the radix
tree on subsequent decode. Raises if already registered.

Speculative-decoding draft-slot lifecycle for the radix page table.

MTP / Eagle draft tokens are appended past ``length`` into the request's
PRIVATE page tail (rollback-on-exhaustion); the verify pass trims the
rejected tail. The radix tree only learns about accepted drafts at the
engine's next ``commit_full_pages``. Mixed into
:class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`.

Extend ``req_id`` by ``k`` draft slots in private pages; rollback-on-fail.

Returns the ``k`` slot indices, or ``None`` on pool exhaustion
(after eviction has been tried). On ``None``, every page taken
during this call is returned and the request's
``private_pages`` / ``length`` are restored.

``[(src_slot, dst_slot), ...]`` to compact a SCATTERED accept.

The radix table's half of the same operation
:meth:`~arbi_serve.cache.pagetable.FlatPageTable.draft_path_moves`
performs, and it exists separately only because the two tables
address a logical index differently: draft slots live past the
shared prefix, in this request's PRIVATE pages. Which ordinals
move — and the three ways a path can fail to be a path through
this tree — is the shared
:func:`~arbi_serve.cache.pagetable.draft_path_ordinal_pairs`, so
the guard set cannot drift between the two.

Call BEFORE the trim: the ordinals resolve against the current
length. RoPE is not re-applied — a depth-``d`` node was written at
absolute position ``P + d``, which is where its destination
ordinal sits, so this is a move of finished KV.

Trim trailing ``k_rejected`` draft slots after the verify pass.

Accepted draft slots stay in private pages and become real KV
for the next decode step. Trailing private pages that became
fully empty are returned to the pool. The radix tree is
untouched — accepted draft tokens are NOT promoted to shared
nodes here; that happens at the engine's next
:meth:`commit_full_pages` once the request's output token list
extends across a new page boundary.

``k_accepted + k_rejected`` must equal the

Enable/disable the SPMD lockstep-eviction guard (see ``__init__``).

Called once per rank when the engine binds this table to the SPMD
driver. With it set, no rank ever runs the rank-asymmetric wall-clock

Register a request from a BROADCAST radix-match result (SPMD worker).

A worker rank has no scheduler and must NOT re-run ``match_prefix``
— its tree is in lockstep (every rank drives the same admit /
``commit_full_pages`` / ``remove_request`` order), but the LRU
match decision itself is the one non-derivable bit, so rank 0's
``prefix_match_len`` + matched ``prefix_page_ids`` ride the
:class:`AdmitRow` wire. This binds EXACTLY those pages as the
request's shared prefix and ref-ups the matching tree nodes so the
worker's eviction accounting / ``remove_request`` stay symmetric
with rank 0's.

The worker re-walks its own (lockstep) tree to RESOLVE the node
objects to ref-up, then asserts the page ids it found match the
broadcast ones — a divergence here is a lockstep break (the
worker's tree drifted from rank 0's), which raises LOUD rather than
silently aliasing the wrong pages.

Matched cache pages backing ``req_id``'s shared prefix.

Rank 0 reads this right after the scheduler's ``add_request`` so
the matched page ids ride the :class:`AdmitRow` to the worker (see
:meth:`add_request_spmd`). Empty on a cache miss.

Replay rank 0's eviction: free EXACTLY ``page_ids`` (SPMD worker).

The wall-clock LRU pick is the one non-derivable eviction bit, so
rank 0's victim ids ride the :class:`SlateDelta` and the worker
frees those exact in-tree leaf pages here — never re-running the
``last_access`` scan. Fed the lockstep admit/commit/evict order the
worker's tree holds the SAME ref_count-0 leaves, so each broadcast
page id resolves to one evictable leaf. A page id that does NOT
resolve to such a leaf is a lockstep break and raises LOUD (the
worker's tree drifted from rank 0's), mirroring
:meth:`add_request_spmd`.

Resolved ONE page at a time (re-scanning after each free): rank 0's
planner evicts iteratively, so evicting a leaf can promote its
parent to a new evictable leaf the NEXT broadcast id then targets —
a stale up-front snapshot would miss it.

New private pages :meth:`allocate_slots`\ (req_id, k) would alloc.

Mirrors that method's page-demand arithmetic EXACTLY (same
``end_off`` / ``n_need`` / ``n_have`` derivation) without mutating
state — the SPMD eviction planner sums this over the slate to know
how many LRU pages rank 0 must reclaim BEFORE deriving, so the plan
rides the delta and never depends on when ``allocate_slots`` runs.

New private pages a ``grow_tokens``-token extension would alloc.

The shared page-demand arithmetic behind :meth:`pages_needed` (a
single ``allocate_slots(req_id, grow_tokens)``) AND the SPMD verify
planner (:meth:`draft_pages_needed`, a tail slot + K draft slots =
``1 + K`` tokens). ``allocate_slots`` and ``allocate_draft_slots``
both extend the request's private tail by page boundaries off the
SAME ``(length, shared_len, private_pages)`` state, so the combined
page demand of any sequence of extensions summing to ``grow_tokens``
equals a single ``grow_tokens`` extension (page arithmetic is
additive across a contiguous tail). No state mutated.

Private pages an MTP verify step allocates for ``req_id``.

A verify step's per-row derivation (``build_verify_plan`` /
``derive_verify_tensors``) allocates a tail slot
(``allocate_slots(req_id, 1)``) AND ``step_k`` draft slots
(``allocate_draft_slots(req_id, step_k)``) — a ``1 + step_k``-token
private-tail extension. :meth:`plan_evictions_spmd` sums THIS (not the
plain :meth:`pages_needed`) on a spec step so the pre-plan reclaims
the DRAFT pages too; without it the per-step ``allocate_draft_slots``
drives its own rank-asymmetric LRU eviction and the trees diverge.
``step_k <= 0`` degrades to the plain tail demand.

Pre-evict for the slate's page demand; return the victim ids.

Rank-0-only. Sums the per-row page demand over the slate, and if that
exceeds the pool's free (non-evictable) pages, reclaims the shortfall
via the LRU evictor UP FRONT — recording each victim. The returned ids
ride ``SlateDelta.evict_pages`` so workers replay the SAME frees
(:meth:`evict_spmd`) before deriving; the subsequent per-step
allocation on both ranks then draws from an identical free-list and
never triggers its own (rank-asymmetric) eviction. Returns ``[]`` when
no eviction is needed.

Demand model — MUST cover what ``derive`` actually allocates on this
step, or the plan under-reclaims and the per-step allocation refuses
under the lockstep guard:

  - plain step (``spec_active=False``): a single
    ``allocate_slots(rid, n)`` per row → :meth:`pages_needed`.
  - MTP verify step (``spec_active=True``, ``step_k`` the broadcast
    uniform K): the tick is MIXED at any real concurrency
    (``_spmd_split_mixed_slate``) — MTP-opted decode rows allocate a
    tail slot + ``step_k`` draft slots (:meth:`draft_pages_needed`),
    while prefill / non-MTP rows run the legacy sub-pass and allocate
    the full ``allocate_slots(rid, n)`` (n up to chunk_prefill). The
    planner cannot see the split (it has only ``max(pages_needed(rid, n),
    draft_pages_needed(rid, step_k))`` — an UPPER BOUND on either
    route. Over-planning is symmetric-safe (the victims ride the
    delta to every rank; a reclaimed-but-unused page just returns to
    the free pool as cache-miss pressure), while ignoring ``n`` would
    under-plan every prefill row on a mixed verify tick (e.g.
    planning ``1 + step_k`` tokens for a row that allocates a much
    larger chunk).

Unknown-rid defense: this planner is the ONE page-table read on the
SPMD rank-0 delta-build path that runs OUTSIDE the per-step
try/except — a slate rid with no ``_reqs`` entry (a finished request
the scheduler transiently re-emitted) must not raise here, or the
whole engine loop dies. Estimate such a rid at its FRESH demand
(``cdiv`` of its token count): over-evicting is symmetric-safe (the
victims ride the delta to every rank), while raising — or
under-planning and letting each rank's per-step allocation evict on
its own — is not. The bogus row itself still refuses cleanly inside
the guarded step.

Refuse-before-mutate gate for one SPMD step's slot allocations.

``slate`` is ``[(rid, grow_tokens)]`` — EXACTLY the extensions the
caller is about to run through ``allocate_slots``. Under the
lockstep-eviction guard the per-step allocation path cannot evict
locally, so if the pool's REAL free list (no evictables — those are
only reclaimable via the rank-0 plan, which already ran) cannot cover
the summed demand, the step must fail BEFORE the first
``allocate_slots`` mutates anything. A mid-slate failure would leave
earlier rows extended (length advanced, private pages taken) while
the step is refused — per-rank-identical TODAY, but any later
asymmetric read of the half-mutated state (rank 0's Request objects
vs the mirror) turns it into tree divergence. Raising HERE keeps a
refused step ATOMIC: every table on every rank is byte-identical to
its pre-step state, so a subsequent broadcast ``evict_spmd`` still
resolves identically on all ranks.

No-op when the lockstep guard is off (single-process: the local LRU
evictor inside ``_alloc_pages_with_eviction`` handles pressure) or
when the demand fits. Rows without a ``_reqs`` entry are skipped —
the caller's own per-row mirror lookup raises the loud lockstep
KeyError for those before any allocation.

DEFERS the rows that do not fit; never fails the ones that do.
Returns the trimmed slate — the caller derives from THAT. Failing
every in-flight request because one could not be placed is the
wrong blast radius: a page shortfall is a scheduling delay for the
rows that did not fit, not an error for the rows that did.

Deterministic, therefore rank-symmetric: every rank sees the same
broadcast free list and the same slate, and drops in the same order
— largest demand first (fewest rows dropped to make the step fit),
ties broken by descending rid. A dropped row is untouched: it holds
no new pages, its length is not advanced, and it stays in the
scheduler's queues to be re-slated next step.

Rank-comparable state digest for the SPMD lockstep audit.

Returns ``(free_list_hash, tree_hash, free_count, node_count)`` —
four int62 components a debug-mode driver all-reduces across ranks
per tick to pinpoint the FIRST tick two ranks' page state diverges
(and which component: the pool free-list ORDER vs the tree's
chunk→page bindings). Deliberately EXCLUDES ref counts and
``last_access`` — both are rank-asymmetric BY DESIGN (rank 0 refs
at request arrival, the worker at slate admission; stamps are
wall-clock) and harmless. Debug-only cost: O(nodes + free pages)
Python per call.

Multi-tier (VRAM <-> RAM <-> disk) prefix-page demotion / restore.

An LRU-evicted refcount-0 leaf that passes the break-even gate is
snapshotted to the warm tier instead of dropped; a later lookup faults
the page back and splices it into the tree. Mixed into
:class:`~arbi_serve.cache.radix_pagetable.RadixPageTable`.

Snapshot an evicted prefix page → the warm tier (if gated in).

Called from :meth:`_detach_and_free` just BEFORE the page returns
to the pool's free list, so the D2H gather reads the still-valid
KV content. No-op when no tier store is configured. The
break-even gate decides whether the prefix is long enough to be
worth keeping: a prefix below the warm crossover is simply dropped
(re-prefilling it is cheaper than restoring it). Errors are
swallowed (logged) — a tiering failure must never break the LRU
eviction critical path; the page is freed regardless.

Fault page-aligned prefix pages back from the tier into the tree.

Continues a hot (in-VRAM radix) match that stopped at
``start_pages`` pages: for each subsequent page-aligned chunk,
looks the page up in the tier store (warm → cold), and on a hit
allocates a fresh GPU page, H2D-restores the snapshot into it, and
splices a fresh radix node into the tree (refcount left at 0 — the
caller ref-ups, identical to a hot match). Stops at the first miss
or the first chunk whose cumulative prefix is below the break-even
gate. Short pages are restored when they are required ancestors of
an eligible cached prefix. Returns
``(restored_pages, restored_nodes, restored_tokens)``;
empty tuple of lists + 0 when no tier store / no restore.

Byte-identity: the restored page is the SAME KV content that was
snapshotted on eviction (``snapshot_pages`` / ``restore_pages`` are
exact inverses), so a request that hits a restored prefix decodes
identically to one that hit the page while it was still hot.

Recurrent state pool — Mamba-2 / GDN.

Per-layer flat slab tensors of shape ``(max_num_seqs, *state_geom)``
sized at construction. Per-request state is identified by a slab row
index handed out by :meth:`alloc_for_request` and reclaimed by
:meth:`free_for_request`. **Not paged** — each row holds one full
recurrent-state vector for the request that owns the row, freed when
the request finishes.

State layout
------------

For each layer with ``LayerSpec.kind == LayerKind.MAMBA``
(Mamba-2, NemotronH chunk-scan path) we allocate two slabs per pool:

  * ``conv_state``: ``(max_num_seqs, conv_dim_local, conv_kernel)``
    rolling buffer holding the most recent ``conv_kernel`` time-step
    inputs to the depthwise causal conv over the ``[hidden ‖ B ‖ C]``
    stream. Channels-first to match ``causal_conv1d_update``'s
    ``(batch, dim, state_len)`` argument layout — kernel updates the
    slab IN-PLACE on rows selected by ``state_indices``.
  * ``ssm_state``:  ``(max_num_seqs, num_heads_local, head_dim,
    state_dim)`` — multi-head SSM hidden state, matching
    ``mamba_chunk_scan_combined`` final state shape. Updated in-place
    each step.

For each layer with ``LayerSpec.kind == LayerKind.GDN``:

  * ``recurrent_state``: ``(max_num_seqs, num_v_heads_local,
    head_v_dim, head_k_dim)`` — the matrix our vendored fused-decode
    Triton kernel (:func:`arbi_serve.kernels.fused_sigmoid_gating.    fused_sigmoid_gating_delta_rule_update`) consumes as ``h0`` /
    produces as the final state. The kernel's pointer math walks
    ``i_hv * V * K + o_v[:, None] * K + o_k[None, :]`` — i.e. V is
    the outer head axis and K is the inner. Matching the slab to
    that layout natively lets the decode hot path skip the K↔V
    transpose + ``.contiguous()`` round-trip on every layer (vs. the
    alternative ``(HV, K, V)`` layout); FLA's chunk-prefill kernel
    ``chunk_gated_delta_rule`` consumes ``(HV, K, V)`` so the prefill
    branch transposes K↔V at the slab gather/scatter boundary, which
    is far rarer than decode and a net win on the per-step launch
    budget.
  * ``conv_state``: ``(max_num_seqs, conv_dim_local, conv_kernel)``
    rolling buffer, channels-first to match HF's
    ``causal_conv1d_update`` semantics.

``num_heads_local`` / ``num_v_heads_local`` / ``conv_dim_local`` are
the per-rank sizes after slicing by ``parallel_cfg.tp_size`` — we
mirror the parallel-linear convention (out-features sharded along
dim=0).

Zero-sentinel slab row (defensive hardening)
---------------------------------------------

Slab row 0 is reserved as a permanently-zero "sentinel" row and is
NEVER handed out by :meth:`alloc_for_request`. Slabs are allocated
``(max_num_seqs + 1, *state_geom)`` — usable rows are ``1..max_num_seqs``,
row 0 stays at the boot-time zero allocation for the engine's
lifetime. The free list is initialized with rows ``1..N`` (skipping
0) and :meth:`zero_all_rows` zeros every row including row 0 (so the
sentinel survives the boot-capture-sweep wipe — same as any other row).

Why a sentinel row? The actual leak path is closed:
:meth:`Scheduler.add` calls
:meth:`MultiStatePool.alloc_recurrent_state` so each request reserves
a fresh, zero-cleared row before decode reads it. But the per-kind
metadata builders' fallback paths in
:mod:`arbi_serve.backends.sub_builders.state_indices` still default to
``arange(B)`` when ``req_ids`` is missing or ``row_for`` raises
:class:`KeyError`. With B=1 decode that fallback routes to slab row 0
— whatever the previous request just wrote. The production
fallback path RAISES; the safety net is row 0 being permanent zero so
that any future regression that re-introduces a silent fallback hits
"garbage starts from zero" (detectable in tests, non-cross-
contaminating) instead of "garbage is the previous request's state"
(silently wrong output).

The sentinel row costs one row of state PER LAYER, and every copy of
it holds the same thing: zeros. When the host has the cuMem
virtual-memory API, the slabs are therefore allocated out of a
:class:`~arbi_serve.runtime.sentinel_alias_arena.SentinelAliasArena`,
which maps every granule wholly inside a sentinel row onto ONE shared
physical page. The slabs keep their shapes, their row indexing, their
contiguity and their stable ``data_ptr``s — only the physical page a
row-0 address resolves to changes — so no kernel, metadata builder or
captured graph can tell the difference. Aliasing rounds DOWN to the
device granularity: the partial granules at either end of row 0 also
cover live row 1, so they keep private physical. Without the driver
(CPU pools, hosts with no cuMem) the slabs are ordinary
caching-allocator tensors and row 0 is private physical, exactly as
before; ``ARBI_RECURRENT_SENTINEL_ALIAS=0`` forces that layout.

Aliasing makes row 0 SHARED across the layers of one pool, which
changes nothing the sentinel promises. What row 0 must guarantee is
"reads zero, and is never a previous request's state" — and it still
does: the only writer of row 0 is a full-slab ``zero_()``, whose write
through the alias is idempotent (and still required, because the
ragged tail of row 0 beyond the aliased granules is private physical).
The capture sweep's synthetic writes into row 0 land on the shared page
and are cleaned by the same :meth:`zero_all_rows` that cleaned them
before. Across a sleep/resume the arena does better than restore: its
shared pages are re-created and zeroed on wake, so the sentinel is zero
by construction rather than by trusting what was dumped to host.

Capacity bookkeeping: ``max_num_seqs`` is the OPERATOR-facing
usable-rows count (``--max-num-seqs=4`` ⇒ 4 concurrent requests
admissible). The slab is sized ``max_num_seqs + 1`` to cover the
sentinel; ``bytes_total()`` reports the slab footprint less whatever
the sentinel alias removed, because that is the VRAM the pool holds. The profiler-side helper
:func:`smallest_recurrent_slab_capacity` returns slab dim 0 (``N+1``)
so synthetic-batch ``state_indices`` clipping covers the full valid
range; the sentinel row is clean-zero so a clamp into row 0 just
produces a no-op recurrent-state-update.

Layer view contract
-------------------

Each backend's per-layer block reads a single :class:`Mamba2LayerView`
/ :class:`GdnLayerView` instance — its
``conv_state`` / ``ssm_state`` / ``recurrent_state`` attributes are the
flat slab tensors. Kernel access is ``slab[state_indices_tensor[i]]``
where ``state_indices_tensor`` comes from the per-step metadata
builder. Slab rows have stable ``data_ptr`` for the engine's lifetime
— that's the contract that lets the kernel run inside a captured
CUDAGraph (the captured kernel reads the slab at a fixed pointer; per
replay, the engine ``copy_()``s the current step's slot mapping into
the persistent ``state_indices`` buffer the graph references).

Implementation notes
--------------------

The pool keeps a LIFO ``_free_rows`` stack and a ``_row_for_req``
dict. ``alloc_for_request(req_id) -> int`` pops a free row, zeros it
across every layer's slab, and returns the row. ``free_for_request``
returns the row to the free list. Rows are zeroed eagerly on alloc
(rather than on free) so that admission of a new request always sees
a clean slate without depending on a prior tenant's free-time
correctness.

Per-request recurrent-state pool, one per state kind.

Constructed with the model's ``layer_specs``, the kind it serves
(``StateKind.MAMBA`` or ``StateKind.GDN``), ``max_num_seqs`` (the
slab capacity = upper bound on concurrent recurrent rows), and
``parallel_cfg`` so intermediate / state dims are sliced to per-
rank counts.

The pool pre-allocates one slab tensor per (layer, state-kind)
sized for ``max_num_seqs``. Per-request rows are handed out by
:meth:`alloc_for_request` and reclaimed by
:meth:`free_for_request`.

Internal slab dim 0 — usable rows + 1 for the zero sentinel.

Row 0 is the sentinel (always-zero, never handed out). Rows
``1..max_num_seqs`` are user-allocatable. See module docstring
for rationale.

``(view, attr, shape, dtype)`` for every base slab of our kind.

One list, consumed by both the arena path and the plain-tensor
path, so the two can never describe different geometry.

Back ``plan``'s slabs with a sentinel-aliasing cuMem arena.

Returns the committed arena with every slab assigned to its view,
or ``None`` when this host / configuration cannot use one — no
cuMem driver, a non-CUDA device, the alias switched off, or no
slab whose sentinel row reaches the device granularity (aliasing
rounds DOWN, so a row smaller than one granule aliases nothing and
the arena would only add machinery).

FIRST so their offsets stay multiples
of their own uniform size; interleaving the small companion slabs
would push each big slab off the granule grid and pay alignment
padding for what it recovers.

``row_dim`` names the axis the sentinel row sits on; every axis
before it must be length 1 so that row still begins at byte 0 of
the slab (a 1-deep ``snap_X[1, rows, …]`` frame qualifies).

Any failure falls back to plain allocation: the alias is a memory
optimisation, never a correctness dependency.

Allocate one slab per layer of size ``(max_num_seqs + 1, *geom)``.

The leading dim is bumped by 1 to host the zero-sentinel row at
index 0 (see module docstring). Slabs live for the engine's
lifetime and have stable ``data_ptr``s — the contract the
recurrent kernels (and the captured CUDAGraph that wraps them)
rely on.

When a sentinel-aliasing arena can be built, the slabs are views
over its VA and every granule wholly inside a sentinel row
resolves to ONE shared zero page; otherwise they are ordinary
caching-allocator tensors. Either way the shapes, the row
indexing and the boot-zero contents are identical.

Zero ``rows`` across every layer's state attributes (one indexed
write per attribute). See :meth:`PerRequestRowPool.flush_pending_zero_clears`.

Routed through :func:`fill_rows_`, never ``buf[rows] = 0`` -- on the
48-layer hybrid this loop is 96 indexed writes in ONE step, and the
scalar-assignment spelling makes every one of them a device sync.

Zero every slab row across every layer of this kind.

The boot-time cudagraph capture sweep
(``precapture_decode_graphs`` / ``precapture_layer_graphs``)
runs synthetic decode + prefill forwards through
:class:`GDNBlock` / :class:`Mamba2Block` / :class:`ShortConvBlock`,
which write the post-capture recurrent state back into rows
``0..B-1`` of each per-layer slab via ``index_copy_``. The
synthetic content is meaningless but non-zero, and the
per-step ``has_initial_state = seq_lens > 1`` derivation
unconditionally treats every prefill row as "carrying prior
state" — so the first live prefill read picks up the capture-
sweep garbage, the GDN ``h0 * has_init`` mask multiply does
NOT zero it (mask is True), and the kernel's recurrent state
starts off corrupt. A few decode tokens trickle out coherently
before the slab drift triggers a degenerate-token loop.

Eager (``cuda_graphs=False``) skips the capture sweep entirely,
so its slabs stay at the boot-allocated zeros and the same
``has_initial_state = True`` derivation reads correct zeros —
which is why the eager path produces coherent output on the
same prompt while cudagraph replay loops.

Called from :func:`arbi_serve.engine.build.build` after the
capture sweep completes (decode + prefill + piecewise + drafter)
and after each later synthetic forward.

Row 0 — the zero sentinel — is zeroed with every other row. Under
the sentinel alias that write goes through the shared page, which
is idempotent (it writes zeros to a page whose invariant is zeros)
and still necessary: the part of row 0 that no whole granule covers
is private physical like any other row. Skipping row 0 would leave
the capture sweep's synthetic writes in the sentinel for good.

Allocate per-token snapshot buffers sized for ``T <= max_k_plus_1``.

Per GDN layer view, allocate ``snap_X[T, max_num_seqs,
*state_geom]`` for every state attribute the kernel writes per
token. The verify forward writes ``snap_X[t, slab_row]``
post-token-t; partial-accept rollback reads
``snap_X[n_accepted, slab_row]``.

Idempotent: re-calling with a same-or-smaller ``max_k_plus_1``
is a no-op. Re-calling with a larger value reallocates (used today
by the deferred-KV-resize reattach in ``engine.build``). NOTE: this
idempotent-grow is only ONE sub-step of a true runtime serving-depth
raise — the drafter+verify cudagraphs are captured for ``max_k`` at
boot and re-capturing them at runtime on a 27B TP2 engine wedges the
NVIDIA driver, so raising the SERVABLE ``mtp_driver.max_k`` is
reboot-only (admission rejects a per-request K above the boot max_k;
it is never silently clamped). See the ``MtpDriver.max_k`` setter.

MAMBA pools allocate the mirror-image buffers on
:class:`Mamba2LayerView` (base ``snap_{ssm,conv}_state`` + the
per-token ``[x ‖ B ‖ C]`` stream and raw ``dt``); the Mamba-2
recurrence needs no gating pair, so its retained-input footprint
is strictly smaller than GDN's.

Allocate the 1-deep ``snap_*`` base frames collected by the caller.

Each frame is ``(1, *slab.shape)`` of the slab's dtype, so it carries
a copy of the sentinel row too. That row is never read — rollback
only ever addresses rows a live request owns — so the frames share
their own :data:`UNREAD_PAGE`, kept distinct from the slabs' zero
page: a stray write into a snapshot sentinel then cannot reach the
sentinel the pool's safety net depends on.

Allocate the MAMBA rollback buffers — mirror of the GDN branch.

Base frame (leading dim 1) for both modes; then either the
``(N, T, …)`` masked-replay inputs or the ``(T, N, …)`` recompute
inputs, never both. The base frames are appended to ``base_frames``
for :meth:`_attach_base_frames` rather than allocated here.

``"replay"`` or ``"recompute"`` — see ``gdn_mtp_rollback_mode``.

``max_k_plus_1`` is retained for call-site symmetry with
``attach_mtp_snapshot_buffers`` (and the drafter/lifecycle callers
that pass it); the resolution is capture-aware and does not depend
on the snapshot footprint.

Partial-accept rollback — MASKED REPLAY or RECOMPUTE-from-base.

Mode follows the buffers ``attach_mtp_snapshot_buffers`` allocated
(``mtp_snapshot_mode``); both keep only a 1× ``snap_X`` base frame:

* MASKED REPLAY (``replay``, the ``replay_*`` inputs allocated):
  one captured launch set recomputes each partial row's accepted
  prefix ``t <= n_accepted`` from the base frame and commits it
  in-place — see :meth:`_dispatch_masked_replay`. No host-driven
  per-token loop.

* RECOMPUTE (no ``replay_*`` inputs): REPLAY the recurrence
  from the captured pre-verify base over the retained
  accepted-prefix inputs — see :meth:`GDNBlock.rollback_recompute`
  / :meth:`Mamba2Block.rollback_recompute`. Host-driven per-layer
  replay; the eager-decode fallback.

Empty input is a no-op. ``n_accepted`` is bounds-clipped by the
caller (each entry ``<= K``); out-of-range ``slab_row`` entries
silently no-op via a mask.

``no_host_sync`` (the async-hybrid leapfrog). The default path
validates the row range on host (``bool(valid_mask
.all())`` — a D2H sync); the async verify path passes
``no_host_sync=True`` so the reconcile enqueues data-dependent on a
DEVICE ``n_accepted`` with no host pull. Masked-replay staging is
pure device work in both cases (invalid rows route to the
zero-sentinel).

``k_uniform`` is the uniform draft depth ``K``, used to skip
fully-accepted rows wherever the forward already committed their
post-token-K state: the recompute rollback, and Mamba's masked
replay. GDN's masked-replay forward commits nothing, so every row it
advanced has to be replayed.

Identity of every tensor the captured replay graph dereferences.

``data_ptr`` per tensor catches reallocation; the replay depth T
catches an idempotent-grow reattach the allocator happened to
place at the old addresses. Compared on every dispatch — a stale
captured graph holding a freed device address is an IMA, not a
cache miss.

Build the active tree's scan tables now, in the default pool.

The tree verify forward resolves these from inside the compiled
region, so a FIRST resolution landing inside a cudagraph capture
allocates them from that capture's private pool — memory the next
graph captured into the same pool is handed again. Every captured
tree verify then gathers through another graph's activations.
Building them here, off any pool context and ahead of the first
forward, is what makes the address stable for the boot.

Address of the conv-window table the tree replay body gathers through.

Zero when no tree is configured or no GDN view is attached — the
fingerprint only needs it to CHANGE when the tensor does, and a
stable zero on a route that never reads it is exactly that.

MASKED-REPLAY rollback — the ladder-free accepted-prefix commit.

Stage the per-slab-row accepted offsets into the persistent
``_replay_nacc`` buffer (−1 = untouched row), then run ONE captured
launch set that, per layer, (a) recomputes the recurrence from
the 1× base frame with the SAME compiled kernel the forward ran,
steps ``t > nacc`` masked to an exact identity at the inputs
(see ``GDNBlock.replay_masked_commit`` /
``Mamba2Block.replay_masked_commit``), and (b) advances the conv
slab to the accepted window via a pure gather over
[base ‖ saved x_conv].

Zero host tax by construction: per tick the host issues the small
nacc staging ops plus one ``CUDAGraph.replay()`` — the same order
of host work as the ladder's fused accept-commit, unlike the
recompute mode's host-driven per-token per-layer loop this mode
replaces. The first dispatch
(and any dispatch after a buffer reallocation) runs the launch
set eagerly once — correct AND it JIT-warms the kernels — then
captures the identical sequence for every later tick.

``n_accepted`` is the accepted offset (state after committing tokens
``[0..n]``); ``k_uniform`` is accepted for signature parity with the
recompute path and not read here, because the replay-mode forward
commits no state and so every staged row must be replayed. Invalid
rows are routed to the zero-sentinel row, which every replay kernel
skips (``state_idx <= 0``).

All staging is pure device work — no host sync in either the
host-validated or ``no_host_sync`` caller path.

Issue the per-layer masked-replay launches (capture body).

Derives the per-tick shared tensors from the persistent ``nacc``
staging buffer — the inactive-row → zero-sentinel route table, and
either the ``(N, T)`` accepted-prefix step mask or the per-row trip
count that replaces it — then dispatches every layer of OUR kind.
Pure device work off persistent inputs: capture-safe.

Per-layer replay launches for a TREE step (capture body).

Same shape of work as the chain branch — one masked replay per
layer plus a conv-window gather — over ``depth`` steps rather than
the staging buffer's full width: a root-to-leaf path is
``depth`` block rows long however many nodes the tree has, so a
tree's replay is no deeper than the chain's it replaces.

The conv-window table is a pure function of the geometry and the
layer's kernel width, so it comes from the same cached builder the
forward's conv read — one tensor identity per ``(geometry,
conv_kernel, device)``, which is what a captured graph needs.

Physical bytes held by this pool's sentinel-alias arenas.

These live in their OWN VA reservations, not under the pool's
``NamedMemPool`` tag, so no torch or pluggable-allocator counter
sees them — every accounting surface that reports this pool must
add this term explicitly (mirrors the growable KV slab).

Re-map the arenas at the SAME VAs, re-applying every alias.

The shared pages come back freshly zeroed, so the sentinel row is
zero on the far side of a sleep by construction rather than by
trusting whatever was dumped to host.

Release the arenas' physical AND their VAs, now.

Teardown only — every slab view into them dangles afterwards. Dropping
the pool releases them anyway (each slab holds its arena alive until the
last reference goes); this is for a caller that must reclaim
deterministically rather than at the next collection.

Total bytes currently held by every per-layer slab.

The slab is sized for ``max_num_seqs`` rows and is always live
(regardless of how many rows the engine has currently handed
out via :meth:`alloc_for_request`). When MTP rollback buffers
are attached they add a 1× base frame per state attribute plus
the ``(K_max + 1)``-deep per-token INPUT buffers of whichever
mode is live. The number reflects the engine's worst-case
footprint for this kind, which is what the scheduler /
VRAM-watermark surface needs to know.

Recurrent-state savepoint store — the host-RAM checkpoint cache.

Provides recurrent-state checkpointing for prefix
cache on hybrid GDN / Mamba / ShortConv models. The radix prefix
cache (``RadixPageTable``) keys on token IDs and aliases paged-KV
pages — that's enough for pure-attention models, where the only
per-request state is the KV cache. Hybrid models also carry per-
request *recurrent* state (GDN's ``recurrent_state`` matrix +
``conv_state`` rolling window; Mamba's ``ssm_state`` + ``conv_state``;
LFM2 ShortConv's ``conv_state``) that the radix tree can't capture
because the recurrence is non-trivially a function of the token
sequence, not the multi-set of tokens. Without checkpointing the
recurrent state along the prefill, a prefix-cache hit would skip
prefill and decode against zero recurrent state — the
"France is France is France..." bug.

This module is the checkpoint side: a content-addressed savepoint
store that holds post-prefill recurrent state in host RAM, indexed by
``(tenant_id, prefix_hash, num_tokens)``. On admission the scheduler
reads the candidate prompt's per-page digests off the radix tree,
looks up the longest stored prefix they agree with, and queues a
host→device restore of the recurrent state into the request's
freshly-allocated slab row. The restore mirrors the cudagraph-safe
pinned-host ring-buffer pattern that ``flush_pending_zero_clears``
already uses.

This file holds the CPU-side store + unit-testable surface. The
device-side write path (``MultiStatePool.flush_pending_savepoint_resumes``)
and the prefill-commit hook (``SavepointSnapshotter.snapshot_at_boundary``)
follow the captured-graph data_ptr stability and pinned-host staging
invariants of the existing ring-buffer pattern.

Architecture
============

``SavepointKey``
----------------

Three components:

  * ``tenant_id: str`` — mirrors the radix tree's per-tenant
    namespace (closes the multi-tenant existence side-channel; see
    ``RadixPageTable._tenant_roots``). Empty string = global namespace
    (single-tenant deployments retain free cross-request sharing).
  * ``prefix_hash: bytes`` — 16-byte digest of ``token_ids[:num_tokens]``.
    On the radix path it is the tree's per-page Merkle digest at that
    depth, truncated; the flat-table fallback is :func:`hash_token_prefix`.
  * ``num_tokens: int`` — how many leading prompt tokens the snapshot's
    state has absorbed. It is the resume point, and it is always a
    multiple of the KV ``block_size`` (below).

THE RESUME POINT IS PAGE-ALIGNED, ALWAYS
-----------------------------------------

A resume credits ``prompt_consumed = num_tokens``: the recurrent slab
row is restored to the state after exactly that many tokens, and the
radix tree aliases the KV pages for exactly that many tokens. Both
halves must resume from the same token, and the KV half can only be
aliased in whole pages — a page is published to the tree only once
full, so a position inside a page belongs to one writer and is never
shareable. A snapshot whose coverage is not a multiple of ``block_size``
therefore has no KV to pair with and could never be restored honestly:
``put`` refuses it. Nothing rounds: the slab state cannot be rewound
to an earlier position (the recurrence for the tail was swept inside
one kernel and no intermediate state exists), and the tail's KV is not
the resumer's to write.

``SavepointEntry``
-------------------

Carries the per-layer per-attribute host tensors that constitute the
recurrent state for ONE slab row after ``num_tokens_covered`` tokens.
The shape is backend-specific: GDN snapshots both ``recurrent_state``
and ``conv_state``; Mamba1/2 snapshots ``ssm_state`` and
``conv_state``; ShortConv snapshots ``conv_state`` only. The store
treats them as opaque ``dict[(layer_idx, attr), torch.Tensor]``
payloads — the caller (``MultiStatePool.snapshot_recurrent_row``) is
responsible for the shape-correct dict; this module never inspects the
contents.

``RecurrentSavepointStore``
---------------------------

LRU + byte-budget eviction.

  * ``put(key, entry)`` — admit; evict LRU until the post-admit
    footprint fits the budget. Idempotent: re-put under the same key
    bumps recency without re-counting bytes.
  * ``get(key)`` — return the entry (and bump recency) or ``None``.
  * ``best_match_hashes(tenant_id, page_hashes)`` — the radix path:
    given the candidate's per-page digests (read off the tree, no
    hashing), return the key of the LONGEST stored prefix the
    candidate agrees with. ``best_match(tenant_id, token_ids)`` is the
    same selection hashing the candidate itself (flat-table fallback).
  * ``on_radix_evict`` / ``invalidate_prefix`` — eviction-callback
    sinks. Tied to the radix tree's eviction path so a node freed from
    the tree drops the savepoint at that depth atomically (cross-tenant
    leak guard).

Multi-tenant correctness
========================

The same ``prefix_hash`` MUST NOT match across tenants. The store's
hash key is ``(tenant_id, prefix_hash, num_tokens)``; every lookup
filters by ``tenant_id`` first. Mirrors ``RadixPageTable._tenant_roots``
exactly so the two stores stay scope-aligned — the eviction callback
only fires under the tenant whose tree node was evicted, and only
that tenant's savepoints are dropped.

Hard-fail invariants
====================

  * ``cache_enabled=True`` admission with a hybrid model + non-zero
    radix ``match_len`` BUT no savepoint hit ⇒ RAISE in the
    scheduler. The two stores must stay consistent. (This module's
    contribution: the eviction callback drops savepoints in lockstep
    with radix nodes, so a missed lookup post-admission is genuinely
    a bug, not a benign race.)
  * A key or entry whose coverage is not a multiple of the pinned
    ``block_size``, or a key and entry that disagree on it ⇒ RAISE at
    ``put``. Such an entry could only ever be restored wrongly.
  * Budget exhaustion ⇒ caller ("admission gate") MUST honour the
    fallback policy in ``CacheConfig.savepoint_overflow_policy``
    (``"refuse"`` by default — admission falls through with
    cache_enabled=False, paying full prefill; ``"raise"`` for tests).

An operator's explicit byte cap cannot hold one savepoint of this model.

Raised at BOOT, not at the first snapshot, because it is a misconfiguration
with exactly one honest outcome: nothing would ever be retained, every
hybrid prefix-cache hit would re-prefill, and the store would spend its
whole write cost for nothing. Overruling the operator silently is what the
old ``min_entries`` behaviour did and is the defect this replaces; serving
on with a store that cannot hold anything is the other way to be silent.

Refuse a stated cap smaller than ONE of this model's savepoints.

Names all three numbers, because the operator cannot fix it from any two:
what they set, what the model's savepoint actually weighs, and the floor
that used to paper over the gap.

Only for a STATED cap. An unstated one is the shipped default, which is
tuned for a small hybrid and is expected to be too small for a large model
— that is precisely the case ``min_entries`` exists to carry, so it is not
an error. ``slot_bytes <= 0`` means the model holds no recurrent state and
there is nothing to size against.

The grid every savepoint resume point lies on, checked against the kernel.

Two grids meet at a resume point. The KV half is aliased in whole
pages, so the point must be a multiple of ``block_size``. The
recurrent half is bit-identical to a single-sweep prefill only when
every boundary the original prefill was cut on is a multiple of the
kernel's own fold grid — the intra-kernel chunk the recurrence is
reduced over (``GDN_CHUNK_SIZE`` for GDN); a cut anywhere else
leaves a state no uncut prefill produces. The scheduler cuts
truncated chunks on the page grid, so the page grid must be a
multiple of the fold grid, or a page-aligned restore would be a
silently different state. Refuse the pairing at boot instead.

The reduction grid of the pool's recurrent kernel, or ``default``.

ONE place answers "what grid must a cut lie on", because three do the
same arithmetic against it — the boot pairing check
(:func:`savepoint_resume_grid`), the scheduler's chunk cut
(``Scheduler._fold_align_prefill_end``) and the forward's fold split
(:func:`savepoint_fold_split`). Two of them reading the constant and one
reading a literal is exactly the drift ``GDN_CHUNK_SIZE``'s own
docstring warns about.

GDN on CUDA reduces over ``GDN_CHUNK_SIZE``; Mamba-2's SSD scan reduces
over its own ``chunk_size`` (256 on NemotronH), which is the model's, not
a kernel constant, so it is read off the layer spec the pool was built
from. A pool carrying BOTH is answered with the least common multiple:
a cut has to be on every running kernel's grid at once, and the pairing
check below turns an impossible one into a boot refusal rather than a
per-step surprise. A pool with no recurrent view is answered with
``default`` — the caller's own grid — which makes the cut and the split
degenerate to "already aligned" rather than silently picking a grid for a
kernel whose reduction width nobody here has measured.

A NON-CUDA pool is answered with 1, and that is the same rule, not an
exemption: :meth:`~arbi_serve.models.gdn_block.GDNBlock.forward` routes
a non-CUDA batch to the per-token reference fold, which reduces over one
token, so every position is on its grid. Answering 64 there would make
the scheduler cut a chunk for a kernel that is not running.

``pool`` is duck-typed: anything exposing ``_views`` keyed by
:class:`~arbi_serve.models.layer_spec.StateKind`. A pool that cannot
answer (a mock, a per-layer view) is answered with ``default`` too.

The SSD scan's ``chunk_size`` for a Mamba pool's layers.

One value: a model whose Mamba layers disagreed about it would have no
single grid a cut could be on, so a disagreement is a refusal here rather
than a grid that is right for some layers.

A recurrent fold split was asked for at a position the kernel's own
fold grid does not divide.

Raised, never rounded — and neither end of the pair has slack: rounding
the offset DOWN moves the snapshot off the resume grid, where no restore
can splice it against whole KV pages, and rounding it UP names a state
the step has not reached.

WHICH GUARANTEE THIS PROTECTS
=============================
Reproducibility, not validity, and the difference decides what may
replace this refusal. An off-grid split does not produce a wrong state:
the per-chunk update is affine in the carried state, so a different
chunking reassociates the same sum and lands within about one bf16 ulp
of the state's own scale, non-compounding over repeated off-grid cuts
(measured on arbicity/arbi-serve#2189). What it does produce is a state
that depends on HOW the prefill was cut, and a savepoint key is
``(tenant_id, prefix_hash, num_tokens)`` — content-addressed, with no
request identity in it — so one request's savepoint is restored into
another's prefill. A non-canonical state therefore makes the served
output depend on the scheduling of a request that has already finished.

So this is a refusal in defence of a STANDARD, and the standard is
stated where it is set (:func:`savepoint_resume_grid`), not here. If
that standard is ever deliberately relaxed, this becomes a counter
recording the offset's residue rather than an exception — but it is not
the kind of check that may be dropped because "the numbers look small",
because the numbers were never the reason.

The token position this step must snapshot at, or ``None``.

THE ONE PLACE THE BOUNDARY IS DECIDED
=====================================
Two callers need this number and they must never disagree: the forward,
which splits the recurrent fold so the state at the boundary exists at
all, and the commit-side trigger, which labels the snapshot with it. A
label that disagrees with the state it names is the silent-restore
failure the savepoint machinery exists to prevent, so the decision is a
pure function of the step's own numbers and both callers read it here.

``pc_before`` / ``pc_after`` bracket the prompt tokens this step's
prefill row absorbs. ``chunk_size`` is the write-rate governor (one
write per crossing of a multiple of it, wherever the step lands — see
:meth:`~arbi_serve.scheduler.savepoint_admission.SavepointSnapshotter.
snapshot_crossed_boundaries`). ``resume_grid`` is the grid a resume can
be spliced onto (:func:`savepoint_resume_grid` — the KV page size).
``fold_grid`` is the recurrent kernel's own chunk width
(``GDN_CHUNK_SIZE``), the grid on which a cut is bit-identical to no cut
at all.

``None`` means "this step writes nothing", for one of three honest
reasons: it made no progress, it crossed no ``chunk_size`` line, or it
contains no ``resume_grid`` point past ``pc_before``. Each is a normal
step shape, not a fault.

``completes_prompt`` lifts the ``chunk_size`` governor for the step that
absorbs the prompt's last token. The governor rations writes to one per
``chunk_size`` tokens of prompt; the prompt-completing step is the one
write the ration must not skip, because its last resume-grid point is
the deepest state any later admission of this prompt — a chat's next
turn, which resends it whole — can resume from. Dropped, the next turn
resumes from the previous governor crossing, up to ``chunk_size - 1``
tokens short, and re-prefills the difference on every turn. One extra
write per request at most; the end-of-step snapshot already writes it
when the step happens to end on the grid.

Raises :class:`FoldSplitUnaligned` when the boundary lies inside the
step but the offset to it is not a multiple of ``fold_grid``. That is
NOT a normal step shape: it means the step's own start ``pc_before`` is
off the fold grid, i.e. some earlier chunk was cut somewhere the kernel
does not reduce over, and every state from that cut onward is one no
uncut prefill produces. The scheduler cuts prefill chunks onto the fold
grid precisely so this cannot happen; reaching it means that cut stopped
running, and the caller must refuse to snapshot rather than round.

SHA-256 hash (truncated to 16 bytes) of ``token_ids[:n]``.

Deterministic and order-sensitive. Token ids are encoded as 4-byte
little-endian unsigned ints (matches the model's vocab fitting in
~250k → 18 bits, with 14 bits of headroom for future expansion);
the encoding is fixed-width so prepended ids can never collide
with appended ids of a different length.

Returns the raw 16-byte hash; callers store it as ``bytes`` (the
immutable hashable form) — never as ``bytearray``.

Fixed-width 8-byte little-endian encoding of ``token_ids``.

Handles the full cache-key id domain: vocab ids, negative ids
(two's-complement low 64 bits) and media-content salts ≥ 2**63
(which overflow int64 — the exact trap that made the earlier
numpy-int64 vectorization unsafe for multimodal requests).

32-byte chain seed for ``tenant_id`` (length-prefixed fold).

Mirrors the tenant fold of
:func:`arbi_serve.cache.prefix_tier_store.prefix_hash` so
cross-tenant collisions are structurally impossible: tenant A's
chain and tenant B's chain differ from the seed onward even for
byte-identical prompts.

One Merkle link: 32-byte digest for ``parent → page_tokens``.

``parent_phash`` must be the parent's RAW 32-byte digest (root =
:func:`tenant_root_phash`). The page's token count is length-
prefixed and each token fixed-width encoded, so variable splits of
the same concatenation can never collide.

Per-request incremental Merkle-chain hasher (snapshot side).

The commit-side snapshot path takes its snapshots while the
request's pages are still PRIVATE (they only enter the radix tree
at end-of-prefill ``commit_full_pages``), so there is no node to
read a ``phash`` from yet. This hasher carries the chain state
across snapshots instead: each ``digest_at`` call folds only the
NEW pages since the previous one — O(prompt) per request total —
and by construction produces the exact digest the radix node for
the same ``(tenant, prefix)`` will carry once committed (same pure
function).

Rewind-safe: preemption resets ``prompt_consumed`` (and may rewrite
``prompt_token_ids`` to prompt+output), so a later ``digest_at``
with ``n`` below the fed watermark recomputes the chain from the
tenant seed over the CURRENT token ids — a pure function, so the
result always matches the tree.

From-scratch reference for the chain digest of ``token_ids[:n]``.

Truncated to ``_HASH_LEN`` — the ``SavepointKey.prefix_hash``
surface. Every incremental producer (``_RadixNode.phash``,
:class:`ChainPrefixHasher`) must agree with this byte-for-byte;
tests pin that equivalence. Not a hot-path function (O(n) per
call) — use the incremental forms in production.

Composite key — tenant scope + prefix digest + prefix length.

``frozen=True`` → hashable, usable as an OrderedDict key.

The three components map 1:1 to the radix tree's per-tenant
namespace (``tenant_id``), the token prefix the snapshot's state
has absorbed (``prefix_hash``, a digest of ``token_ids[:num_tokens]``)
and that prefix's length (``num_tokens``) — the token a resume from
this snapshot continues at.

Per-row recurrent-state snapshot.

Holds host (CPU) tensors keyed by ``(layer_idx, attribute_name)``.
Attribute names are pool-specific:

  * GDN: ``"recurrent_state"`` + ``"conv_state"``.
  * Mamba1: ``"ssm_state"`` + ``"conv_state"``.
  * Mamba2: ``"ssm_state"`` + ``"conv_state"``.
  * ShortConv: ``"conv_state"``.

The store treats the dict as opaque — it never reads a tensor's
contents, only its ``element_size() * numel()`` for byte-budget
accounting. This decouples the store from the per-pool layout and
keeps the key surface stable as new recurrent kinds (e.g. RWKV)
land.

``num_tokens_covered`` is the number of leading prompt tokens the
state has absorbed — the admission-time ``prompt_consumed`` credit.
``put`` requires it to equal the key's ``num_tokens`` and to be a
multiple of the store's pinned ``block_size`` (the resume point is
page-aligned; see the module docstring).

Raised by :meth:`RecurrentSavepointStore.put` when the
new entry's cost exceeds the configured byte budget AND the
overflow policy is ``"raise"``.

Default policy ``"refuse"`` returns False from ``put`` instead;
admission then falls through to the no-checkpoint path (full
re-prefill — same as today's force-cache-disabled behaviour).

Content-addressed LRU + byte-budget cache for recurrent state.

Thread-safe under a single lock — the admission-side ``put`` and
the eviction-callback-side sinks can race, and the match walks the
index while admission may be replacing LRU entries.

Eviction policy:

  * On ``put``: evict LRU entries (oldest by access time) until
    ``current_bytes + new_entry.nbytes <= max_bytes``. If the
    single new entry exceeds ``max_bytes``, ``put`` returns False
    (or raises under ``overflow_policy="raise"``) WITHOUT
    evicting — protects the LRU set from being thrashed by a
    single oversized request.
  * On ``on_radix_evict``: drop the savepoint at exactly the evicted
    node's depth. The radix tree evicts leaves first, so a deeper
    savepoint under the same path was dropped by its own node's
    eviction and a shallower one still has every page it needs.

Lookup policy:

  * ``get(key)`` — exact-key lookup. Used by the admission-side
    resume path once a match has selected the longest prefix.
  * ``best_match_hashes(tenant_id, page_hashes)`` — probe the
    tenant's STORED prefix lengths, longest first, against the
    candidate's per-page digests. The first hit is the longest
    prefix the candidate agrees with. ``best_match(tenant_id,
    token_ids)`` is the same walk hashing the candidate itself.

Bookkeeping:

  * ``self._entries: OrderedDict[SavepointKey, SavepointEntry]`` —
    LRU order; ``move_to_end`` on each touch.
  * ``self._index: dict[(tenant_id, prefix_hash), set[int]]`` —
    which prefix lengths are stored under each digest; the
    eviction sinks drop by digest without scanning ``_entries``.
  * ``self._lengths: dict[tenant_id, dict[int, int]]`` — how many
    entries the tenant holds at each prefix length; the match
    walks these lengths instead of a grid, so it costs one dict
    probe per DISTINCT stored length rather than one per chunk of
    the candidate.

Metrics emitted via :mod:`arbi_serve.cache._metrics`:

  * ``arbi_serve.savepoint.hit_total{tenant}``
  * ``arbi_serve.savepoint.miss_total{tenant}``
  * ``arbi_serve.savepoint.eviction_total{reason}`` —
    ``"lru_pressure"`` from put / ``"radix_eviction"`` from the
    invalidation callback / ``"explicit_flush"`` from admin flush.
  * ``arbi_serve.savepoint.bytes_inflight`` — gauge of current
    footprint.
  * ``arbi_serve.savepoint.entries_inflight`` — gauge of count.

Truncated (``_HASH_LEN``) chain digest of ``token_ids[:n]``.

``n`` must be page-aligned (a multiple of ``block_size``) — a
savepoint's resume point always is (see the module docstring),
so every position the snapshot trigger asks for qualifies.
Monotone ``n`` folds only the new pages; a rewind recomputes
from seed.

Move ``entry``'s bytes on/off the byte budget AND the host ledger.

ONE place, because the store admits and drops entries from many call
sites and a byte moved on one counter and not the other is invisible
until something is starved by it.

A ring-backed entry contributes to the byte budget but NOT to the host
ledger: its tensors are views into a ring slot the ring already booked.
Booking them again would put one reservation under two owners.

Caller holds the lock.

The byte cap the CONFIGURATION asks for, for entries of ``nbytes``.

Two sources, and which one wins is the operator's call, not ours:

``savepoint_max_bytes``, when an operator STATED it
    It binds, verbatim. A flag named as a byte cap that a floor can
    silently overrule is not a cap; an operator who sets it to bound
    host RAM is entitled to have host RAM bounded. ``min_entries``
    does not raise it, and the boot refuses outright — see
    :func:`refuse_a_cap_that_admits_nothing` — if what they stated
    cannot hold ONE savepoint of this model's size, because that is a
    misconfiguration worth stopping for rather than overruling.

``min_entries``, when nobody stated a cap
    A savepoint's size is a property of the MODEL (layers x heads x
    state dims), not of the workload, so the DEFAULT cap buys ~1400
    entries on a small hybrid and single digits on a large one — far
    too few for the prefix cache to ever short-circuit a prefill,
    which is the store's whole purpose. So with no cap stated, the
    default is a floor and the store grows to hold ``min_entries`` of
    whatever this model's savepoint turns out to weigh, bounded by
    ``hard_max_bytes``.

Computed from ``_floor_bytes`` — the CONFIGURED value — and never from
the current cap, so re-resolving is idempotent rather than ratcheting.

Set the effective byte cap for entries of ``nbytes``. Caller holds the lock.

Two bounds, applied in order: what the configuration asks for
(:meth:`_target_cap`), then what the HOST can hold.

The host bound is the one with teeth and it is not optional for either
source. This cap governs page-locked host RAM on the route where the
store owns its own arenas, and it is reached during SERVING — at the
first snapshot, long after every boot-time budget has closed. A cap
past what the host can hold does not degrade: the kernel's global OOM
killer chooses a process. So the target is priced through
:func:`~arbi_serve.runtime.pinned_host_budget.fit_pinned_host` against
the host AS IT IS AT THAT MOMENT and clipped to what fits. Clipping an
operator's stated cap does not contradict it: a cap is a maximum, and
holding less than a maximum is always permitted.

``host_owned=False`` skips the host bound: those entries are views into
the savepoint ring's slots, whose bytes the ring already booked, so
this cap governs no host RAM of the store's own.

Install ``effective`` as the cap and say so when it moved.

Silence is not an option in either direction: a cap that grew and a cap
the host clipped are both facts an operator has to be able to read off
one boot log, and the second one is why a configured retention figure
is not the one being served.

Admit ``entry`` under ``key``.

Returns True on admit, False on refusal (entry too large for
the budget under ``overflow_policy="refuse"``). Raises
:class:`SavepointBudgetExceeded` under
``overflow_policy="raise"``.

Raises ``RuntimeError`` when the key and the entry disagree on
the coverage, or when the coverage is not a multiple of the
pinned ``block_size`` — an entry that could only be restored
wrongly is refused rather than stored.

Idempotent on re-put: the new entry replaces the old, byte
accounting refreshes, recency bumps. Idempotent under the
same value too — safe to re-call from a retry path.

Drop ``key`` iff its pinned-host ring slot has been handed out again.

True when the entry was dropped (the caller must treat the key as a
miss). Caller holds the lock. Cheap: one bool on the host route, where
every entry is live by construction.

The key for ``(tenant, ph, n)`` iff it is stored AND live.

Walks PAST a ring-backed entry that has lost its slot: a dead deep
match must not mask a live shallower one, or a bounded ring would
silently cost every resume rather than the depth beyond the bound.

Return the key of the longest stored prefix ``token_ids`` agrees with.

Hashes the candidate once per DISTINCT prefix length the tenant
has stored (longest first, first hit wins) — the flat-table
fallback; the radix path is :meth:`best_match_hashes`, which
hashes nothing.

Multi-tenant: ``tenant_id`` is the FILTER. Tenant A's match
never sees tenant B's savepoints (the index key is
``(tenant_id, prefix_hash)``).

Returns ``None`` on no match. Does NOT bump LRU recency on the
scanned-but-unmatched lengths — a missed scan is observation,
not access.

Zero-hash match: probe PRE-COMPUTED per-page digests.

``page_hashes[i]`` is the truncated chain digest for the
candidate's ``(i + 1) * block_size``-token prefix, read off the
radix tree's per-node Merkle chain by
:meth:`RadixPageTable.savepoint_page_hashes` — no hashing
happens here or there, only index lookups. The list may be
SHORTER than the candidate's page count: the tree walk stops at
the first uncached page, and a savepoint whose radix pages are
gone is unusable anyway (the admission-side KV-coverage
cross-check would drop it — the eviction listener keeps the two
stores in lockstep, so such an entry has already been dropped).

Probes the tenant's stored prefix lengths longest first; the
first digest that agrees wins — the longest honest match.

Eviction-callback sink shaped for :class:`RadixPageTable`.

``path_token_ids`` is the root→evicted-node concatenated token
sequence and ``page_phashes`` the per-page 32-byte chain
digests along it (root→node order, one per page, as delivered
by the eviction callback). The savepoint whose resume point IS
the evicted node's depth just lost the last KV page it needs,
so it is dropped — by READING the node's stored digest, zero
hashing. That is the whole set to drop: the tree evicts leaves
first, so a deeper savepoint under this path was dropped when
its own node went, and a shallower one keeps every page it
depends on.

Returns the number of savepoint entries dropped.

Per-layer state views exposed by :class:`RecurrentStatePool`.

The two dataclasses below are the flat-slab views the pool hands to
each recurrent block. They are re-exported from
``arbi_serve.cache.recurrent_pool`` so both import paths resolve to the
same classes. See the ``recurrent_pool`` module docstring for the slab
layout and lifetime contract these views obey.

Per-layer view exposed by :class:`RecurrentStatePool` for Mamba-2.

Both state attributes are flat slab tensors keyed by row index. The
pool owns tensor lifetimes — never keep a reference past
:meth:`RecurrentStatePool.free_for_request`.

  * ``conv_state``: ``(max_num_seqs, conv_dim_local, conv_kernel)``
    — channels-first, matching ``causal_conv1d_update`` on the
    ``[hidden_states ‖ B ‖ C]`` concatenated stream.
  * ``ssm_state``:  ``(max_num_seqs, num_heads_local, head_dim,
    ssm_state_size)`` — multi-head SSM hidden state, matching
    ``mamba_chunk_scan_combined`` final state shape.

``conv_dim_local`` is ``(intermediate_size + 2 * n_groups *
ssm_state_size) // tp_size``. The block reads
``slab[state_indices[i]]`` per batch row; updates are in-place.

Partial-accept rollback carries the same 1-deep base frame +
retained-input contract as :class:`GdnLayerView` (no per-draft-
position A Mamba
MTP head ships today: ``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-
A3B-NVFP4-DSpark`` is a DSpark speculative-decoding checkpoint for
this family, so :meth:`RecurrentStatePool.attach_mtp_snapshot_buffers`
allocates the fields below on a MAMBA pool and
:meth:`Mamba2Block.rollback_recompute` /
:meth:`Mamba2Block.replay_masked_commit` consume them.

Per-layer view exposed by :class:`RecurrentStatePool` for GDN.

Two flat slab tensors — ``recurrent_state`` (the
``(max_num_seqs, num_v_heads_local, head_v_dim, head_k_dim)``
matrix our vendored fused-decode Triton kernel reads / writes
in-place; V is the outer head axis to match the kernel's pointer
math) and ``conv_state`` (the ``(max_num_seqs, conv_dim_local,
conv_kernel)`` buffer the depthwise causal conv rolls each step).
Same in-place update contract as :class:`Mamba2LayerView`; the
pool owns lifetimes. FLA's chunk-prefill kernel consumes the
legacy ``(HV, K, V)`` layout — the prefill branch in
:class:`GDNBlock` transposes K↔V at the slab boundary; the
decode branch passes the slab through unchanged.

Per-token verify snapshot fields mirror :class:`Mamba2LayerView`.

ShortConv per-request state pool (LFM2 family).

LFM2's :class:`ShortConvBlock` is a 1-D causal convolution along the
token axis. Across autoregressive steps the kernel needs the last
``conv_kernel - 1`` channel-wise activations from the previous step;
those live in a per-request, per-layer dense buffer here.

Sibling to :class:`RecurrentStatePool`: same per-request dense state
shape, no paging, freed on request finish. ShortConv carries only the
conv state plus an ``has_initial_state`` flag distinguishing
first-chunk prefill from continued prefill / decode.

Slab layout: one tensor per ``(state-kind, layer)`` of shape
``(max_num_seqs + 1, conv_dim_per_rank, conv_kernel - 1)``. The
leading dim is bumped by 1 to host the zero-sentinel row at index 0
— see :class:`RecurrentStatePool` module docstring for the full
rationale. ``max_num_seqs`` stays the operator-facing usable-rows
count; row 0 is permanently zero and is never handed out by
:meth:`alloc_for_request`. The kernel expects "dim-first" layout
(channels before window) — see ``causal-conv1d``'s
``conv_state_indices`` semantics.

Per-layer per-request handle the :class:`ShortConvBlock` reads.

Wraps the layer's slab tensor + a back-pointer to the pool so the
block can resolve a request id to a slab row when needed.

Per-request, per-layer state pool for ``StateKind.SHORT_CONV``.

Allocates one slab tensor per ShortConv layer at construction time
of shape ``(max_num_seqs, conv_dim_per_rank, conv_kernel - 1)`` in
the model dtype (LFM2: bf16). Per-request rows are handed out by
:meth:`alloc_for_request` and reclaimed by :meth:`free_for_request`.

The slab is allocated with channels-before-window layout; the LFM2
kernel path (``causal_conv1d_update``) takes ``conv_state_indices``
pointing into the leading dim. If a fork eventually wants
"window-before-channels" layout the caller can transpose the view
in :class:`ShortConvLayerView` — the slab itself is layout-agnostic
once the kernel-facing call site is consistent.

Zero every slab row across every layer.

The boot capture sweep dirties ShortConv
slab rows just like it does for GDN / Mamba; called from
:func:`arbi_serve.engine.build.build` after the capture
sweep completes so the first live prefill reads zeros.
See :meth:`RecurrentStatePool.zero_all_rows` for the full
rationale.

Suspended-job KV store — host-RAM + disk tier for parked batch jobs.

The scheduler suspends a batch job the instant interactive demand
appears (or the paged-KV pool comes under pressure). A suspended job
keeps its place in ``Scheduler.running`` but produces no token. Until
now its KV pages stayed resident on the GPU — fine for a handful of
jobs, but a HUGE pooled backlog of suspended batch jobs would exhaust
the pool and either OOM or 429 fresh work.

This store lets a suspended job's KV working set spill off the GPU so
its pages free up:

    GPU pages ──D2H──▶ pinned host ──disk write──▶ NVMe blob

and fault back on resume:

    NVMe blob ──read──▶ host ──H2D──▶ freshly-alloc'd GPU pages

so the pool can hold a backlog whose cost is DISK, not VRAM (park for
up to ``ttl_seconds``, default 24h). On resume the content lands in a
NEW page allocation (the original pages were freed and may have been
reused), which is why the store carries the page CONTENT, not page ids.

Built on :class:`arbi_serve.cache._lru_byte_budget_store._LruByteBudgetStore`
(host-RAM LRU + byte budget with pinned-host D2H staging, lazy-confirm
event barrier, and an optional disk-spill tier). The suspended-job
specifics over that base:

  * Keyed by **request/job id**, not a prefix hash. A suspended job is
    a single live request; there is no content-addressed sharing.
  * Snapshots a request's **paged-KV page content** (the actual KV for
    its allocated pages) AND its **recurrent state** if hybrid/GDN
    (reusing the recurrent_savepoint snapshot for the GDN part).
  * The disk blob is an unencrypted ``torch.save`` (cold suspended-job
    KV) with a small **persisted index** so suspended jobs survive a
    server restart.

Tiering policy
==============
  * ``stage(job_id, ...)`` — admit a suspended job's snapshot into the
    HOST tier. Evict host-LRU to disk until the host byte budget fits.
  * Cold host entries spill to disk (a disk blob + index entry; the host
    bytes free). Disk is bounded by a max-bytes cap (LRU evict) AND a
    TTL sweep (drop blobs older than the TTL).
  * ``fault_in(job_id)`` — return the snapshot whether it lives in host
    RAM (fast) or on disk (read blob back to host first). Resume always
    works regardless of tier.
  * ``drop(job_id)`` — the job finished / was cancelled; free host bytes
    and unlink any disk blob.

Restart durability
==================
The disk index is a JSON sidecar (``index.json``) rewritten on every
disk mutation. On construction the store reloads it, dropping expired
or missing-blob entries, so suspended jobs that were on disk at the
last shutdown are faultable after a restart.

Thread-safety
=============
A single re-entrant lock guards the in-RAM maps and the index. Disk
I/O (blob read / write) happens OUTSIDE the lock — the lock only
covers the bookkeeping mutation. The disk-write worker is the caller's
(``asyncio.to_thread`` from the engine); this module's ``spill`` /
``fault_in`` are synchronous and meant to be called from a worker
thread.

One suspended job's full KV working set.

Carries everything ``Scheduler`` needs to resume the job's decode
bit-identically:

  * ``page_tensors`` — ``{tag: host_tensor}`` from
    :meth:`MultiStatePool.snapshot_pages`. The paged-KV content of
    the job's allocated pages, in page order.
  * ``num_pages`` — how many pages the job held (so resume can
    ``alloc_pages`` the same count and scatter back in order).
  * ``recurrent_tensors`` — ``{(layer_idx, attr): host_tensor}`` from
    :meth:`MultiStatePool.snapshot_recurrent_row` for hybrid models;
    empty dict on pure-attention models.
  * ``meta`` — opaque small dict the scheduler stashes to reconstruct
    the request's progress on resume (token ids, prompt_consumed,
    ...). The store never inspects it; it is JSON-serialised into the
    disk index for restart durability.

``pending_event`` mirrors the lazy-confirm contract: when the
snapshot was issued async (non-blocking D2H), the recorded CUDA event
must fire before the host tensors are read or freed.
:meth:`synchronize` waits on it once, then clears it.

Two-tier (host-RAM + disk) LRU store of suspended-job KV snapshots.

Host tier: bounded by ``host_max_bytes`` (LRU spill to disk).
Disk tier: bounded by ``disk_max_bytes`` (LRU unlink) AND
``ttl_seconds`` (age sweep). Both tiers are keyed by job id.

Lookup (``fault_in``) checks host first, then disk (reading the blob
back to host). ``stage`` admits to host. ``drop`` removes from both.

Force ``job_id`` from host → disk (frees host bytes).

Used by the engine's background spiller to proactively cool the
host tier (e.g. when the host budget is approached). No-op /
False when the job isn't host-resident or no disk tier exists.

Serialize an entry's host tensors to a single blob file.

Format: ``torch.save`` of a dict carrying the page tensors and
recurrent tensors (CPU). Simple + robust; the per-tensor dtype /
shape ride along so fault-in needs no external manifest. Returns
a small manifest recorded in the index (tags + shapes) for
observability / restart validation.

Offline TKV calibration against the *served* (AWQ-INT4) model.

The engine loads weights `transformers` can't (compressed-tensors
group_size=0), so calibrating through the engine measures fidelity against
the exact weights we deploy — not a bf16 proxy.

This package is import-light and offline-only: the production serving path
never imports it, and its runtime touch-points (the ``_attn_eager`` KV seam
check, and ``model_runner``'s captured-graph-replay exclusion for an active
seam — see ``kv_seam.is_injecting()``) are single global-``None``-backed
checks that compile away when inactive.

transformers-free tokenizer adapter for the calibration corpus builders.

The slim production image ships only the Rust ``tokenizers`` library (via
:mod:`arbi_serve.tokenizer`), not ``transformers`` — see the pyproject note
"NOT a runtime dep — the server ships its own tokenizer". The calibration
pipeline (``engine_capture`` / ``engine_solve``) runs inside that slim image
at boot (the autocal prepare phase), so it must not import ``transformers``
either, or autocal ``ImportError``s on a correctly-slimmed image.

This thin adapter wraps arbi-serve's own :class:`~arbi_serve.tokenizer.Tokenizer`
and exposes exactly the small surface tkv's ``build_corpus_tokens`` /
``build_balanced_basket`` call against an HF tokenizer:

  - ``tok(text, add_special_tokens=..., return_tensors="pt").input_ids`` →
    ``(1, T)`` long tensor
  - ``tok.encode(text, add_special_tokens=...)`` → ``list[int]``
  - mutable ``pad_token`` / ``eos_token`` attributes

so those builders work unchanged with no HF-modeling dependency.

Process exits that do not run interpreter finalization.

A child that built an in-process Engine allocates through the cuMem
``CUDAPluggableAllocator`` (torch ``MemPool``). At normal interpreter
exit, torch's C++ static teardown destroys the pluggable allocator
before the caching allocator releases the private-pool blocks;
``DeviceCachingAllocator::release_block`` then makes a virtual
``raw_delete`` call through the dead allocator — ``__cxa_pure_virtual``
→ SIGABRT (or SIGSEGV via a stale vptr, build-dependent). The child's
work is already complete and flushed by then, so the crash only
corrupts the exit code the parent sees. ``os._exit`` skips the doomed
teardown entirely. The same hazard reaches every process that built an
engine, not only children: the ~8 cuMem-backed pools are deliberately
parked in module globals so their VRAM is reclaimed by the OS at exit,
which puts CUDA objects in the path of ``Py_FinalizeEx``. A boot refusal
that unwinds instead of hard-exiting therefore reports SIGSEGV/SIGABRT
where it meant to report its own exit code, and an operator, a
supervisor or a CI lane reads a deliberate refusal as a crash.

:func:`exit_refusal` is the seam every in-process engine build refuses
through, and :func:`exit_cli_process` the entrypoint backstop for the
paths that do not build one directly. Stdlib-only; safe to import
anywhere.

Push every buffered log record to its sink. Best-effort, never raises.

``os._exit`` skips ``atexit`` and every library's own teardown, so
whatever a sink is still holding is simply lost — and the records a
hard exit most needs delivered are the ones explaining WHY it is
exiting, written moments before. Three sinks hold records:

  * stdlib ``logging`` handlers (``logging.shutdown``);
  * the loguru sinks, including the OTLP bridge
    (``logger.complete()`` — a no-op for the synchronous stdout
    sink, and the awaited drain for any enqueued one);
  * the OTEL Logs SDK's ``BatchLogRecordProcessor``, which batches on
    a background thread and needs an explicit ``force_flush``.

``otlp_timeout_ms`` bounds only the network-facing flush: a collector
that is down must not turn a deliberate exit into a hang, so the
export is given a bounded chance and then abandoned.

A :func:`hard_exit` a HOSTING process asked to receive rather than die of.

Derived from ``BaseException``, not ``Exception``, and that is the whole
point: every one of these exits is raised from inside an ``except`` block
or a ``contextlib.suppress(Exception)``, so an ``Exception`` subclass
would be swallowed by the very code that was about to exit — turning a
silent death into a silent no-op, which is worse.

Receive :func:`hard_exit` as :class:`HardExitRequested` in this process.

For a process that HOSTS a server rather than being one — the test
session, and nothing else. It does not make the exit safe; it makes it
VISIBLE. The state a hard exit exists to avoid unwinding (an
uninterruptible build thread, a half-built private pool) is still there
afterwards, and a host that catches this should report and stop, not
carry on as though the boot had succeeded.

Process-wide rather than per-call because the eleven call sites are
reached through code that neither the host nor the exit can see between
them, and a hook threaded through all of it is a hook that one path
forgets to carry.

Report a boot refusal, then leave without interpreter finalization.

Call from inside the ``except`` block that caught the refusal: the
traceback is logged from the live exception context, the boot-progress
state is flipped so ``/health`` reports the reason, and the process
ends at :func:`hard_exit` rather than unwinding through a finalization
that a live cuMem pool can fault (see the module docstring).

Run a CLI entrypoint and end the process on its outcome.

The backstop under :func:`exit_refusal`: any FAILING outcome — a
non-zero return, an escaping exception, a non-zero ``SystemExit`` —
leaves via :func:`hard_exit` so the code the caller sees is the code
the CLI chose, never a teardown signal. A successful run keeps stock
interpreter teardown, so ``atexit`` hooks and buffered writers behave
exactly as they always have.

TP>1 drift-seam control op — activate the inject seam on worker ranks.

At ``tp_size > 1`` each rank owns a kv-head shard and ``o_proj`` all-reduces
them, so a drift measurement that only injects codecs into rank-0's
process-local :mod:`arbi_serve.calibration.kv_seam` mixes rank-0's quantised
heads with the workers' *clean* heads — the measured drift comes out ≈
``1/tp_size`` of the real drift and the drift-scales are silently too small
(see ``calibration/inprocess.py::_assert_calibratable_topology``).

The fix (design: ``docs/tp-ep-inprocess-calibration-design.md``) is to activate
the same inject seam symmetrically on every rank. Rank 0 drives it through the
driver's FIFO shm ControlOp ring (``distributed_driver._broadcast_control`` —
it does not touch NCCL, so it is deadlock-free and shares one FIFO order with
the forward-batch broadcasts). This module is the worker side of that op plus
the small (de)serialization helpers rank 0 shares, so both sides build a
byte-identical codec.

Three actions, all carried by ``ControlOp(op="drift_seam", payload=...)``:

* ``load_bundle`` (once, drift-stage start) — rank 0 writes the stage-1 bundle
  (centroids + initial scales) to a temp path and broadcasts the path; each
  worker loads it rank-locally from disk (the swap's rank-local load pattern —
  not ``dist.broadcast`` on the bundle) and caches it. Centroids never change
  during stage 2, so this ships the big data once.
* ``install`` (per measure) — ships only the small per-channel *scales* (which
  stage 2 mutates; centroids don't). Each worker applies the scales onto its
  cached bundle, builds its shard's codec via
  ``tkv.calibration.empirical_anchor._build_layer_codec`` (exactly as
  ``EngineDriftSource._get_codec`` does — head-agnostic, rank-invariant
  centroids ⇒ byte-identical to rank-0's codec) and sets
  ``kv_seam._SEAM = _InjectSeam(codecs)``.
* ``clear`` (per measure) — each worker sets ``kv_seam._SEAM = None``.

FIFO ⇒ every worker applies ``install`` → the candidate's forward deltas →
``clear`` in exactly that order, so the seam is active precisely during the
candidate's forwards. Rank 0 keeps its own local ``with kv_seam.inject(...)``
unchanged; it merely also brackets the measure's forwards with install/clear.

``install`` payload for ``alloc`` against ``bundle``.

Ships the allocation and only the small mutated per-channel scale vectors
for the layers under test (centroids reached the workers once, via
``load_bundle``). Mirrors what ``EngineDriftSource._scale_key`` /
``_build_layer_codec`` do on rank 0, so the worker rebuilds an identical
codec.

Worker-side dispatch for a ``drift_seam`` :class:`ControlOp`.

Called from the one shared worker admin handler
(``DistributedEngineDriver._apply_worker_control_op``) so both the legacy
and SPMD worker loops route it identically. ``device`` is this rank's engine
device.

Load the stage-1 bundle rank-locally from ``path`` (no NCCL).

``device`` is this rank's engine device — codecs must live on it so the
seam's round-trip tensors are on the same device as this rank's kv-head
shard (under torchrun rank-1's KV is on ``cuda:1``).

Arm the parameterless fp8-e4m3 seam on this worker.

The fp8-KV drift probe (NOT a servable KV mode): unlike ``install`` this
needs NO bundle, centroids, or scales — every rank round-trips ITS
kv-head shard through the same fixed scaled-e4m3 transform, so the
all-reduced attention output is fully fp8-quantised (a rank-0-only seam
would measure ≈ Shares the ordering guard: a non-None seam here
means a prior measure was never cleared.

Disarm the seam and drop this measure's codecs (bounds GPU memory).

Dropping ``_SEAM`` first releases the ``_InjectSeam``'s references (or
the parameterless ``_Fp8Seam``), then the cache clear frees any codecs'
GPU-resident centroids/scales — the next install rebuilds from the
cached bundle (cheap). Safe for the fp8 seam too (its cache is empty).

In-process prefill capture driver for AWQ-faithful TKV calibration.

Boots an in-process :class:`~arbi_serve.engine.engine.Engine` against the
*served* (AWQ-INT4) model and drives a **teacher-forced prefill** over a
small calibration corpus, recording per-layer K/V through
:mod:`arbi_serve.calibration.kv_seam` and per-position reference logits
through a cudagraph-replay-safe raw-logits tap on the model runner (see
:class:`_LogitsSink`). The captured K/V feed turbo-attn's
``tkv.calibration.calibrate_centroids.fit_centroids_from_captured_kv`` to
produce a stage-1 centroids bundle fit on the *exact* weights we serve
(raw ``transformers`` cannot load Qwen3.6-27B-AWQ-INT4).

Prefill only — sidesteps the GDN decode-path state-index bug. Each
calibration sequence is one prefill forward with ``max_tokens=1,
temperature=0``: the single emitted token comes from the prefill of the
last prompt position; no autoregressive recurrent-state-slab decode read
ever happens. ``_paris_smoketest`` verifies prefill is sane before any
capture (if even prefill is garbage, calibration must stop).

This module is the C2 component of the calibration plan. It builds on the
C1 seam already wired in ``qwen3_5.py::_attn_eager``.

TP: the kv_seam global is process-local, so under torchrun each rank
captures its own kv-head shard. Centroids are fit per (layer, side) over
``head_dim`` (rank-invariant — every rank sees the full head_dim), so
rank-0's capture is a complete, deployable stage-1 bundle. The
``--tp-size`` smoke confirms capture survives the TP build; merging
across ranks is unnecessary for centroids.

Calibration corpus as token-id rows, one row per run/seed.

Returns ``(rows, labels)``. The labels are per-row provenance and are
stamped through to the run report, so which corpus a bundle was fit on is
readable off the bundle rather than inferred from how it was launched.

The ``balanced`` arm fails LOUD when the materialised real corpus is
absent. The solver's own basket builder falls back to c4 in that case,
which is right for a drift objective but wrong here: a silent fallback
would make the treatment arm of a capture-corpus A/B secretly identical to
its control, and the comparison would report "no effect" for a change that
never ran.

Captures per-forward reference logits via the model runner's raw-
logits tap (:meth:`~arbi_serve.runtime.model_runner.EagerModelRunner.
set_raw_logits_sink`), NOT a ``lm_head`` forward hook.

A forward hook is a Python ``nn.Module.__call__`` callback: it fires
during eager forward and during cudagraph CAPTURE (both real Python
calls) but never during a captured-graph REPLAY, which launches only
the recorded CUDA kernels with no Python in the loop — so a hook-based
sink silently stops firing the moment ``cuda_graphs=True`` starts
replaying. The runner's ``_finish_execute`` chokepoint hands us the
same tensor a hook would have seen (the model's final, post-
``collapse_dual_head_logits`` last-position logits) whether that step
replayed a graph or ran eager, because a captured replay already
``.clone()``s its persistent output buffer back to a plain tensor
before returning up the call stack. This also makes the old plain-vs-
dual-head ``resolve_lm_head_module`` resolution unnecessary here: the
tap sees the SAME collapsed tensor regardless of architecture.

We record the detached CPU logits per forward; the driver concatenates.

Captures per-step K/V via the model runner's raw-KV tap
(:meth:`~arbi_serve.runtime.model_runner.EagerModelRunner.
set_raw_kv_sink`), NOT :mod:`arbi_serve.calibration.kv_seam`'s
``_CaptureSeam.apply()`` live hook.

``_CaptureSeam.apply()`` is a Python method called from ``attn.py``'s
``AttentionBlock.forward`` (via ``kv_seam.maybe_apply``) — a real
Python call, so it fires during eager forward and during cudagraph
CAPTURE, but never during a captured-graph REPLAY (only the recorded
CUDA kernels launch, no Python runs). A chunk that happens to hit a
captured bucket therefore silently contributes NOTHING to the capture.
Real-verified against this engine (Qwen3-0.6B, real 4-row/2048-token
basket): under ``cuda_graphs=True`` this dropped exactly half the
captured tokens (1024 of 2048), deterministic and reproducible, not
noise — every downstream centroid fit was silently built from half
the intended data.

The fix mirrors :class:`_LogitsSink`'s: read a persistent, stable-
address buffer instead of relying on a live forward-time callback.
Here that buffer is the paged KV pool itself
(``eng.pool.layer_view(layer_idx)``) — every attention kernel
scatters newly-computed K/V into it as a normal side effect of
running, baked into a captured graph's kernel sequence the same as
any other op, so it holds the real data for this step's tokens
whether the step replayed or ran eager. ``batch.slot_mapping``
(handed to the sink by ``_finish_execute``) says exactly which flat
``page*page_size + offset`` slots those tokens landed in.

Scoped to the standard (non-MLA) fused single-plane bf16 tkv-bypass
layout ``arbi_serve.backends.tkv_bypass_backend`` documents:
``(num_pages, page_size, 2*H_kv*D)``, ``K = [..., :H_kv, :]``,
``V = [..., H_kv:, :]`` after viewing as ``(*, 2*H_kv, D)`` — the
backend ``build_engine`` always selects for calibration. MLA models
(a single shared latent per token, a differently-laid-out uint8 slab
— see ``cache/mla_pool.py``) are NOT covered by this reader and raise
rather than silently mis-slice; ``attach`` refuses up front instead.

The ``lm_head`` submodule to hook for reference-logits capture.

The common case is a plain top-level ``model.lm_head`` attribute (every
single-head dense/hybrid arch). Dual-head composite models — the
``LOGITS_HEAD_NAMES`` capability marker
(:class:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel` /
``NemotronVoiceChatBackboneModel``, see that class's docstring) — don't
expose a top-level ``lm_head`` at all (theirs lives nested under
``backbone.stt_model.lm_head``), so they instead expose
``model.calibration_lm_head`` pointing at the SAME head the generic
single-tensor engine dispatch already samples
(``collapse_dual_head_logits`` picks ``LOGITS_HEAD_NAMES[0]``, i.e.
"text") — same duck-typed capability-check convention as
``LOGITS_HEAD_NAMES`` itself, so this function never needs an
``isinstance`` check on a concrete arch.

Return the per-layer ``layer_types`` list for ``text_cfg``, or None.

Three spellings resolve to the same vocabulary: an explicit
``layer_types`` (Qwen3.5, Gemma-4, LFM2), a ``layers_block_type`` list,
and NemotronH's ``hybrid_override_pattern`` string. Returns None when the
config declares none of them, leaving the caller to assume a dense stack.

Return ``([(layer_idx, layer_type), ...], num_kv_heads)`` for the
KV-bearing (full/sliding attention) layers.

Mamba / GDN / linear-attention / MLP-only layers carry no KV cache and are
excluded — they are exactly the layers
``fit_centroids_from_captured_kv`` must not see (and the kv_seam never
fires for them, since it lives in the attention block's ``_attn_eager``).
A hybrid stack whose non-attention layers were counted as KV-bearing fails
the capture-completeness check in
:func:`~arbi_serve.calibration.inprocess.calibrate_engine_inplace`.

MLA configs report one KV head: the cache holds a single shared latent
per token, whatever ``num_key_value_heads`` says about the unfolded
attention heads.

Two resolution paths:

* ``eng`` given (the in-process ``--calibrate-only`` path, which always
  has a live built model before this is called): read
  ``eng.model.layer_specs`` directly and filter to
  :attr:`~arbi_serve.models.layer_spec.LayerKind.ATTENTION`. This is the
  model's OWN real per-layer topology — correct for every architecture,
  including hybrids (NemotronH/NemotronVoiceChat Mamba+attention+MLP,
  LFM2, Qwen3.5 GDN+attention) with no per-arch special-casing, and the
  only path validated against a checkpoint whose ``config.json`` is not
  even the standard HF shape (NemotronVoiceChat's is a raw NeMo Hydra
  training config — see ``models/nemotron_voicechat.py``'s module
  docstring).
* ``eng`` omitted (the offline ``_drive`` CLI, which calls this BEFORE
  building an engine): fall back to the config.json path below.
  ``_layer_types_from_config`` resolves the three HF spellings —
  ``layer_types``, ``layers_block_type``, and NemotronH's
  ``hybrid_override_pattern`` — so a shipped NemotronH-family checkpoint
  now resolves its real Mamba/attention/MLP split here. A config
  declaring NONE of the three (NemotronVoiceChat's raw NeMo Hydra
  config) still falls through to "every layer is full_attention", which
  is wrong for a Mamba/MLP-mixed hybrid; the offline driver has no model
  to introspect before this call, and ``arbi-serve --calibrate-only``
  uses the in-process path above.

Construct + build an in-process Engine for calibration.

MTP disabled (no speculative decode needed), tkv-bypass paged KV backend
(calibration reads the K/V *into* the seam, not from a quantised cache).

``cuda_graphs=True`` is the calibration default. Getting there took
three real, independently-verified fixes — read on, since removing any
one of them silently reopens a real correctness gap:

1. Reference-logits capture goes through the model runner's raw-logits
   tap (:class:`~arbi_serve.calibration.engine_capture._LogitsSink` /
   :meth:`~arbi_serve.runtime.model_runner.EagerModelRunner.
   set_raw_logits_sink`) instead of a ``lm_head`` forward hook, so it
   survives cudagraph replay (a forward hook does not — replay
   launches only the recorded CUDA kernels, no Python
   ``nn.Module.__call__``).

2. The solver's real per-candidate measurement
   (``EngineDriftSource.ameasure`` / ``ameasure_decode``) installs a
   real, per-allocation codec round-trip via ``kv_seam.inject`` — a
   captured graph replays the EXACT kernels recorded at boot warmup
   (seam always off then), so replaying one under injection used to
   silently skip the injected round-trip and return the seam-off
   reference's own logits for every candidate. Real-verified the hard
   way: before ``model_runner`` gained an ``is_injecting()``-gated
   eager fallback (see ``_forward_or_replay`` / ``kv_seam.
   is_injecting``), the search's per-candidate prefill objective came
   back EXACTLY 0.0 for every one of 6 distinct real allocations under
   ``cuda_graphs=True`` — silently wrong, not a crash, and worse than
   the old hook's empty-``torch.cat`` crash. That gate is load-bearing;
   do not remove it.

3. ``capture_corpus``'s stage-1 K/V capture used
   ``kv_seam._CaptureSeam.apply()`` — the SAME class of live-Python-
   hook bug as #1, just for K/V instead of logits: a Python method
   called from ``attn.py``'s eager forward, never invoked during a
   captured-graph REPLAY. Real-verified: on a real 4-row/2048-token
   basket, this silently captured exactly HALF the tokens (1024 of
   2048) under ``cuda_graphs=True`` — every downstream centroid fit
   silently built from half the intended data, and that (not any
   "cudagraph build" effect — a clean isolation test with everything
   else forced eager showed BIT-IDENTICAL results) is what produced
   wildly allocation-dependent KL-drift distortion end to end (one
   genuinely mild K6V6 allocation read 0.0017 nats reference-vs-
   candidate all-eager, 0.90 nats under the unfixed
   ``cuda_graphs=True``). Fixed the same way as #1: the model
   runner's raw-KV tap (:meth:`~arbi_serve.runtime.model_runner.
   EagerModelRunner.set_raw_kv_sink`, read from :class:`_KvSink`)
   reads K/V back from the paged KV pool itself
   (``eng.pool.layer_view``) via that step's ``batch.slot_mapping`` —
   a persistent, stable-address buffer every attention kernel
   scatters into as a normal side effect of running, unlike a
   forward-time callback, so it holds the real data whether the step
   replayed or ran eager.

``EngineDriftSource._aprefill`` / ``_adecode``'s ``codecs is None``
(reference/jitter) branch ALSO routes through ``kv_seam.inject({})`` —
a real, empty-map injection, a true numeric no-op, but one that takes
the ``is_injecting()``-gated eager path fix #2 installed — so the
reference and every candidate are always measured through the
identical regime. This matters independently of #3: even with #3
fixed, comparing a captured-graph-REPLAYED reference against a
forced-eager candidate measures real (if small, ~0.01-0.14 raw-logit,
ordinary bf16 replay-vs-eager kernel-dispatch) noise that the top-k-
restricted KL metric can amplify — eliminated by keeping both sides on
the same footing rather than chasing the kernel-level noise itself.

With all three in place: real-verified end to end on a real 6-
allocation prefill + decode-accumulated search (Qwen3-0.6B) — every
one of 12 real (allocation × measurement-kind) drift values is
EXACTLY bit-identical (0.000000 diff) between an all-eager run and a
``cuda_graphs=True`` run, and raw captured K/V + fitted centroids are
bit-identical too. ``kv_seam.inject`` measurement forwards correctly force eager either
way (fix #2), so they cost the same as before — no regression, just no
further speedup available there without recapturing per candidate
(which would cost far more than it saves). ``cuda_graphs=False``
remains available via ``--no-cuda-graphs`` for debugging.

Submit a single teacher-forced prefill (max_tokens=1, greedy) and
return ``(first_token_id, decoded_token)`` for the prefill-of-last-
position token.

The token *id* is read from ``output_token_ids`` (set synchronously on
the step that finishes the request) and decoded via the engine
tokenizer — not from ``output_text``, which is populated by the async
detokenizer task and races with our post-finish read (empty-string
reads were the symptom).

One throwaway prefill to warm the engine.

The first prefill submitted to a freshly-built engine is unstable —
its sampled token differs from every subsequent (identical) submission
(observed: a cold "The capital of France is" greedily emits "The"; the
same prompt run warm deterministically emits " Paris"). The cause is
first-forward warmup state (initial Dynamo trace + cold GDN recurrent
savepoint store), not the GDN decode-path state-index bug and not the
kv_seam — once a single warmup prefill has run, all prefills are
deterministic and correct (France→Paris, Italy→Rome, Japan→Tokyo all
verified). We run the warmup outside the capture window so its K/V is
never recorded.

Prefill 'The capital of France is' → must be ≈ ' Paris'.

Guards against the GDN decode-path bug corrupting even the prefill
path. Requires a prior :func:`_warmup` (the first-ever prefill is
unstable). Returns the decoded token for the caller to assert / report
on.

Move every device-resident clone in ``sink`` to host, in place.

Called after each prefill so the capture's VRAM residency is bounded by
one prefill rather than the whole basket. Idempotent: tensors already on
host are left alone, so a second call is free.

The device copies are dropped here, which returns their blocks to torch's
caching allocator. Every prefill in the basket allocates the same shapes,
so the next one reuses those exact blocks — no ``empty_cache()`` needed
(and none wanted: it would hand the memory back to the driver and force
the serving pools to re-acquire it).

Run ``token_rows`` as prefills under an active raw-KV tap + a
raw-logits tap on the model runner. Returns ``(captured_k, captured_v,
ref_logits)``.

``captured_k[layer]`` / ``captured_v[layer]`` are lists of
``(N_tokens, num_kv_heads, head_dim)`` tensors (one per prefill); the
centroid-fit core treats each list as a multi-run pool and concatenates
on the token axis. Logits are one ``(1, vocab)`` row per prefill.

Read served-codec reconstruction MSE per (layer, side, bit) straight
out of the multi-bits bundle. ``served_mse`` is the better-of
(empirical-Lloyd, analytic-Gaussian) reconstruction MSE on the real
pooled activations — the quantity that must decrease with more bits.

Returns ``{layer_str: {bit: {"k": mse, "v": mse}}}``.

Report whether the sparse K=7/K=8 centroid tails resolved on the pool.

If the capture corpus is too small the high-bit codebooks starve: their
reconstruction MSE fails to fall below the low-bit codebooks (the tails
have no support). We summarise served_mse(b8)/served_mse(b4) per side — a
healthy, tail-resolved capture has this ratio well below 1.0 for nearly
every layer (more bits => strictly lower MSE).

Inject-forward drift source for AWQ-faithful TKV calibration (C3b).

The one new calibration primitive: given a candidate per-layer K/V bit
allocation against a stage-1 centroids+scales bundle, measure the
decode-weighted, top-k-restricted KL drift of the *served AWQ* model's
logits-under-codec vs a clean (seam-off) reference forward — using the
arbi-serve :class:`~arbi_serve.engine.engine.Engine` prefill, not a
``transformers`` forward.

This mirrors :func:`tkv.calibration.empirical_anchor.evaluate_candidates`:
the same codecs (``_build_layer_codec``), the same forward-KL definition
(top-k restricted, renormalised over the ref's top-k support), the same
jitter-floor normalisation — only the forward is swapped from HF to the
AWQ engine via :func:`arbi_serve.calibration.kv_seam.inject`.

Why one logit per prompt
------------------------
The Qwen3.5 engine computes ``lm_head`` only on the last prefill position
(``qwen3_5.py``: ``last_hidden = hidden[last_idx]``) — a serving engine
only needs the next-token distribution. So each prefill yields exactly one
reference next-token distribution (the end-of-prompt position). The drift
basket is therefore N independent prompts → N reference logits; drift is
the mean over prompts of the top-k restricted KL at that one position.
This is the prefill-only ``decode_start == T-1`` slice of the established
metric (forward KL, top-k=64, jitter-normalised). Prefill only by design
(sidesteps the GDN decode bug, as the capture driver does).

The :class:`EngineDriftSource` is the injected "drift backend" the
turbo-attn solver/scale-search consume through the
:class:`tkv.calibration.drift_source.DriftSource` protocol — the engine
cannot resume-from-layer (the forward is opaque), so ``probe == measure``
(a full prefill per candidate) and ``promote`` is a no-op.

Recursively convert nested lists to nested tuples so the value is
hashable. Scale vectors may be flat (``head_dim`` floats) or nested
(per-head), and a flat ``tuple(...)`` of a nested list is still unhashable
— so recurse rather than assume the shape.

Run each token row as one teacher-forced prefill and collect the
last-position lm_head logits. Returns ``(N, vocab)`` fp32 CPU tensor.

Uses the model runner's raw-logits tap (same sink the capture driver
uses — see :class:`~arbi_serve.calibration.engine_capture._LogitsSink`
for why this is NOT a ``lm_head`` forward hook: a hook never fires
under cudagraph replay) rather than the sampled token, so we get the
full next-token distribution at the end-of-prompt position regardless
of sampling.

Like :func:`_prefill_logits`, but continued past the prompt: each row
runs its teacher-forced prefill AND ``decode_steps`` real autoregressive
decode steps, collecting the lm_head logits at every position along the
way. Returns ``(N, decode_steps + 1, vocab)`` fp32 CPU — position 0 is
the prefill-only ``_prefill_logits`` slice (end-of-prompt), positions
1..decode_steps are real accumulated-decode positions.

Same raw-logits-tap mechanism as :func:`_prefill_logits` (this module's
established seam — see
:class:`~arbi_serve.calibration.engine_capture._LogitsSink` for why it
is a runner-level tap and not a ``lm_head`` forward hook), just left
armed for ``max_tokens = decode_steps + 1`` instead of 1.
``ignore_eos=True`` so a real EOS mid-basket can't truncate a row short
and silently misalign the chunking below.

One real forward per real output token is NOT the whole story, though
(live-verified against this engine): the step scheduler pipelines one
extra lookahead forward past a request's LAST real token before
noticing ``max_tokens`` was hit, so a request that emits ``k`` real
tokens fires the raw-logits sink ``k + 1`` times — argmax-verified: for
``decode_steps=2`` (3 real output tokens), 4 sink fires arrived with
argmax ids ``[tok0, tok1, tok2, <discarded>]``, the first 3 exactly
matching the real ``output_token_ids`` and the 4th matching nothing
that was ever emitted. The trailing fire is discarded, never the
leading/middle ones, so it is always safe to keep a row's FIRST
``decode_steps + 1`` fires and drop anything after — hence resetting
the sink per row (submissions are sequential, one row fully finishes
before the next starts) rather than assuming a single global count.

Per-row forward KL(ref ‖ cand) restricted to the REF top-k support.

``ref_logits`` / ``cand_logits`` are ``(N, vocab)`` raw logits. For each
row: take ref's top-k ids, gather both sides' logits at those ids,
log-softmax each over the K-subset, then KL(ref_sub ‖ cand_sub).
Numerically the per-position form of
:func:`tkv.calibration.bit_allocation_solver._per_position_kl_topk`.
Returns ``(N,)`` KL in nats.

Fail loud when ``eng``'s prefill chunk size doesn't match its KV page
size (``cache.block_size``) — the exact regression a real incident
traced to a driver that hand-built ``ServerConfig``/``Engine`` directly
instead of reusing :func:`arbi_serve.calibration.engine_capture.
build_engine` (which always sets the two equal; see that function's
"fresh-diagonal fix" comment for the full derivation this mirrors).

Why the check lives HERE (at :class:`EngineDriftSource` construction)
rather than in the shared ``Engine.build()`` path: ``chunk_prefill ==
block_size`` is NOT a general engine invariant — real serving routinely
(and correctly) runs the two unequal (e.g. the default chat profile's
chunk_prefill=2048/block_size=256, or embed/rerank's deliberate
chunk_prefill>=max_context single-chunk-prefill forcing in
``cli/config_builder.py``). It matters ONLY for the specific measurement
``EngineDriftSource`` performs: a single teacher-forced PREFILL scored at
the last-prompt-position logit, with the candidate codec swapped in via
:func:`arbi_serve.calibration.kv_seam.inject`. Per
``arbi_serve/models/attn.py``'s ``_attn_eager`` — "Prefill + inject needs
the real fresh-diagonal split" — a chunk's OWN intra-chunk self-attention
is always computed on FRESH, un-injected K/V (this mirrors real
production's Phase-1 prefill bypass: there is nothing quantized to read
back yet within a chunk's own diagonal). The candidate codec is only
ever exercised by a LATER chunk's cross-chunk read of an EARLIER
chunk's already-quantized cache. So when ``chunk_prefill`` is large
enough that a basket prompt never crosses a chunk boundary within ONE
prefill (the common case for a short calibration/search basket — a
handful of thousand tokens at most), the candidate codec is never read
back at all, and every candidate allocation measures identical
(near-)zero drift. A solver walking that flat, degenerate objective can
converge on an allocation that looks free in the broken proxy while
scoring far worse on the real served ladder — a real ~3.5h GPU search
over exactly this broken proxy once produced an allocation ~778% worse
than plain uniform K4V4 on the real canonical ladder. (Decode-based
measurement — ``_adecode`` — is NOT affected by this: a decode step
always reads back the already-committed, already-quantized cache
regardless of how prefill was chunked, per the same comment's "Decode
doesn't need this split." The check still applies unconditionally here
because every :class:`EngineDriftSource` also builds its seam-off
PREFILL reference via the same code path.)

Skips silently when ``eng`` doesn't duck-type a real, built engine (a
stub/fake ``eng`` — several unit tests pass ``eng=None`` or a bare
``SimpleNamespace``/``object()`` to exercise logic that never reaches a
real forward) — there is no live config to check and no real forward
will run.

Drift backend backed by the arbi-serve AWQ engine inject-forward.

Conforms to the :class:`tkv.calibration.drift_source.DriftSource`
protocol the solver / scale-search call. Built once per run: it owns the
token basket, the per-prompt reference (seam-off) logits, and the
self-jitter floor; ``measure(alloc)`` injects the candidate codecs and
runs a fresh prefill of the whole basket.

Parameters
----------
eng
    A built, running :class:`~arbi_serve.engine.engine.Engine` (prefill
    only; MTP disabled, eager for bring-up). Caller owns its lifecycle.
    Must have ``cfg.batch.chunk_prefill == cfg.cache.block_size`` (see
    :func:`_reject_chunk_prefill_block_size_mismatch`) — enforced here,
    not just documented, since a driver that skips
    ``engine_capture.build_engine`` can otherwise silently reintroduce a
    degenerate (always-zero-drift) search objective.
bundle
    Stage-1 (or scale-tuned) calibration bundle with ``centroids`` +
    ``per_channel_scales`` + ``head_dim`` + ``num_kv_heads``.
token_rows
    List of token-id lists (the calibration basket). One prefill each →
    one reference logit row each.
device
    Tensor device for the codec round-trip + KL (``cuda:0``).
topk
    Top-k vocab restriction for the KL numerator (64 = established).
driver
    The multi-rank ``DistributedEngineDriver`` (``calibrate_then_swap``'s
    ``swap_target``) at ``tp_size > 1``, else ``None``. When set, every
    candidate inject-forward is bracketed by a ``drift_seam`` install/clear
    broadcast on the driver's FIFO control ring so the worker ranks quantise
    their kv-head shard too; the seam's own maybe_apply stays rank-local.
    ``driver is None`` (single-rank, or EP-only ``tp_size==1``) means no
    new broadcasts and behaviour is byte-identical to before — the whole
    attention output is rank-0's, so rank-0's local seam is complete.

Adapts an :class:`EngineDriftSource` to the FULL
:class:`tkv.calibration.drift_source.DriftBackend` protocol
(``probe``/``measure_full``/``promote``/``reprime``/``invalidate_codecs``/
``n_chunks``, plus ``measure_components`` for the measure-abs path) so
the SOLVER's hot search loop walks on decode-accumulated drift instead
of the cheap prefill-chunk-boundary drift ``EngineDriftSource`` itself
exposes through that same protocol.

NEXT chunk's prediction -- cheap, but not what real autoregressive
serving actually experiences. A controlled A/B on the SAME 8-prompt
basket, ranking all 12 (layer,side) candidates from the DISAGREE -- and a real
oracle-scored allocation from the prefill-chunk-driven solver came back
significantly WORSE than uniform K4V4 (mean KL +7.36%, 95% CI entirely
positive). This backend makes the search itself optimize the same
decode-accumulated signal the promotion gate
(:class:`DecodeAccumulatedRealVerifyBackend`) already uses to verify
candidates AFTER the fact, instead of two different, disagreeing
objectives at search-time vs verify-time.

Every method is mechanical -- forwards straight to
:meth:`EngineDriftSource.measure_decode` /
:meth:`EngineDriftSource.measure_components_decode` -- because
``EngineDriftSource`` has NO resume capability at all (the engine
forward is opaque; see its class docstring: "the engine cannot
resume-from-layer... so probe == measure... and promote is a no-op").
``changed_layers``/``capture``/``from_layer`` are accepted for protocol
conformance and ignored, exactly as ``EngineDriftSource``'s own
prefill-based protocol methods already do.

REQUIRES the caller to have already run
:meth:`EngineDriftSource.abuild_reference_decode` (sync:
:meth:`EngineDriftSource.build_reference_decode`) on ``src`` before
handing this adapter to the solver as its ``backend=``. This is NOT
checked eagerly at construction (the reference may legitimately be
built after the adapter object exists, before first use) -- every
method below delegates to ``measure_decode``/``measure_components_decode``,
which themselves raise ``RuntimeError`` loudly (not silently produce
garbage) if the decode reference was never built. Constructing this
adapter is cheap and side-effect-free; the guard lives where the real
work happens, matching this module's established fail-loud convention.

Adapts an :class:`EngineDriftSource` to turbo-attn's real-verify-gate
contract (``tkv.calibration.bit_allocation._promotion_gate``): an object
exposing ``measure_components(alloc)`` and ``n_chunks``, the same
protocol ``_holdout_gate.py``/``test_promotion_gate.py`` exercise.

The promotion gate's whole point is to catch cases where the search
objective's own proxy predicts a win that doesn't hold up on a more
faithful measurement. Wiring the SAME prefill-only ``EngineDriftSource``
the search already optimizes against as its own gate (turbo-attn's
``_orchestrator.py`` auto-reuse default for an injected backend) can't
catch that class of failure — comparing an allocation against itself on
the identical objective it was chosen to minimize is close to a no-op.
This adapter instead routes the gate through
:meth:`EngineDriftSource.measure_components_decode` — a genuinely
different, decode-accumulated measurement — so the gate can actually
disagree with the prefill-only search that produced the candidate.

Caller must have already run :meth:`EngineDriftSource.abuild_reference_decode`
(sync: :meth:`build_reference_decode`) on ``src`` before this adapter's
``measure_components`` is called.

Drive a coroutine to completion from a non-loop thread (the solver
runs synchronously on its own thread; the engine loop runs elsewhere).
Falls back to ``asyncio.run`` when no engine loop was supplied.

``asyncio.run`` (not ``asyncio.get_event_loop().run_until_complete``,
the previous fallback): the no-loop case is exactly "plain sync code,
no ambient loop at all" -- ``asyncio.run`` is the version that
actually handles that without depending on thread-local event-loop
state some prior, unrelated async work may have already torn down.
``get_event_loop()`` is a real, live footgun here: it silently
creates an implicit loop the first time (masking the missing
``loop=``), then hard-raises ``RuntimeError: There is no current
event loop`` on a LATER call in the same thread/process once
anything else has closed the thread's event loop in the meantime —
e.g. any pytest run where an earlier async test already ran and
cleaned up (live-caught: CPU CI's full suite failed here, isolated
single-file runs didn't, because isolation is exactly what hid it).

Hashable fingerprint of the per-channel scales ``lid``'s codec would
be built from right now.

Part of the codec cache key because ``_build_layer_codec`` reads
``bundle["per_channel_scales"]["k"|"v"][lid]`` at build time and bakes
it into the codec via ``set_per_channel_scales``. Stage 2 tunes by
mutating those vectors in place between measurements, so a key of
``(lid, k_bits, v_bits)`` alone made every scale point after the first
reuse the codec built from the previous one — the sweep compared a
value against itself and every point tied, so ``best_s`` fell out of a
tie at the neutral 1.0 and stage 2 silently produced no information.

Keyed on the scales themselves rather than a version counter so
correctness does not depend on a caller remembering to call
``invalidate_codecs`` after each mutation: a cache must be keyed by
everything the cached object was built from.

The cache key for ``lid``'s codec at ``(kb, vb)`` as of now.

Single source of truth for the key shape — :meth:`_get_codec` stores
under it and :meth:`_evict_codecs_except` matches on it, so the two can
never disagree about what identifies a cached codec.

Drop cached codecs not used by ``alloc``.

Each codec holds GPU-resident centroids/boundaries/scales. The solver
probe path measures many distinct allocations without an
``invalidate_codecs`` between them, so without eviction the cache grows
one (layer, k_bits, v_bits) entry per explored combination and never
frees — an unbounded GPU accumulation across candidates. Keep only the
keys the just-measured alloc needs; the next forward rebuilds anything
evicted (cheap).

Matches the full key including the scale fingerprint (see
:meth:`_scale_key`), so a superseded scale's entry for a (layer, bits)
combination still in ``alloc`` is dropped too. Matching only the
``(lid, k_bits, v_bits)`` prefix would keep one stale GPU-resident
codec per scale point visited — reintroducing, along the scale axis,
exactly the unbounded accumulation this method exists to prevent.

TP>1 only: ship the stage-1 bundle (centroids + current scales) to
the worker ranks once, at drift-stage start, so each per-measure
``install`` need only carry the small mutated scale vectors.

Rank 0 writes the bundle to a temp path and broadcasts a ``load_bundle``
``drift_seam`` op; each worker loads it rank-locally from disk (the
swap's rank-local load pattern — not ``dist.broadcast`` on the bundle).
Returns the temp path (caller deletes it after the drift stage) or
``None`` when there is no driver (single-rank / EP-only tp1 ⇒ no-op).

FIFO guarantees every worker dequeues + loads this before the first
``install`` op, so the temp file is still present when it is read.

Broadcast an ``install`` so every worker arms its shard's seam.

No-op when there is no driver (single-rank / EP-only tp1) — the
TP1-byte-identical guarantee. Enqueued on the same FIFO ring as the
forward batches, from the engine loop thread, so it lands strictly
before the forwards it must bracket.

Broadcast an ``install_fp8`` so every worker arms the parameterless
fp8-e4m3 seam on ITS kv-head shard (rank symmetry — a rank-0-only
fp8 seam would measure ≈ ``1/tp_size`` of the real drift). No-op with no
driver (single-rank / EP-only tp1); enqueued on the SAME FIFO ring as
the forwards it must bracket, so it lands strictly before them.

Like :meth:`_aprefill` but continued past the prompt: runs the
basket's prefill AND ``decode_steps`` real decode steps per row.
Returns ``(N, decode_steps + 1, vocab)`` fp32 CPU. Same seam
install/inject/clear bracketing as :meth:`_aprefill` — the codec
stays armed for the WHOLE decode continuation, not just the prefill,
since the point is to see drift accumulate over real decode reads of
the (candidate-codec) KV cache, not just the one post-prompt logit.

Run the basket once with the parameterless fp8-e4m3 seam armed on
every rank; return ``(N, vocab)`` fp32 CPU last-position logits.

The fp8-KV drift arm: bracket the forwards with a worker
``install_fp8`` / ``clear`` (no-op at tp1) so the workers' kv-head shard
is fp8-round-tripped too, and arm rank-0's local ``kv_seam.fp8_inject``
for the same forwards. One FIFO ring ⇒ install_fp8 → forwards → clear in
order.

Establish the seam-off reference logits and the self-jitter floor.

The jitter floor is the run-to-run noise of the engine's own AWQ
forward — two clean prefills of the same basket. The HF path induces
jitter via a batch-size delta (different cuBLAS reduction order); the
arbi-serve prefill is one-sequence-per-step and eager, so two clean
passes are bit-identical and the floor genuinely collapses to 0.

That is not a bug — it means the engine has no meaningful self-jitter
in this mode, so the natural objective is the absolute top-k KL drift
(nats), not a ratio. We detect the collapse and set the per-prompt
denominator to 1.0 nat, making ``measure()`` return absolute drift in
nats (× a fixed unit). When a future stochastic forward does produce a
floor (> ``_JITTER_COLLAPSE_EPS`` mean), we use the ratio exactly as
the solver's ``_drift_from_logits`` does — same metric, same units.

Top-k KL drift of a naive fp8-e4m3 KV cache vs the seam-off reference.

The fp8-KV baseline arm: run the SAME basket through the parameterless
fp8 seam (no bundle/codec) and score it with the SAME ``_topk_kl_rows``
/ ``topk`` / jitter handling as :meth:`ameasure`, so the returned nats
are directly comparable to the fitted-codec drift (e.g. tkv k4v4). With
the deterministic engine's jitter collapsed the denominator is 1.0, so
this returns ABSOLUTE nats. Log-only — never written into a bundle.

Like :meth:`abuild_reference` but continued past the prompt:
establishes the seam-off reference logits AND self-jitter floor
over ``decode_steps`` real decode positions per prompt, not just
the one post-prefill position.

A decode-accumulated candidate measurement needs a decode-
accumulated reference -- comparing it against the prefill-only
``self._ref_logits`` would score every decode position against the
wrong (earlier, single-position) reference. Kept as a separate
reference from :meth:`abuild_reference` rather than replacing it,
so the cheap prefill-only path this class's other callers (the
solver's hot loop) use stays exactly as fast as before.

Decode-accumulated counterpart to :meth:`ameasure_components`:
per-prompt ``(mean_drift_nats, mean_jitter_nats, n_tokens)`` where
both means are taken over ``decode_steps + 1`` real positions
(prefill's end-of-prompt logit plus every decode step) and
``n_tokens = decode_steps + 1`` -- so a token-weighted pooling
(e.g. turbo-attn's promotion-gate ``_pooled_diff_mbits``) weights
this measurement by how many real decode positions it actually
covers, the same contract shape :meth:`ameasure_components` uses
with ``n_tokens=1``. Requires :meth:`abuild_reference_decode`.

G4: full AWQ-faithful TKV calibration through the arbi-serve engine.

Runs the drift-greedy byte_budget_table solver against the *served AWQ*
model via the engine inject-forward drift backend
(:class:`arbi_serve.calibration.engine_drift.EngineDriftSource`), starting
from a centroids bundle produced by ``engine_capture`` (stage 1).
The result is a complete **deployable bundle with a byte_budget_table** fit on
the exact weights we serve.

Architecture
------------
The engine runs an asyncio ``run_forever()`` loop on the main thread. The
turbo-attn solver is synchronous, so it runs on a worker thread;
``EngineDriftSource`` schedules each forward onto the engine loop via
``run_coroutine_threadsafe`` (the ``loop=`` it was handed). The reference +
jitter floor are built on the engine loop before the worker starts.

There is one turbo-attn solver — this driver only injects the engine drift
backend. No HF / transformers forward is involved.

Prefill only (MTP off) — same engine config as the capture driver. Runs
with cudagraph on by default — real-verified end to end (see
``engine_capture.build_engine``'s docstring for the full story with real
numbers): the per-candidate ``kv_seam``-injected measurement forwards
correctly force eager regardless of this flag
(``model_runner.is_injecting()``), and stage-1 K/V capture is
cudagraph-replay-safe (a raw-KV tap, not the old live-hook seam) — so
enabling it is a real, measured speedup with no correctness cost.

Synchronous body — runs on the worker thread. Builds the byte-budget
table against the engine drift backend. Returns the solver meta.

Callers that need no table (uniform serving reads only
``bundle["centroids"]``) do not call this at all.

``real_verify_decode_steps > 0`` arms turbo-attn's real-kernel promotion
gate (``tkv.calibration.bit_allocation._promotion_gate``, needs
``TKV_SOLVER_REAL_VERIFY_GATE=1`` too) against a decode-accumulated
:class:`~arbi_serve.calibration.engine_drift.DecodeAccumulatedRealVerifyBackend`
built HERE (not by the caller): ``EngineDriftSource``'s sync methods
dispatch via ``run_coroutine_threadsafe`` onto the engine loop and block
the CALLING thread on the result, so building it must happen on this
worker thread (same as every other ``src.*`` sync call in this
function) — calling it from the engine-loop thread itself (e.g. right
after ``await src.abuild_reference()`` in the caller) would deadlock.
``0`` (default) is a real no-op: no extra decode-reference cost paid
when the caller isn't going to arm the gate anyway.

``search_backend`` is the object the SOLVER's own hot search loop calls
(``backend=`` of ``solve_byte_budget_table``). ``None`` (default) keeps
today's behaviour: the solver walks on ``src``'s own cheap
prefill-chunk-boundary drift. Pass a
:class:`~arbi_serve.calibration.engine_drift.DecodeAccumulatedSearchBackend`
wrapping ``src`` to make the search itself walk on decode-accumulated
drift instead — the caller must have already run
``src.build_reference_decode(decode_steps)`` (same requirement
``build_real_verify_backend`` has for the gate). This is independent of
``real_verify_decode_steps``/the promotion gate: the gate re-verifies an
already-chosen candidate against a more faithful signal, this changes
what the search itself optimizes while picking that candidate.

Build+arm the decode-accumulated real-verify backend for ``src``'s
basket, for wiring into :func:`_run_stages_sync`'s ``real_verify_backend``.

Synchronous (runs ``src.build_reference_decode`` via its own
``_run_sync`` dispatch, same as every other sync entry point on
``EngineDriftSource``) so callers on the solver's worker thread can call
this directly, same as they call ``src.measure`` etc. Building the
decode reference is a real extra cost (``decode_steps`` extra forwards
x 2 reference passes) — callers should only call this when the gate is
actually going to be armed (``TKV_SOLVER_REAL_VERIFY_GATE=1``).

Stamp solve-stage run params + corpus fingerprint + drift_validation.

drift_validation records the solver's per-rung final drift (the absolute
top-k KL the engine objective measured) and an explicit monotonicity flag,
so the bundle carries its own proof it is non-degenerate before deploy.

Mark the bundle deployable.

The bundle already carries the canonical schema from the centroids stage.
For smart-mix the byte_budget_table (per-layer k/v bits per BPE key) is what
makes it deployable — the load-time gate
``arbi_serve.backends.tkv_layer_bits._check_finalized`` refuses a smart-mix
bundle without ``finalized=True`` (unless ARBI_ALLOW_UNFINALIZED_CALIBRATION=1).

For uniform serving there is no byte-budget table: the codec reads
``bundle["centroids"]`` directly and never consults the smart-mix gate, so a
centroids+scales bundle is already deployable. Callers on that path pass
``require_byte_budget_table=False`` — we still stamp ``finalized`` + schema so
the bundle is a first-class deployable, just without the table.

Flip finalized=True on an already-solved bundle (no GPU, no re-solve).

Validates the bundle carries a non-empty byte_budget_table, calls
:func:`_finalize_v8`, and writes --out atomically. Used for a bundle solved
by an earlier run that never got the finalized flag.

Engine identity stamped into a TKV calibration bundle, and checked on load.

A KV-codec bundle is a set of Lloyd-Max centroids, boundaries and
per-channel scales FITTED to the K/V statistics one particular engine
produces. Applying one fitted for a different engine does not crash and
does not raise a shape error — the tensors are the right shape, the
numbers are simply wrong, and output quality degrades silently. That is
the failure this module exists to make impossible.

Torch-free: stdlib + the checkpoint's ``config.json``. The CLI resolves
bundles before any torch import, so the fingerprint must be computable
there too.

WHAT IS IN THE KEY (each field answers: does changing it change the
correct codec?):

``model_name``
    Checkpoint directory basename. Two fine-tunes of one base model can
    ship byte-identical configs and identical weight shapes; the
    directory name is then the only cheap thing that separates them.
``architectures``, ``num_hidden_layers``, ``num_attention_heads``,
``num_key_value_heads``, ``head_dim``, ``hidden_size``
    The shape and stack the codec was fitted against. A per-channel
    scale vector is indexed by head-dim channel and a per-layer table is
    indexed by layer, so these dims ARE the bundle's coordinate system.
``kv_lora_rank``, ``qk_rope_head_dim`` (MLA only)
    The MLA latent is the tensor an MLA bundle quantises; its width is
    the codec's channel count.
``weight_quant``
    Weight quantisation propagates its error into every activation, so
    the K/V distribution an AWQ-INT4 checkpoint emits is not the one its
    bf16 twin emits, and the centroids that minimise error on one do not
    minimise it on the other.
``dtype``
    The dtype the K/V tensors arrive in bounds their dynamic range,
    which is what the per-channel scales normalise.
``weight_bytes``
    Total size of the checkpoint's weight files. The one field that
    catches a checkpoint replaced IN PLACE under an unchanged path and
    an unchanged config.

WHAT IS DELIBERATELY NOT IN THE KEY (a key so wide that nothing ever
hits gets switched off, and then there is no guard at all):

served KV bit widths
    A calibration fits the whole 2..8 ladder, so one bundle is
    legitimately servable at any width. Keying on the served width would
    miss on every bundle. The width is already checked exactly and
    structurally: the ``k_bits_N`` / ``v_bits_M`` section the served spec
    names must be present, or the apply path raises.
``tp_size``
    The bundle is per-layer and per-head-dim-channel, and every rank
    holds the same head_dim, so each rank slices its own kv-head shard
    out of the same full-tensor table. Calibration is rank-0-driven and
    produces one bundle for all topologies by construction; keying on
    tp_size would refuse a bundle that is correct.
``max_context``
    The centroids describe the per-token distribution of K/V vectors, a
    per-token property. A longer window draws more samples from the same
    distribution, it does not change the distribution. max_context is
    also a routine operational knob that ``auto`` resolves differently on
    different cards, so keying on it would invalidate bundles on a pool
    resize.
activation quantisation (W4A8 vs W4A16)
    Perturbs activations at fp8 granularity, well below the resolution
    the centroid fit can express, and is a runtime knob rather than part
    of the checkpoint's identity.
kernel / tkv library version
    A kernel consumes the centroids, it does not change which centroids
    are correct. The two changes that DO reinterpret the stored bits —
    a schema change and a codec-family (``lloyd`` vs ``vq2``) change —
    already have exact guards (``schema_version`` and
    ``_guard_vq_transition``).

KNOWN LIMIT, stated rather than papered over: a fine-tune that keeps the
base model's config, weight sizes AND directory name is indistinguishable
from its base here. Separating those needs a hash of the weights
themselves, which is a multi-terabyte read at every boot.

Fingerprint the engine that ``model_path`` will boot, or None.

None means the checkpoint could not be read (no ``config.json``, a
remote id that has not been materialised yet). Callers treat that as
"cannot verify", never as "verified".

Fields where the two disagree, as ``(field, bundle_value, engine_value)``.

Compared over the engine's fields plus any KNOWN field the bundle
carries that the engine does not; unknown bundle fields (a newer
producer) are ignored.

Record the producing engine's identity in ``bundle``.

Refuses to overwrite a DIFFERENT existing stamp: a multi-stage
pipeline that captured against one checkpoint and solved against
another has already produced a bundle no engine should serve.

Compare an unstamped bundle's own header against ``engine_fp``.

Returns ``(diffs, compared)`` — the differing fields in
:func:`fingerprint_diff`'s shape, and the names of the engine fields that
were actually compared. ``compared`` is the evidence a caller needs to say
what it checked; an empty list means the bundle carried nothing comparable
and only then is "unverified" the truthful word.

Refuse ``bundle`` if it was fitted for an engine other than this one.

Three outcomes:

  * stamped and equal — silent, the bundle is verified;
  * stamped and different — ``ValueError`` naming every differing
    field and both values;
  * unstamped, header agrees — load. The bundle's own header names the
    model, head_dim and kv-head count it was fitted for, so it is
    verified on those and the load says which fields it compared and
    which the missing stamp left out. Every bundle on disk today is in
    this class, so refusing here would break every running deployment
    on its next boot;
  * unstamped, header DISAGREES — ``ValueError``. Same silent quality
    loss a mismatched stamp catches, found without one.

``require=True`` (``ARBI_KV_BUNDLE_REQUIRE_FINGERPRINT=1``) promotes any
missing stamp to a refusal for operators who want the guard absolute.

``require`` defaults to the env flag.

Non-raising form for AUTO-DISCOVERY: may this bundle be picked up?

Returns ``(accept, reason)``. A stamped bundle is accepted only on an
EXACT fingerprint match — discovery is a guess made from a filename,
so a bundle that names itself for another engine is skipped rather
than served. An unstamped bundle is accepted (``reason`` says it is
unverified) unless ``ARBI_KV_BUNDLE_REQUIRE_FINGERPRINT=1``.

In-process TKV calibration against a live engine (no child, no re-load).

The offline ``--tkv-calibrate`` CLI (:mod:`engine_capture` + :mod:`engine_solve`)
each **build their own** single-purpose engine, calibrate, and shut it down —
so the model is loaded once per stage and, for auto-cal, a *second* time by the
serving parent. For a model too large to fit on one GPU (a native-EP MoE like
Qwen3.6-35B-A3B), single-process TP1 calibration is impossible outright.

This module runs the identical capture → fit → drift-solve pipeline against an
**already-built serving engine** — booted in ``tkv-bypass`` / eager so capture
reads raw K/V — then the boot sequence swaps the attention backend to
``tkv-k4v4`` and resizes the KV pool for serving. One model load, any TP/EP
topology, no child process. The building blocks are reused verbatim from
:mod:`engine_capture` / :mod:`engine_solve`; only the orchestration differs:

  * no ``build_engine`` — the caller passes the live ``eng``;
  * no ``eng.shutdown()`` — the engine survives into serving;
  * calibration is **rank-0-driven**, not run in lockstep by every rank (the
    retired CLI's model). See "one calibration path" below.

One calibration path (single-rank and torchrun multi-rank/EP, identical code).
The retired ``arbi calibrate`` CLI ran on every rank in lockstep, so it could
use out-of-band collectives (``broadcast_object_list`` for the stage-1 handoff,
``barrier`` before the swap). The in-process path cannot and must not: worker
ranks are inside the serving loop (``DistributedEngineDriver.run_forever``),
dequeuing per-tick SPMD deltas — they never reach an out-of-band collective, so
one would deadlock rank 0 against them. Nothing needs one:

  * **forwards** — capture / drift both drive the engine through
    ``eng.asubmit``, so every forward is scheduled by rank 0 and broadcast to
    the workers as a ``SlateDelta`` they execute in lockstep. This is the
    ordinary serving path; calibration traffic is just traffic.
  * **capture** — centroids are fit per ``(layer, head_dim)`` and head_dim is
    rank-invariant, so rank-0's capture is a complete, deployable stage-1
    bundle at any topology (see the TP note in ``engine_capture``). No gather.
  * **capture → solve handoff** — both run in the same rank-0 process on the
    same in-memory ``bundle``. No broadcast.
  * **the swap** — ``aswap_attention_backend`` is already shipped to the
    workers as a ``ControlOp`` that they apply rank-symmetrically, and it
    carries ``calibration_path`` so each rank loads the bundle rank-locally.
    That op is the barrier — no extra one is needed.

``calibrate_engine_inplace`` assumes ``eng`` is already running its
``run_forever`` loop and warmed — the caller (the boot sequence) owns that
lifecycle, since the same engine goes on to serve.

True iff ``driver`` can broadcast the ``drift_seam`` control op.

The multi-rank ``DistributedEngineDriver`` exposes ``_broadcast_control``
(the FIFO shm ControlOp ring). A plain ``Engine`` (single-rank swap target)
does not — and does not need to, since single-rank attention is not
sharded. Duck-typed rather than an ``isinstance`` to keep this module free
of the driver import.

Fail loud if this topology's drift stage would silently under-measure.

Calibration runs rank-0-driven (see the module docstring). Two of the three
engine-touching stages are exact at any topology — the capture is complete
on rank 0 (rank-invariant head_dim), and the forwards are broadcast to the
workers as ordinary slate deltas. The drift stage is the exception: it
measures logits-under-codec by installing codecs into the process-local
``kv_seam``, so only rank-0's KV heads get the codec.

  * ``tp_size == 1`` (single-rank, or EP-only like ``--ep-size 2``):
    attention + KV are replicated on every rank, so rank-0's seam covers
    the whole attention output. Exact.
  * ``tp_size > 1`` with a drift_seam-capable ``driver``: each rank owns a
    kv-head shard, but rank 0 now broadcasts the seam install/clear to every
    worker over the FIFO control ring, so every rank quantises its own
    shard during the measure and the drift is measured whole. Allowed.
  * ``tp_size > 1`` with no such driver: rank-0's injected shard would be
    all-reduced with the workers' un-injected heads — the measured drift a
    fraction of the real one, so the solver's byte-budget allocations (and
    the fp8 comparison) would be chosen on a number that is quietly wrong.
    Refuse rather than write a bundle that looks fine and is not.

Log-only fp8-KV drift measurement, apples-to-apples with the tkv codec.

Measures the top-k-KL drift of a NAIVE fp8-e4m3 KV cache (no calibration,
no bundle — a bare ``x.to(float8_e4m3fn).to(x.dtype)`` round-trip) against
the SAME seam-off reference, on the SAME basket, in the SAME nats as the
fitted tkv codec — so the number sits directly next to the k4v4 drift. fp8
is NOT a servable KV mode; this writes NOTHING into the bundle. At ``tp>1``
the fp8 round-trip fires on EVERY rank's kv-head shard via the ``drift_seam``
``install_fp8`` broadcast (a rank-0-only seam would report ≈

``(n_seqs, per_seq)`` for the capture/drift basket at ``max_ctx``.

The pooled total is ``tkv``'s ``recommended_pooled_tokens()`` unless the
caller names its own ``num_prompts × target_tokens``; either way it is
delivered as ``n_seqs`` sequences of ``per_seq <= max_ctx`` tokens, because
the centroid statistics come from the token COUNT and the live KV of a short
sequence is a few pages. The pooled total is gated against the same K=8
tail-coverage floor ``tkv.calibration.calibrate_centroids.run_calibration``
enforces for the HF CLI: this path calls ``fit_centroids_from_captured_kv``
directly — below the core, past that check — so the gate is re-stated here or
an under-provisioned basket silently produces high-bit sections whose tail
centroids are placed from a few outlier scalars each.

The KV-pool sizing that precedes calibration reads the same shape, so the
pool the capture runs against is sized from the basket it will actually see.

``(k_widths, v_widths)`` already fit in ``bundle['centroids']`` — the set
of N for which every layer carries a ``k_bits_N`` / ``v_bits_N`` codebook
(a width is 'present' only if complete across layers). ``(set(), set())``
for a missing/empty bundle.

Read the auto-cal cache bundle at ``out_path``, or None if there is none.

The cache path is derived from the checkpoint's directory BASENAME alone,
so two different checkpoints that share a directory name collide on one
file. A cached bundle whose fingerprint names a different engine is
therefore REFUSED rather than used: serving it would apply another model's
centroids silently, and extending it would merge this model's freshly-fit
widths into another model's bundle. An unstamped (pre-fingerprint) cache
entry is used, with a warning naming what could not be checked.

Both readers of this cache — the CLI preamble that decides whether the
bundle is already complete enough to serve directly, and the boot path
that extends it incrementally — come through here, so the check cannot be
present on one and absent on the other.

Copy the freshly-fit ``k_bits_N`` / ``v_bits_N`` codebooks for ``k_widths``
/ ``v_widths`` from ``src`` into ``dst['centroids']`` per layer — leaving
existing widths and the width-independent ``per_channel_scales`` untouched.
This is the seamless incremental update: a later k6v4 boot fits only k6 and
appends it to the cached bundle.

Record which tensor the bundle was fitted on.

An MLA bundle carries the shared NoPE latent; a GQA bundle carries
per-head K and V. Nothing in the numbers distinguishes them, so the
apply path reads this stamp rather than inferring from lengths.

Run one long synchronous fit on a worker thread, on the engine's device.

THE HTTP SURFACE IS THE ONLY THING WATCHING THIS BOOT. The server binds
before the engine can serve so ``/health`` and ``/metrics`` answer through
a cold start — that is what the serving gate exists to make safe
(:mod:`arbi_serve.server.boot_gate`) — and the in-process calibration runs
inside the boot, on that same event loop. A fit called straight from here
holds the loop for its whole duration: the probes stop answering, and an
orchestrator watching liveness cannot tell a process fitting codebooks
from one that has wedged. The stage this file already runs off the loop
(the byte-budget solver) is the same shape and is there for the same
reason.

``engine_build_window`` because the scrape that can now reach the process
during the fit must not query the caching allocator while the fit is
churning it — the lease answers False and the sample is skipped, which is
a gap in a gauge rather than two threads meeting inside the allocator.

``torch.cuda``'s current device is THREAD-LOCAL, so the worker is handed
the engine's: a thread that never set one fits on device 0 whatever card
the engine is on.

Calibrate/extend a deployable TKV bundle against the live ``eng``.

``eng`` must be booted in a raw-K/V capture mode (``tkv-bypass``, eager) and
already running its ``run_forever`` loop. Incremental: ``existing_bundle`` is
the cached bundle (or ``None``). The needed centroid widths are diffed
against what the cache already has —

  * ``fit_full_ladder`` (the default) needs every width in
    :data:`_LADDER_BIT_WIDTHS`, so one capture yields a bundle re-servable at
    any width without re-calibrating. Stage-1 Lloyd is a small fraction of a
    calibration run, so the whole ladder costs a fraction of the capture it
    shares;
  * ``fit_full_ladder=False`` needs only the served widths — ``{N}`` (K) and
    ``{M}`` (V) for uniform kNvM. The cheap arm, for a fast test run;
  * **smart-mix** (``run_solver``) needs the full ladder + a
    ``byte_budget_table``, whatever ``fit_full_ladder`` says.

Which widths are SERVED is independent of this: the served codec reads the
``k_bits_N`` / ``v_bits_N`` sections named by the served spec and ignores the
rest of the ladder.

If everything needed is present it's a **cache hit** (no forwards). Otherwise
we capture once and fit only the missing widths; when the cache already holds
the width-independent ``per_channel_scales`` we just merge the new codebooks
in (a later k6v4 boot fits only k6). Smart-mix additionally runs the
byte-budget solver.

Rank-0-driven at every topology (see the module docstring): this runs only on
global rank 0, whose engine forwards are broadcast to the worker ranks as
ordinary slate deltas, and whose capture is already a complete bundle. It
issues no out-of-band collective, so it cannot deadlock against workers
sitting in the serving loop. Returns the in-memory bundle. Does not shut
``eng`` down — the caller swaps it to the served codec and serves.

Unmap the bypass member's KV before the swap builds its replacement.

The swap prepares the serving member BESIDE the active one, so without
this the card must hold two KV pools at once. It cannot: measured on a
4090 serving Qwen3.8-27B with a DFlash drafter, the prepared member was
left 22.5 MiB for a step whose smallest admissible slate costs 49.2 MiB,
and the boot ended either in a clean ``MemoryBudgetError`` (eager) or,
under cudagraphs, an out-of-bounds gather in a replayed verify graph.

Nothing needs that KV after capture. The corpus has been read, the bundle
is written to disk, and the new member starts from an empty cache — so
the pages are discarded rather than offloaded, which is the difference
between a free release and one that copies GiB to pinned host RAM.

What this gives up: the bypass backend can no longer keep serving if the
swap refuses. That fail-safe was costing more than it protected — it is
the reason the swap had no room, so it turned "degraded serving" into
"no serving at all". The bundle is on disk either way, and a boot that
reads it back serves in ~2 min against the 30+ min this path spends.

Single-rank only. Calibration is rank-0-driven while worker ranks sit in
the SPMD loop, so a rank-0-only unmap would desync the group; at
multi-rank the swap's own ControlOp remains the only thing that touches
every rank.

Boot-sequence hook: calibrate on the live bypass engine, then hot-swap the
attention backend to the served TKV codec — the whole in-process flow.

Precondition: ``eng`` is built + running in a raw-K/V capture backend
(``paged_kv:tkv-bypass``, eager) with its ``run_forever`` loop up but not yet
accepting traffic. Steps:

  1. :func:`calibrate_engine_inplace` — capture → fit missing widths →
     (smart-mix: solve) → write/extend the bundle at ``out_path`` (the discovery
     dir). Rank-0-driven; its forwards reach the workers as slate deltas.
  2. :func:`aswap_attention_backend` to ``served_spec`` (e.g.
     ``"paged_kv:tkv-k4v4"``) — builds the codec backend as a prepared
     variant, discovers the just-written bundle, sizes the tkv KV pool,
     drains, and swaps. It fails safe (bypass keeps serving) if the bundle
     can't load, so a bad calibration never bricks the boot.

After this returns the engine serves the calibrated codec on one model load,
no child process, at whatever TP/EP topology it booted with. Called only on
rank 0 — the worker ranks stay in their serving loop throughout, servicing
the calibration forwards and then the swap control op.

Process-global KV seam for offline TKV calibration.

Read once inside :meth:`_Qwen3_5AttentionBlock._attn_eager` — the eager
attention call that runs *outside* the captured cudagraph blocks (split-attn).
Placing the seam there (not in the compiled ``_pre_attn``) keeps the
production cost at a single ``_SEAM is None`` test that Dynamo specialises to
a constant and dead-code-eliminates: **zero production overhead when off.**

Two modes, installed only during an offline calibration run (which runs the
model eager, so the seam is never traced):

* **capture** — record ``(layer_idx, k, v)`` so turbo-attn can fit centroids
  on the exact post-norm/post-RoPE K and raw V the deployed model attends to.
* **inject** — replace ``(k, v)`` with a per-layer codec round-trip, so a
  forward yields *logits-under-codec* (the drift signal stages 2 & 3 minimise).

``k`` / ``v`` arrive shaped ``(N_tokens, num_kv_heads, head_dim)`` — under TP
``num_kv_heads`` is the per-rank shard (see ``qwen3_5.py`` ``// cfg.tp_size``);
both modes are per-head-vector so the seam is rank-local and TP-correct.

MLA layers enter through :func:`maybe_apply_mla` instead, from
``MLAAttentionBlock.forward``. They carry one shared ``(N_tokens,
kv_lora_rank)`` latent rather than per-head K and V, and it is presented to
both sides of the sink because MLA quantises K and V with one codec.

This module imports nothing from turbo-attn: codecs are opaque objects with a
``compress_k``/``decompress_k`` (and ``_v``) duck-interface, built by the
turbo-attn glue and handed in. arbi-serve stays decoupled from the codec.

Replaces K/V with a plain, parameterless fp8-e4m3 round-trip.

The measurement counterpart to :class:`_InjectSeam`: it mimics a naive
fp8-KV cache so the drift probe can report fp8-KV top-k-KL drift in the
SAME units, at the SAME seam, as the fitted-codec drift — WITHOUT any
calibration bundle, centroids, or codec (fp8 has none). fp8 is NOT a
servable KV mode; this seam exists only to inject the fp8 round-trip into
the calibration drift seam for an apples-to-apples number.

The round-trip is turbo-attn's canonical fp8 baseline
(``tkv/calibration/bit_allocation/_activation_cache.py``):
``x.to(float8_e4m3fn).to(x.dtype)`` — a bare cast, NO absmax/448 scaling
and NO per-vector scale. Every layer, K and V alike. Head-agnostic and
rank-local like ``_roundtrip`` (the cast is elementwise), so under TP each
rank round-trips its own kv-head shard.

Hot-path entry for the MLA shared-KV latent, ``(N_tokens, kv_lora_rank)``.

The post-``kv_a_layernorm`` NoPE latent is the only tensor an MLA layer
quantises; the RoPE tail is stored bf16 and has nothing to calibrate.
``_SEAM is None`` in production → constant-folded out.

Whether the active seam actually MUTATES K/V (``_InjectSeam`` /
``_Fp8Seam``), as opposed to ``_CaptureSeam`` (records, passes through
unchanged) or no seam at all.

Two call sites, both correctness gates rather than pure optimizations:

* ``_attn_eager`` uses this to gate the fresh-diagonal prefill path (see
  its call site): capture mode must stay a pure pass-through recorder —
  no reason to pay the extra forward — so this must be False for it.
* ``arbi_serve.runtime.model_runner._forward_or_replay``/``execute``
  use this to force the whole step eager, skipping every captured-
  graph replay rung. A captured graph replays the EXACT kernel
  sequence recorded at capture time (boot warmup, seam always off in
  production) — replaying one under injection silently re-runs the
  original, un-injected attention kernels, so the "candidate" logits
  come back identical to the seam-off reference regardless of which
  codec was installed. Live-verified: under ``cuda_graphs=True`` with
  this gate absent, the search's per-candidate prefill objective
  (``EngineDriftSource.ameasure``) came back exactly 0.0 for every one
  of 6 distinct allocations — a plausible-looking but completely wrong
  measurement, not a crash.

Activate the parameterless fp8-e4m3 KV round-trip for the duration.

The fp8 analogue of :func:`inject` — no codecs, no bundle. Used by the
fp8-KV drift probe to measure logits-under-fp8 against the seam-off
reference. Shares the reentrancy guard: only one seam active at a time.

Whether this forward must run EAGER because a seam hook lives in Python.

Every seam here -- capture, inject, fp8, and the OSCAR accumulator -- hangs
its hook off ``_attn_eager`` or the attention op's return. A captured CUDA
graph replays the exact kernel sequence recorded at capture time and runs
no Python, so under replay those hooks do not fire. What that costs
depends on the seam, and both costs are silent:

* INJECT / FP8 -- the replay re-runs the original, un-injected attention,
  so a candidate's logits come back identical to the seam-off reference
  whatever codec was installed. Measured: an allocation search returned
  exactly 0.0 drift for all 6 candidates.
* CAPTURE -- the recorder sees nothing, so a replayed bucket contributes
  no K/V at all. Calibration then has a bundle fit on whatever fraction
  ran eager.
* OSCAR -- the Gram accumulates from a fraction of the corpus, and the
  rotation is fit on that fraction.

This is ONE question -- "would a replay skip the hook?" -- and it is not
the question ``is_injecting`` answers. That one asks whether the seam
MUTATES K/V, which ``_attn_eager`` needs for its fresh-diagonal gate and
which is correctly False for a pass-through recorder. Using it here read
"capture keeps the captured-graph speedup", which is a claim about a
DIFFERENT property than the one it was standing in for.

SCOPE. ``capture()`` has no production caller today: stage-1 K/V capture
moved off this seam in #1398 to the model runner's raw-KV tap, which
reads the paged KV pool and is replay-safe by construction. The live
caller of this predicate is the OSCAR accumulator. The capture arm is a
guard against re-introduction, and is NOT the cause of "only 0/16 KV
layers captured" -- that was the MTP seed path forwarding without the
step's calibration taps at all (#1999).

Feed a forward's ``Q`` and attention output into an OSCAR Gram fit.

Holds ``tkv.calibration.oscar_rotation.OscarGramAccumulator`` rather than
reimplementing the accumulation: the offline fitter
(``turbo-attn scripts/compute_oscar_rotations.py``) drives the SAME
object, so the two routes cannot drift into fitting different bases from
the same statistics. arbi-serve owns only the capture.

WHY THIS ROUTE EXISTS AT ALL. The offline fitter loads its model through
``AutoModelForCausalLM``, which cannot read an exl3 checkpoint -- so for
the format this server actually serves, an in-process capture is not a
convenience, it is the only way to fit a rotation on the served weights.

Hot-path entry. ``_OSCAR_SEAM is None`` in production → folded out.

Called with the attention output the layer actually returns, AFTER the
kernel wrote it -- never with the shadow output of the inject path's
second pass, which is a codec artifact and not what reaches the residual
stream.

Accumulate OSCAR statistics for the duration of a calibration forward.

Independent of the K/V seam's reentrancy guard: a rotation fit observes
and mutates nothing, so it composes with a capture run rather than
competing with it. It must NOT be combined with an inject seam -- the
statistics would then describe attention over round-tripped K/V, i.e. the
basis would be fit on the error it exists to reduce.

Why a calibration bundle could not be read — stated once, for every gate.

Three places need to refuse a bundle path that is not there: the CLI gate that
runs before any torch import, the engine-build gate that runs after the weights
are on the GPU, and the live backend swap. They had one message between them,
on the LATEST of the three, so the cheapest refusal was also the least useful
one — and the two copies that did exist could drift.

Stdlib only, on purpose: the CLI gate's whole value is that it fires before
torch is imported, so anything it calls has to be importable without one.

``None`` when ``path`` is a readable file; else why it is not.

``flag`` is how THIS caller was asked for the path (``--calibration``, a
query parameter, an env var), so the remedy names the thing the reader
actually typed rather than a flag they never used.

The three failures are distinguished because their remedies are different:

* **parent directory missing** — the calibration volume is not mounted at
  all (an NFS mount that did not come back after a reboot). Nothing about
  the filename is wrong; listing siblings would print an empty list and
  send the reader hunting for a typo that isn't there.
* **file missing** — almost always a name that drifted (a date stamp, a
  bpe suffix), so the directory's actual contents ARE the answer.
* **present but unreadable** — a permission or a mount flag, and the file
  list would be a red herring.

Fit an OSCAR rotation IN PROCESS, from the weights this server serves.

turbo-attn ships an offline fitter
(``scripts/compute_oscar_rotations.py``) that drives the same accumulator
through an HF forward. This module exists because that route cannot reach the
checkpoint formats this server actually runs: the offline fitter loads through
``AutoModelForCausalLM``, and an exl3 checkpoint is not loadable that way at
all. For an exl3 deployment an in-process capture is not a convenience, it is
the only way to fit a basis on the served weights.

WHAT IT PRODUCES, and what it does NOT. This writes a rotation file and
returns its sha256. It does not fit codebooks. A rotation alone serves
nothing: ``tkv.oscar.check_bundle_rotation`` refuses a bundle whose
``oscar_rotation_sha256`` does not match the configured rotation, in BOTH
directions, because a vq2 codebook is only valid in the basis it was fit in.
So a fresh rotation invalidates every existing bundle rather than upgrading
it, and the centroids must be re-fit against it -- which is the step
:mod:`arbi_serve.calibration.inprocess` already knows how to do, once
``TKV_OSCAR_ROTATION`` names this file.

THE ORDER IS NOT NEGOTIABLE: rotation first, centroids second, both stamped
with the same sha. Fitting them the other way round produces a bundle whose
codebooks describe a basis that no longer exists.

``(kv_layer_ids, n_kv_heads, n_q_heads, head_dim)`` for a built engine.

Read off the LIVE model's ``layer_specs``, and divided by ``tp_size`` the
same way :class:`~arbi_serve.models.attn.AttentionBlock` divides them
(``num_heads = layer_spec.num_heads // cfg.tp_size``). The seam observes
the tensors that block produces, so the accumulator must be sized in the
RANK-LOCAL head counts; the config's are the unsharded ones and would
build per-head bases for heads this rank never sees.

The layer list still comes from ``_kv_layers_and_heads`` -- it owns the
filtering that excludes GDN/Mamba/MLP-only layers and collapses MLA to one
shared latent, and a second copy of that would drift.

The accumulator class, or a refusal naming what is missing.

arbi-serve floors turbo-attn at a PUBLISHED PyPI version, and
``tkv.calibration.oscar_rotation`` arrived after the current floor.
Without this the flag dies on a bare ImportError deep inside a boot that
has already loaded a 27B model -- a stack trace naming a module, with
nothing saying which version supplies it or that the floor is the problem.

An :class:`OscarGramAccumulator` sized for this engine.

Accumulates on the ENGINE's device by default. The Grams are small --
``n_kv * padded_dim^2`` fp64 per layer per side, tens of MiB for a 27B --
and the rotation fit runs while the KV pool is CALIBRATION-SCOPED, which
frees GiB rather than MiB, so the card has room at exactly this moment.

Keeping them on the host instead would move every observed ``q`` and
attention output across PCIe and then do fp64 GEMMs on CPU cores. Device
accumulation avoids the transfer entirely; the fp64 rate on a consumer
card is poor but still far above a host BLAS. ``device=`` remains
available for a card with genuinely no headroom -- the accumulation is
numerically identical either way, only the placement differs.

Run ``token_rows`` as prefills under the OSCAR seam; write the rotation.

Returns the file's sha256 -- the value that must be stamped as
``oscar_rotation_sha256`` into every bundle fit against it. It is returned
rather than merely written because a caller who forgets to record it has
produced exactly the un-pairable artifact that has already stranded one
27B bundle in this fleet.

The forward runs EAGER: :mod:`arbi_serve.runtime.model_runner` forces it,
because the seam lives in ``_attn_eager`` and a replayed cudagraph would
skip it silently -- under-sampling the Gram without any error.

Stamp the basis these codebooks were fit in, and verify the pairing.

THE STAMP IS WHY IN-PROCESS OSCAR CALIBRATION DID NOT WORK. Both bundle
validators call :func:`tkv.oscar.check_bundle_rotation`, which refuses in
BOTH directions -- a bundle carrying a sha with no rotation configured,
and a rotation configured with no sha on the bundle. Neither validator
stamped, so a bundle freshly fit UNDER a rotation failed its own
validation for lacking the sha of the rotation it had just been fit in.

One seam rather than a stamp at each validator: the two are a stamp and a
check of the same fact, and splitting them is how one site comes to stamp
while its neighbour, arriving later, only checks.

No rotation configured => nothing to stamp and nothing to check; the
bundle stays Hadamard-basis and ``check_bundle_rotation`` passes it.
Returns the sha stamped, or ``None``.

Fit the rotation, then make it the basis the NEXT codebook fit uses.

ONE seam for both boot paths -- the offline ``--calibrate-only`` generator
and the serve-time calibrate-then-swap. They differ in what happens
afterwards, not in what a rotation is or how it is installed, and a second
copy of the install is how one path comes to set the env while its
neighbour, arriving later, only fits.

The install is an ENVIRONMENT assignment because that is how tkv resolves
it: ``tkv_flags()`` keys its cache on a snapshot of every ``TKV_*`` var, so
this assignment forces its own re-parse and every later reader -- the
codebook fit, the bundle stamp, and the served codec after the swap --
sees the same rotation without being handed it.

Whether ``bundle`` was fit in the basis ``rotation_path`` defines.

The auto-cal cache check asks "does this bundle cover the width I want".
Under OSCAR that is not sufficient: a Hadamard-fit bundle covers every
width and is still unservable in an OSCAR basis, and serving it would be
REFUSED at load rather than silently wrong -- but refused after a boot,
which is a worse way to find out than not reusing the cache.

sha256 of a rotation file.

Computed here rather than delegating to turbo-attn's helper: this runs at
BOOTSTRAP, deciding whether a cached bundle is reusable, which is before
any engine exists and on installs whose turbo-attn predates the fitter
module entirely. A cache check that ImportErrors is a boot failure over a
question whose answer is a file hash.

Where this model's OSCAR rotation lives, for EVERY route that wants it.

One derivation, not two that happen to agree. The boot flag and the admin
flip both resolve the artifact through here: a console flip that looked in
a different place than ``--oscar`` would fit a second rotation beside the
first, and the second one invalidates the bundle paired to the first --
silently, because both files are individually valid.

Quality-gate a multi-bit calibration bundle (centroids).

Raises :class:`QualityGateFailed` when any
``(layer, side, bit_width)`` section has ``improvement_pct`` below
``min_improvement_pct`` and no fallback marker.

The upstream cal-emit guard (when enabled) zeroes out fitted
centroids that don't beat the Gaussian default — those sections
surface as ``improvement_pct == 0`` with a ``fallback_reason``
field. Those count as passing for this gate (the operator
explicitly accepted Gaussian fallback).

Multi-bit bundle is the canonical schema; single-bit JSONs are
accepted but their ``improvement_pct`` lives at a different path so
we walk both.

Decide what a boot can serve from a model's calibration store, and merge
back what it had to fit.

THE EXPECTATION THIS ENCODES. A store is never finished. The set of
(basis, quantizer, width) a model might be served at is open-ended, so an
absent calibration is a normal state and not a failure: the boot fits what is
missing, merges it into the same file, and serves. The next boot anywhere
finds it and starts in seconds. That is the whole contract -- an operator
names what they want served and waits out a delay ONCE per configuration,
rather than hunting for a file that happens to have been baked with the right
combination of flags.

WHY THE DECISION IS A FUNCTION AND NOT A GATE. The v42 shape had exactly one
answer to "this bundle does not cover your request": refuse the boot. Under a
container restart policy that is a restart loop rather than a message, and it
is the wrong answer anyway, because the missing work is well-defined and the
engine that is booting is the only thing in the fleet that can do it.
:func:`resolve` returns SERVE-or-FIT plus the work list, and the caller acts.

TWO LEVELS OF CHECK, DELIBERATELY. This runs at CLI-resolution time, before
any weights load, off the store's own recorded layer set. The authoritative
check is the engine's, which knows the real KV layer ids. Same policy, checked
early so a doomed configuration costs a JSON read instead of a model load --
the same reason :func:`arbi_serve.cli.kv_resolve._reject_unresolvable_explicit_tkv`
exists next to the deep engine gate.

The store module, or a refusal naming what is missing.

arbi-serve floors turbo-attn at a PUBLISHED PyPI version and
``tkv.runtime.calibration_store`` arrived after the current floor. Without
this the boot dies on a bare ImportError with nothing saying which version
supplies the module or that the floor is the problem.

The ONE file this model's calibrations live in.

Derived, not configured, and derived in one place: a console flip and a
boot flag that looked in different places would each fit a calibration the
other could not see, and the fleet would accumulate near-duplicate stores
that never converge.

What a boot should do about the calibration it asked for.

``action`` is ``"serve"`` when the store already covers the request and
``"fit"`` when it does not. ``missing`` is the work list in the second
case -- a list, not a count, because an operator reading a boot log needs
to see WHICH rungs are being fit to know whether the delay is one width or
a whole ladder.

SERVE this configuration from the store, or FIT the gap first.

``layer_ids`` defaults to the layer set the store itself records for this
calibration, which is what the CLI has before any model is loaded. The
engine passes the real ids, so the same policy is checked twice and the
authoritative check is the one that knows the model.

A store that does not exist yet is a FIT with the whole configuration
missing, not an error -- a first calibration of a new model needs no
separate initialisation step.

Fold a freshly-fit calibration into the model's store; returns
``(basis_key, calibration_key)``.

Read-modify-write under the store's own lock, so a second engine
calibrating the same model at the same time keeps its work. A calibration
costs a boot's worth of time; losing one to a last-writer-wins overwrite is
the expensive failure, not a torn file.

``basis_pre`` is the fitted rotation, or ``None`` for the Hadamard basis.
It goes IN the store rather than beside it: the artifact that lives beside
the file is the artifact that goes missing, and a bundle whose basis is
gone is unservable forever.

Which family these cells belong to, read off the cells themselves.

Derived rather than passed, because a caller that could name the family
independently of the data is a caller that can name it WRONG -- which is
the v42 defect this schema exists to remove. A cell carrying a group-2
codebook is vq2; one that does not is scalar Lloyd. Mixed input is refused
rather than resolved: the two families cannot share an entry, and guessing
which one the caller meant would file half the cells under a lie.

The capture-pool max-not-sum invariant, from the boot log.

The piecewise prefill sweep records every layer capture on ONE shared
mem-pool so the transient working sets OVERLAY: the pool's *reserved* footprint
stays ~FLAT as rungs are added (≈ the single largest capture) instead of
growing every rung toward the SUM of all captures.

The boot emits, once per rung (``cudagraph_admin.py``, boot-only, no hot-path
cost)::

    piecewise rung <N>: pool[capture.cudagraphs] allocated=<A> MiB reserved=<R> MiB (<G> graphs total)

This module parses those lines and decides the invariant. It is pure text →
verdict so it is unit-tested on CPU against saved SUM-shaped and MAX-shaped
logs; the live lane feeds it the real boot log.

Outcome of the max-not-sum check.

``active`` is False when the boot log carries no ``pool[capture.cudagraphs]``
rung lines at all — the telemetry is not present on this build, so
the invariant is unprovable and MUST NOT be reported as a pass or a fail.
The gate reports "inactive" and moves on; it activates the instant the
telemetry ships.

Decide the max-not-sum invariant from a piecewise boot log.

Healthy (max-not-sum): the reserved footprint is ~FLAT across rungs, so
``max(reserved) ≈ final(reserved)`` and the ratio of the largest reserved to
the smallest post-first reserved is within ``flat_tolerance``. A SUM
regression grows reserved every rung, so the ladder's reserved spans a large
multiple — ``growth_ratio`` (max / first-substantive) blows past tolerance.

``min_rungs``: with fewer than this many rungs the max-vs-sum distinction is
not observable (one capture's max IS its sum), so the check stays inactive
rather than pass vacuously.

Per-``model_type`` chat-template overrides.

Some checkpoints bundle a ``chat_template`` that is NOT the format the
model was trained to serve: Step-Audio-2 ships a generic Qwen
``<|im_start|>`` tool template, while its real conversation format —
``<|BOT|>{role}\n{content}<|EOT|>`` with ``user`` rendered as
``human`` and audio parts as ``<audio_start><audio_patch><audio_end>``
— is built in code by the reference clients. The tokenizer consults
:data:`CHAT_TEMPLATE_OVERRIDES` by the checkpoint's ``config.json``
``model_type`` and serves the override instead, so
``/v1/chat/completions`` renders the true training format.

The audio part emits ONE ``<audio_patch>`` marker inside its
start/end frame; the OpenAI glue's placeholder expansion
(:func:`arbi_serve.multimodal.openai.expand_audio_placeholders`)
replaces it with the per-chunk token runs.

NVIDIA NemotronLabs VoiceChat 11B is a THIRD, distinct case: its HF
repo (``/mnt/k8scache/models/NVIDIA-NemotronLabs-VoiceChat-11B``)
ships no ``chat_template`` at all — no ``chat_template`` field in any
``tokenizer_config.json`` (only ``rnnt_tokenizer/tokenizer_config.json``
exists there, and it configures the ASR sub-tokenizer, not chat), and
no sidecar ``chat_template.jinja``. Its own README points readers at
an EXTERNAL template in the NeMo-Speech reference repo instead
(``github.com/NVIDIA-NeMo/Speech``, branch ``nemotron-labs-voicechat``,
``examples/speechlm2/function_calling/template.jinja``) — that is the
real trained tool-calling format, reproduced verbatim below as
:data:`NEMOTRON_VOICECHAT_CHAT_TEMPLATE`.

Registry-key caveat: ``config.json`` for this checkpoint is a NeMo
Hydra training config (top-level keys ``data``/``exp_manager``/
``hf_export_dir``/``model``/``trainer``/``_rnnt_merge_info``), not a
standard HF config — it carries neither ``model_type`` nor
``architectures``, so :func:`chat_template_override` cannot resolve
anything from the raw checkpoint dir as shipped. ``"nemotron_voicechat"``
below is a placeholder key for the ``model_type`` the HF-style
``config.json`` (synthesized by the model loader when it exports/reads
this checkpoint — see ``arbi_serve/models/__init__.py``) is expected to
carry; the override becomes live the moment that loader work lands.

CLI entry point: ``arbi-serve`` console script + ``python -m arbi_serve``.

KV-cache codec is selected via ``--kv-cache-dtype``:

  --kv-cache-dtype tkv          # Turbo Attention codec; bit widths from env
                                 # TKV_BITS=4.0 + TKV_CALIBRATION_FILE (smart) | TKV_BITS=4 (uniform)
  --kv-cache-dtype auto          # default — bf16 paged KV (no codec)

Non-TKV state-kind backends still use the ``kind:name`` form
(``--backend`` repeatable):

  --backend paged_kv:tkv-bypass          # default — lossless bf16 on turbo-attn
                                         #   kernels (Turbo prefill + split-K decode),
                                         #   all head dims, capturable incl. MTP
  --backend mla_shared:mla       # MLA shared-KV (DeepSeek V2/V3)
  --backend mamba:mamba          # Mamba SSM
  --backend gdn:gdn              # Gated Delta Net

``--kv-cache-dtype`` and ``--backend`` compose: the TKV spec is appended
to the ``--backend`` list internally, so multi-backend hot-swap (bf16
co-resident with tkv) works the same way as before — just spell the
TKV slot via the dtype flag, not as a ``paged_kv:tkv-...`` literal.

``--tp-size`` and ``--ep-size`` are explicit; the constraint
``tp_size * ep_size == world_size`` is validated at parse. They are
orthogonal factors of the world size: ``--ep-size N`` gives N independent
TP replicas, each holding a distinct expert partition and each
replicating the attention weights + KV cache.

``--enable-expert-parallel`` (equivalently ``--moe-ep-size <tp-size>``)
is the other axis and the one that scales attention: it shards the MoE
experts across the ranks of a TP group, so attention and the dense
projections stay TP-sharded over all ``--tp-size`` ranks while each rank
still holds only its slice of the experts. On 2 GPUs,
``--tp-size 2 --enable-expert-parallel`` serves the same expert split as
``--ep-size 2`` with half the attention compute and half the KV bytes
per rank.

Subcommands:
  - (default, no subcommand) — boot the server.
  - ``calibrate`` — run calibration + optional bit-allocation solver,
    write a multi-bit bundle, exit. Same code path as ``--autotune``
    without booting the engine.
  - ``calibrate-mtp-k`` — sweep MTP K∈{0..5} once at deploy time, pin
    the winner to ``~/.cache/arbi-serve/mtp-k/<key>.json``. Subsequent
    boots auto-load the K when ``--mtp-n-draft`` is not passed.
  - ``warmup`` — build the engine, populate the Turbo prefill CuTeDSL JIT cache
    (and Triton kernels) with one forward, optionally capture
    CUDAGraphs, then exit. Idempotent when ``/cache`` is mounted.

This is a package; the genuine public surface (``main``, ``build_parser``,
``cfg_from_args``, …) is re-exported here so ``[project.scripts]
arbi-serve = "arbi_serve.cli:cli_entry"`` and every importer keep working.
``cli_entry`` is the PROCESS entry and ``main`` the plain function callers
and tests invoke in-process; only the former decides how the process
leaves (arbi_serve.calibration.child_exit). The
implementation is partitioned across torch-free submodules
(``parse_helpers``, ``kv_resolve``, ``parser``, ``config_builder``,
``calibrate``, ``warmup``, ``bootstrap``) so ``arbi-serve --help`` and the
OpenAPI dump stay engine-free.

``python -m arbi_serve.cli`` entrypoint.

Equivalent to ``python -m arbi_serve`` and the ``arbi-serve`` console
script: dispatches to :func:`arbi_serve.cli.main`. The CUDA allocator and
thread-cap env are set in :mod:`arbi_serve` (imported before this module
runs), so every entrypoint sees the same runtime configuration.

Apply env-var defaults that are universally helpful for arbi-serve.

Cap inductor compile threads to half the cores (idle-thread hygiene)
and bind torch's distributed comm to loopback when there's nothing
else to do. No HuggingFace offline default: hub access is online by
default (a hub-id ``--model`` downloads on first boot); exporting
``HF_HUB_OFFLINE=1`` opts a deployment out of all hub traffic.

Each var is only set when not already present, so the caller can
still override.

Seed ``TORCH_EXTENSIONS_DIR`` from the image's build-time bake.

The slim image bakes tkv split-K decode / MTP-verify / compress-store
kernels at build time (GPU-free nvcc, arch + shapes known) into
``/opt/cache-baked/torch_extensions``. tkv loads them through
``torch.utils.cpp_extension.load_inline``, which resolves under
``$TORCH_EXTENSIONS_DIR`` (= the version-keyed
``/cache/torch_extensions/<torch-py-cu>`` once /cache is present; the
seed runs with the image's own torch, so the key matches the runtime
that consumes the bake) — a different dir than the bake landed in, so
without this seed the runtime re-JITs every kernel at boot. The exl3 /
xgrammar prebake short-circuit (``_prebake_loader``) wraps ``ce.load``,
not ``ce.load_inline``, so it does not cover tkv — this seed does.

Copies only the per-kernel build dirs missing in the target, so it is
idempotent and cheap on warm boots / pre-warmed /cache volumes. Best
effort: a copy failure leaves the runtime to JIT, never blocks boot.

A tkv dir is copied only when the bake's recorded source fingerprint
agrees with the installed tkv (:mod:`arbi_serve.tkv_prebake`). The
extension NAME is this cache's entire key, so a bake carried across to a
differently-sourced tkv would serve machine code compiled from other
source with nothing on the path to notice -- and a production
``TKV_NO_JIT=1`` boot will not rebuild over it. Refusing to seed is the
safe direction: the runtime then builds from what is installed.

Gated by its OWN flag (``ARBI_SERVE_PREBAKE_SEED``), not by
``ARBI_SERVE_PREBAKE_DIR``. A tkv bind-mount needs this seed off, and
pointing the shared dir at a nonexistent path to achieve that also
disabled the exl3/xgrammar ``ce.load`` short-circuit above -- which a tkv
mount does not affect -- so every bind-mount boot re-JITed exl3 from the
image's own unchanged sources.

One boot line pinning which arbi_serve is running.

An editable install wins over PYTHONPATH (PEP-660 finder), so an
experiment overlaying a checkout can silently run different code
than intended. The path + git rev in the log make that ambiguity
impossible to miss.

``arbi-serve`` console-script entry: run :func:`main`, end the process.

Separate from :func:`main` because ``main`` is a plain function that
tests call in-process and that returns its code; only the PROCESS entry
may decide how the process leaves. Every failing outcome leaves without
interpreter finalization (see
:func:`arbi_serve.calibration.child_exit.exit_cli_process`).

Entry point for the ``arbi-serve`` CLI.

Dispatches the operational subcommands (warmup, bench, ...) before
argparse runs, then otherwise parses the serve form and boots the
server. Returns the process exit code.

Multi-rank entry under torchrun.

Every rank constructs and builds a :class:`DistributedEngineDriver`.
The build is COLLECTIVE — ``init_process_group`` runs inside
:class:`Engine.build` via :func:`init_distributed_environment` — so
all ranks rendezvous inside it before any rank starts serving
requests / waiting for plans. Where each rank runs it differs:

  - Rank 0 starts uvicorn FIRST and hands :func:`build_app` a driver
    FACTORY; the lifespan builds the driver on the boot thread while
    the HTTP surface already answers ``/health`` and ``/metrics``.
    A TP boot is therefore watchable from its first second, including
    a rendezvous stalled on a peer that has not arrived. Every
    per-step broadcast originates here.
  - Worker ranks 1..N-1 bind no HTTP, so they have nothing to bring
    up ahead of the build: they build inline, then run
    :meth:`DistributedEngineDriver.run_forever` directly under
    ``asyncio.run`` so their event loop services the broadcast
    receive + worker-side forward.

Return the active recipe for ``args`` (CLI-explicit or
auto-detected). Returns ``None`` if neither path resolves a recipe.

Auto-detection is GENERATION-scoped: recipes tune the generation
serving stack (KV page counts, MTP depth, codec calibration), and
the ``architectures:`` fallback match is family-coarse — e.g.
``Qwen3ForCausalLM`` matches Qwen3-Embedding/Qwen3-Reranker too, and
a 4B generation recipe's ``--num-pages`` then fails the memory
budget on a 0.6B embed boot. An embed / rerank instance therefore
skips auto-detection entirely; an EXPLICIT ``--recipe`` is still
honoured verbatim (operator override wins, as everywhere).

The merge into args is deferred to the caller so detection is
testable independently of the merge.

Resolve the model-sampling-default policy for ``--generation-config``.

Precedence: explicit ``--generation-config`` CLI value > the
``ARBI_GENERATION_CONFIG`` env knob (docker compose) > default
``"auto"`` (backfill request defaults from the model's
``generation_config.json``). Rejects an unknown value LOUDLY rather
than silently falling back.

Resolve the KV-cache codec spec and apply the default backend.

Mutates ``args`` in place: resolves the ``--kv-cache-dtype`` sentinel,
translates ``tkv`` into the internal spec, reconciles the calibration
source, and fills in the default attention backend when none was given.

Build the ParallelConfig, validating the TP/EP launch requirements.

Raises ``SystemExit`` when ``--tp-size``/``--ep-size`` > 1 without a
torchrun launch, when ``tp * ep`` disagrees with ``WORLD_SIZE``, or
when the MoE expert sharding does not partition the TP group.

Resolve ``--enable-expert-parallel`` / ``--moe-ep-size`` into the
number of expert shards within each TP group.

The two knobs name the same quantity at different altitudes: the flag
is the ergonomic "shard the experts over my TP ranks" phrasing, the
size is the explicit hybrid. They are reconciled here rather than
left for :class:`ParallelConfig` so a contradiction is a CLI error
naming both flags, not a dataclass ValueError naming neither.

Read through ``getattr`` with the parser's own defaults — this
module's convention for every optional flag (``mtp_n_draft``,
``cudagraph_max_shapes``, ``mtp_capture_full_k_ladder``, ~35 more).
``cfg_from_args`` accepts any Namespace, and several suites hand-build
a minimal one; requiring every such fixture to enumerate each new flag
would make adding a flag break unrelated tests. The defaults below
are the parser's, so a real CLI parse is unaffected.

Cross-validate the MTP flags and build the :class:`MtpConfig`.

Resolves the drafter source (bundled head / external draft model /
DFlash), the ``ARBI_MTP`` env override, and the tkv K-cap. Raises
``SystemExit`` on any inconsistent combination.

Whether ``--chunk-prefill`` carries a width the OPERATOR chose.

The test is "the value differs from the stock default", NOT "the flag
appeared on the command line", and the difference is load-bearing here:
every shipped compose file passes ``--chunk-prefill=${ARBI_CHUNK_PREFILL:-
2048}`` unconditionally, so flag PRESENCE is true for every containerized
boot and would read the stock deployment as an operator pin. A width equal
to the default is a width the engine would have picked anyway, so reading
it as the engine's costs nothing; a width that differs cannot have come
from anywhere but a deliberate ``--chunk-prefill`` / ``ARBI_CHUNK_PREFILL``.

The one case this under-reads — an operator who pins the default value
explicitly — resolves toward the engine-chosen behaviour, i.e. toward
what the boot does today. The ambiguity can only ever cost a refusal, never
cause one.

Resolve the final context / batch / token-budget sizing.

Handles the ``"auto"`` sentinel for ``max_context``, the flat per-step
token budget, and the embed/rerank-specific overrides (which also coerce
the backend list on ``args``). Mutates ``args`` for the embed/rerank
backend + default-backend override.

``--draft-vocab-tail`` as ``cfg.draft_vocab_tail``: ``derive`` or a width.

Parsed here rather than by ``type=int`` on the flag because ``0`` is a
distinct, meaningful setting — the plain prefix an A/B's control arm names
— so it cannot double as "unset" and the default cannot be an int.

Build a :class:`ServerConfig` from parsed serve-form CLI arguments.

Resolves the KV-cache dtype spec, applies backend defaults, validates
distributed (TP/EP) launch requirements, and maps each flag onto the
config. Raises ``SystemExit`` when ``--model`` is missing.

Resolve ``--model`` Hugging Face hub ids to local snapshot dirs.

A hub id (``org/name`` or ``org/name@revision``) resolves to a local
snapshot directory under the HF cache — served from cache when present,
downloaded once otherwise. A local path passes through untouched.
Everything downstream of this module only ever sees a local model
directory. Hub access is online by default; exporting
``HF_HUB_OFFLINE=1`` opts a deployment out of all hub traffic (cached
hub ids still resolve, an uncached one refuses loudly).

Torch-free; ``huggingface_hub`` is imported lazily so parser construction
and ``--help`` stay light.

True when ``model`` reads as a hub id rather than a local path.

Requires exactly ``org/name`` (optionally ``@revision``) in the HF
repo-name alphabet AND that no such path exists on disk — an existing
one-slash relative directory stays a local path.

Resolve ``--model`` to ``(local_dir, display_name)``.

A local path returns ``(model, None)`` untouched. A hub id returns the
local snapshot dir plus a display name (the repo name segment, with the
``@revision`` suffix when given) for ``--served-name`` defaulting:
cache hit first (works offline), else a one-time download of the
servable files into the HF cache. Refuses loudly — with the fix — when
the download is needed but blocked (offline mode, auth, missing repo).

Resolve ``args.model`` in place before anything downstream reads it.

A hub id is rewritten to its local snapshot dir, and ``--served-name``
(when the entrypoint has one, unset) defaults to a stable public name: the
hub-derived name for a hub id, or the final path component for a local
checkpoint. This keeps container/host paths out of ``/v1/models`` and the
residency routing keys. Never mutates the environment: ``HF_HUB_OFFLINE``
stays exactly what the operator exported.

Resolve the SPEC-NAME ``(k_bits, v_bits)`` for ``--kv-cache-dtype tkv``.

Single user-facing knob: ``TKV_BITS``. Two forms:

  * Float, e.g. ``"4.0"``: smart-mix (TRUE per-layer). Reads
    ``TKV_CALIBRATION_FILE`` (deployable bundle), looks up
    ``byte_budget_table[<TKV_BITS>].layers`` and returns the
    ``(max_k, max_v)`` across the per-layer table. This pair is used
    ONLY to name the internal ``paged_kv:tkv-k{K}v{V}`` spec and as
    the backend's fallback width; the engine applies the genuine
    per-layer (k_bits, v_bits) per attention layer (sizing each
    layer's cache slab + codec + kernel dispatch to its own bits) by
    re-reading the same env + bundle inside ``TkvBackend`` (see
    ``arbi_serve.backends.tkv_layer_bits``).
  * Int, e.g. ``"4"``: uniform K=V=N (kept for unit-test / CI smoke
    runs without a calibration bundle; production should use float).

Per-side bits may be any integer in ``[2, 8]`` — the codec + kernels
handle arbitrary per-layer-varying widths in that range.

True when tkv is explicitly selected but has NO calibration source.

The prepare-phase auto-cal fires only for an EXPLICIT
``--kv-cache-dtype tkv`` (or recipe-selected tkv) with no bundle —
never for the unset/auto-discover path (that intentionally falls back
to bf16). ``--no-calibration`` and ``TKV_CALIBRATION_FILE`` both
opt out (the operator chose uncalibrated / pointed at a bundle).

argparse type-validator for ``--backend``.

Rejects the deprecated ``paged_kv:tkv-k{K}v{V}`` codec form at parse
time so the operator sees a clear migration path. Internal spec
strings (used by tests / engine consumers that build a Namespace
directly) are unaffected because they bypass argparse.

``paged_kv:tkv-bypass`` (raw bf16, no codec) shares the ``tkv-``
brand prefix but is a plain backend name, not a codec declaration —
it stays CLI-addressable. The codec grammar is precisely
``tkv-k<K>v<V>`` (see ``backends._parse_tkv``), so ``tkv-k`` is the
exact rejection signature.

Return the dir to scan for auto-discovered TKV bundles, or None.

``ARBI_SERVE_CALIBRATION_DIR`` (single-sourced via runtime_flags):
unset → ``~/.cache/arbi-serve/calibrations``; a literal empty string
disables auto-discovery. Cf. the grammar / budget cache dirs for the
same unset-vs-empty convention.

Dirs scanned to auto-discover a TKV bundle, in precedence order.

1. The curated discovery mount (``ARBI_SERVE_CALIBRATION_DIR`` / ``/cal``)
   — operator-supplied bundles.
2. The persistent autotune cache (``ARBI_SERVE_CALIBRATION_CACHE`` /
   ``/cache/calibrations``) — bundles a prior auto-gen wrote. Covering it
   here means a bundle already generated on a previous boot is FOUND on the
   next boot instead of re-triggering auto-gen (and, at TP>1, instead of
   the no-bundle fail-loud).

A literal empty ``ARBI_SERVE_CALIBRATION_DIR`` disables auto-discovery
entirely (the curated dir AND the cache), preserving the explicit opt-out.

Return calibration JSONs across the search dirs whose name carries the
model basename (case-insensitive), ordered by (dir-priority, newest-first).

The bits filter is applied later so we only resolve TKV_BITS (which imports
tkv) when a name even matches. The ordering lets discovery prefer a
curated-mount bundle over a cache-dir (auto-generated) one regardless of
mtime.

Pick the best TKV bundle from ``candidates`` for the declared bits.

A candidate is eligible when its name carries the ``k{K}v{V}`` tag
(case-insensitive); e.g. model ``Qwen3.5-0.8B`` at k4v4 considers
``qwen3.5-0.8b_k4v4_pcsfix_20260714.json``. ``candidates`` arrives
pre-ordered by (dir-priority, newest-first), so the first ACCEPTED
candidate wins — a curated-mount bundle beats a cache-dir one, and within
a dir the most recent.

Names only narrow; they never decide. A bundle stamped with a
:mod:`arbi_serve.calibration.fingerprint` block is accepted only on an
EXACT match against the engine this model boots — a filename is a claim,
and the calibration a wrong bundle applies is not a crash but a silent
quality loss. A stamped bundle for a different engine is skipped (loudly)
and the search continues; an unstamped bundle is accepted with a warning
unless ``ARBI_KV_BUNDLE_REQUIRE_FINGERPRINT=1``. Returns the bundle path,
or None when nothing is eligible and accepted.

Parse ``path``, failing loud if it is not a current-schema TKV bundle.

Surfaces a clear SystemExit at the CLI boundary rather than letting a
wrong file crash deep in codec install. Mirrors the schema-version +
top-level-key checks ``tkv.runtime.calibration.load_calibration_file``
enforces. Returns the parsed bundle so the caller can inspect it without
a second read.

Make ``--model X`` default to calibrated tkv, hands-free.

Runs when ``--kv-cache-dtype`` was NOT passed (sentinel ``None``).
Auto-discovers a calibration bundle for the model + declared bits; if
found, defaults the codec to ``tkv`` and pins the discovered bundle
as ``args.calibration`` (so the engine applies its centroids). If no
bundle is found we fall back to bf16 with a loud one-line WARNING —
never silent UNcalibrated tkv, which would be a quiet quality loss.

Precedence (all left untouched here):

  * explicit ``--kv-cache-dtype auto`` / ``tkv`` (dtype not None);
  * an explicit ``--calibration`` / ``--autotune`` result already on
    ``args.calibration``;
  * ``TKV_CALIBRATION_FILE`` set in the env.

Unify the one calibration-bundle source for active TKV.

Kills the footgun where a bundle could be pointed at via two
disconnected channels that had to agree: ``--calibration`` (read by
the apply path + the boot reject gate) and ``TKV_CALIBRATION_FILE``
(read by the smart-mix per-layer-bits resolver). Passing only the env
+ an explicit ``--kv-cache-dtype tkv`` is rejected at boot ("no
calibration") because the env never reaches ``args.calibration``.

After the dtype is resolved, when tkv is active we make the two
agree — one bundle, set once, however the operator named it:

  * both set to different files  -> hard SystemExit (loud, no guessing);
  * only ``--calibration``       -> export ``TKV_CALIBRATION_FILE`` from it
                                    (so smart-mix bit resolution sees it);
  * only ``TKV_CALIBRATION_FILE``-> pin ``args.calibration`` from it
                                    (so the gate + apply path see it);
  * neither                       -> untouched (the reject gate /
                                    ``--no-calibration`` / ``--autotune``
                                    path decides).

No-op unless ``--kv-cache-dtype tkv`` is active. Idempotent.

Hard-fail at CLI-resolution time when explicit tkv has no way to be calibrated.

``_reject_uncalibrated_tkv_without_optin`` (``engine/build_config_resolve.py``,
enforced from ``build_phases_load.py`` deep inside ``Engine.build()``) is the
authoritative, always-on gate: it refuses to activate an uncalibrated TKV
backend. But it fires AFTER the model's weights are already loaded onto the
GPU — a doomed boot wastes minutes and VRAM before the operator finds out.
This is the SAME policy, checked here at CLI-resolution time (before any
torch/GPU work happens) so the refusal is immediate and precise instead of
arriving after a full model load. Torch-free: only stdlib + the same
name-matching scan ``_resolve_kv_cache_default`` already does.

Must run AFTER :func:`_resolve_kv_cache_dtype` (needs the resolved bits to
match the auto-discovery ``k{K}v{V}`` tag) and AFTER
:func:`reconcile_tkv_calibration_source` (needs ``args.calibration`` to
already reflect a ``TKV_CALIBRATION_FILE``-only source).

No-op when:

  * ``--kv-cache-dtype`` isn't ``tkv`` (nothing requested — the ``auto``/
    unset-default bf16 path is untouched, by design);
  * a calibration source already resolves (``--calibration``,
    ``--calibration-for``, or ``TKV_CALIBRATION_FILE``);
  * ``--no-calibration`` (explicit opt-in to reduced-quality uncalibrated
    tkv — the same escape hatch the deep engine-build gate honors);
  * a deferred in-process auto-calibration is already staged
    (``args._inprocess_calibration`` — the bootstrap preamble's
    boot-bypass-then-calibrate-then-swap plan for a bare
    ``--kv-cache-dtype tkv`` with no bundle; it fails loud on its own if
    the calibration attempt itself fails, so this gate would only
    duplicate that check less precisely).

Otherwise: try the SAME name-matching auto-discovery
:func:`_resolve_kv_cache_default` uses (so an operator who happens to have
a correctly-named bundle sitting in the discovery dir gets it applied
exactly like the hands-free default path would) — and only when that also
comes up empty, refuse the boot with a precise, actionable message.

Hard-fail when a bundle the operator NAMED is not a readable file.

The gate above proves a calibration source RESOLVES; it never opened it.
``--calibration=/cal/typo.json`` therefore satisfied every CLI check and
died at ``build_phases_load._load_calibration_json`` -- which runs after
the weights are on the GPU, after the warm-cache assertion and after the
auto-max-context solve. On a 27B exl3 checkpoint that is the entire
weight-read + JIT wall (the compose recipe budgets 1200s for it) spent
before a one-line ``stat`` fails. The typo is knowable before torch is
imported, so it is refused before torch is imported.

A NAMED path only. Auto-discovery already stats what it finds, and this
gate must not turn "nothing configured" into an error -- that is the gate
above's job, with its four numbered remedies.

The deep gate stays exactly where it is: it is the authoritative always-on
check for callers that build a ServerConfig without this resolver, and it
raises the SAME message from the SAME builder
(:func:`arbi_serve.calibration.locate.calibration_unreadable_reason`).

Hard-fail before ANY GPU work when the pinned bundle's quantizer cannot
serve the pinned width.

A LADDER bundle carries centroid cells for several widths but names ONE
global ``k_quantizer``/``v_quantizer``. vq2 can only serve the widths in
``tkv.core.VQ2_SERVABLE_BIT_WIDTHS`` -- its pair index is one byte, so 4
bits per channel is the ceiling -- while the same file happily carries
scalar Lloyd cells for 5..8. So ``k_quantizer: vq2`` on such a file is a
true statement about its 2/3/4 rungs and a FALSE one about the wider rungs
it also carries, and pinning one of those wider rungs produces a bundle
that CANNOT be loaded by any code path, ever.

WHY THIS IS ITS OWN GATE AND NOT LEFT TO THE PARSER. The parser refuses it
correctly (``tkv.runtime._calibration_parse._missing_vq2_msg``), but only
once the model's weights are on the GPU -- minutes in. Under a container
``restart: unless-stopped`` policy that is not an error message, it is a
restart loop: the refusal scrolls past at the end of every boot and the
operator sees a service flapping, not a config mistake. This gate is the
same policy at CLI-resolution time, where it costs a JSON read and the
process exits once.

Uniform widths only. A smart-mix bundle assigns widths per layer from its
own byte-budget table, so the pairing is per (layer, side) and cannot be
settled from the env; that case still lands on the parser's refusal, which
now names the width ceiling.

Translate ``--kv-cache-dtype`` into the internal backend spec list.

Appends a ``paged_kv:tkv-k{K}v{V}`` entry to ``args.backends`` when
the operator passed ``--kv-cache-dtype tkv``. The internal engine /
schema / hot-swap surface continues to speak ``kind:name`` strings;
only the user-facing CLI mention of ``paged_kv:tkv-...`` is gone.

The sentinel ``None`` (flag unset) is resolved by
:func:`_resolve_kv_cache_default` before this runs; by the time we
get here ``kv_cache_dtype`` is ``"auto"`` or ``"tkv"``.

Idempotent — calling twice on the same Namespace doesn't double-add.

``_parse_size`` plus the literal ``auto`` sentinel.

``--max-context auto`` defers the value to build time, where it resolves
to the model's ``max_position_embeddings`` ceiling (clamped down by the
realized compressed KV capacity). Any other value parses as a size.

An int CLI value plus the literal ``auto`` sentinel.

Used by knobs whose value is genuinely derived at build time (e.g.
``--mtp-n-draft auto``, which reads the K calibration cache). Any other
value parses as an int.

``--max-batch`` — a concrete concurrency; ``auto`` is refused loud.

Wraps :func:`arbi_serve.config_groups.parse_max_batch` so the refusal
message is the same one the config surface and the admin override raise.
``ArgumentTypeError`` so argparse prefixes it with the flag name.

Parse ``--pool-member 'served_name[,key=value]*'`` repeats.

Each spec is a comma-separated list whose FIRST token is the member's
``served_name`` (the routing key) and whose remaining tokens are
``key=value`` overrides. Recognised keys: ``path`` / ``max_context`` /
``calibration_path`` / ``dtype`` / ``attention_backends`` /
``default_backends``; an omitted key inherits the boot model's value at
startup. ``max_context`` accepts the ``32k`` / ``4096`` size syntax. The
two backends keys take ``+``-separated ``kind:name`` spec lists (comma is
the override separator), e.g.
``attention_backends=paged_kv:tkv-bypass+paged_kv:tkv-k4v4`` — the
cross-architecture member surface (a bf16 dense member alongside a
tkv-calibrated boot model needs its own dtype + backends).

Returns a ``tuple[PoolMemberConfig, …]``. Raises ``SystemExit`` on a
malformed spec or a duplicate served_name — fail-fast at CLI parse time.

Resolve stable-VA residency. On by default — it just works, no flag.

Single-member residency (the no-``--pool-member`` case) is the
byte-identical single-model path at ~zero cost, and it is what makes the
whole runtime-reconfig surface work without any boot flags: live
``/v1/admin/config_override`` (MBT / max_batch / max_context / capture
flags), ``/v1/admin/attention_backend`` (codec swap), and recapture-free
model cycling all require the stable-VA pool to be engaged. Making it an
opt-in flag would silently disable all of that on a stock ``docker
compose up``, so it is on by default.

Off only when cuMem pools are disabled (``--no-cumem-pools`` — the
documented debug escape hatch; residency / sleep-wake / park are
non-functional on the native allocator anyway). The legacy
``--stable-va-residency`` flag and ``--pool-member`` remain accepted (they
just can't turn off what is already the default).

Why a model path does or does not offer a bundled MTP head.

``has_head`` is the answer :func:`detect_bundled_mtp_head` returns.
``reason`` keeps the two very different negatives apart, which is the
point of this type: "the config read fine and declares no head" is a
property of the CHECKPOINT, while "the directory is absent / empty /
has no readable config.json" is a property of the DEPLOYMENT — almost
always a mis-pointed bind mount. Collapsing both into a bare ``False``
makes the boot blame the checkpoint for a wrong mount.

``detail`` is the operator-facing description of the state (it names
the path); it is ``None`` when there is no path to name.

Inspect ``model_path`` for a bundled MTP head, keeping the reason.

Modern checkpoints (Qwen 3.5 / 3.6 / DeepSeek V3 ...) ship the MTP
draft head inside the main safetensors and surface it via one of
several config keys. Forcing operators to also pass a draft flag for
these checkpoints is a foot-gun, so we inspect the config and treat
any of the following as a positive bundled-head signal:

  * ``mtp_num_hidden_layers > 0`` (Qwen 3.5 / 3.6 convention)
  * ``num_mtp_layers > 0`` (alternative spelling some forks use)
  * ``mtp_n_predict`` non-zero / non-empty
  * ``mtp_layers`` non-empty (explicit per-layer index list)
  * ``text_config.mtp_num_hidden_layers > 0`` (multi-modal nesting)

Never raises: every failure to read becomes a ``reason``, so the
caller can fall through to the explicit-flag path (and, at the boot
gate, say which of the two negatives it actually hit). Detection is
intentionally conservative: any field present and truthy -> bundled
head; missing / zero -> no bundled head.

A path that is neither absolute nor ``./``-relative and does not
exist is taken for a Hugging Face hub id, not a broken mount: the
weights are not on disk yet, so the config is simply not inspectable
here.

Return True iff the model's ``config.json`` declares a bundled MTP head.

Thin bool view of :func:`probe_bundled_mtp_head` — see it for the
recognised config keys and for the reason codes a caller needs when a
False has to be explained to an operator.

Resolve ``--compile-on`` precedence vs ``ARBI_COMPILE_ON`` env.

Precedence (CLI wins, env fills the gap, default ON):

  1. ``cli_value is not None`` → take the explicit CLI flag.
  2. Otherwise consult :func:`compile_on_from_env` (default ON).

Compile-on is the perf-best default. ``--no-compile-on`` /
``ARBI_COMPILE_ON=0`` opt out for debug.

Centralized so test setups, the CLI, and any future operator-
facing surface read the env var the same way.

Resolve ``--text-only`` / ``--language-model-only`` against vision env.

The vision tower is driven by ``runtime_flags().enable_vision`` (env
``ARBI_ENABLE_VISION``), and is off by default. ``--text-only`` is the
explicit, discoverable form of that default: it forces vision off by
pinning ``ARBI_ENABLE_VISION=0`` in the environment (the same single
knob the build reads), so the model build's vision-skip path fires and
logs the VRAM saved.

Contradiction guard (fail loud, never silent): combining ``--text-only``
with ``ARBI_ENABLE_VISION=1`` (vision explicitly on) is self-contradictory
— one flag says "skip the tower", the env says "load it". Rather than
silently pick a winner, raise ``SystemExit`` so the operator resolves it.

No-op when ``--text-only`` was not passed: the env is left untouched and
the existing default (vision off unless ``ARBI_ENABLE_VISION=1``) stands.
Idempotent.

Resolve ``--activation-quant {auto,bf16,fp8}`` onto the single
``ARBI_SERVE_AWQ_NO_A8`` tri-state the AWQ loader reads.

The flag is the first-class, discoverable form of that env knob:

  * ``bf16`` → ``ARBI_SERVE_AWQ_NO_A8=1`` (W4A16 everywhere),
  * ``fp8``  → ``ARBI_SERVE_AWQ_NO_A8=0`` (W4A8 on any fp8-MMA arch; the
    GDN ``in_proj_qkv`` auto-exclusion in the loader still applies, and an
    unsupported arch refuses loud at load rather than serving bf16),
  *

Env override channel (fail loud, never silent): ``ARBI_SERVE_AWQ_NO_A8``
stays usable on its own (``auto`` leaves it alone). But passing an explicit
``--activation-quant`` that contradicts an already-set env value is
self-contradictory — one says fp8, the other bf16. Rather than silently
pick a winner, raise ``SystemExit`` so the operator sets exactly one. An
env value that agrees with the flag is accepted (idempotent).

No-op when the flag is ``auto`` (the default): the env is left untouched.
Idempotent.

Shared core for the comma-separated-int-tuple CLI flags.

``None`` (flag absent) / empty string → ``None`` ("unset"; caller reads
env / config default). A disable sentinel — ``off`` / ``0`` / ``false`` /
``no`` / ``none`` plus any ``extra_sentinels`` (case-insensitive, possibly
whitespace-padded) → ``()`` (explicitly empty). Otherwise a sorted,
deduped tuple of positive ints. Each non-sentinel entry must be a positive
integer; a non-int or non-positive value raises ``SystemExit`` so a
malformed flag fails fast at CLI parse time.

Parse ``--prefill-cudagraph-buckets`` into a sorted dedup tuple.

Thin wrapper over :func:`_parse_cli_int_tuple`: ``None`` / empty → unset,
``off`` / ``0`` / ``false`` / ``no`` / ``none`` → ``()`` (disabled), a CSV
of positive ints → the sorted-dedup tuple.

Resolve ``--prefill-cudagraph-buckets`` precedence vs env / default.

Precedence (CLI wins, env fills the gap, default ON):

  1. ``cli_value`` is a non-empty, non-sentinel string → take it.
     ``"off"`` / ``"0"`` / ``"false"`` / ``"no"`` / ``"none"``
     (any case) explicitly disable (returns ``()``).
  2. Otherwise consult :func:`prefill_buckets_from_env` (default
     ``(256, 512, 1024, 2048, 4096)``).

Mirrors :func:`_resolve_compile_on`. Centralized so test setups,
the CLI, and any future operator-facing surface read the env var
the same way.

Parse ``--cudagraph-kv-pages-buckets`` into a sorted dedup tuple.

Thin wrapper over :func:`_parse_cli_int_tuple`: ``None`` / empty → unset,
``off`` / ``none`` / ``()`` / ``0`` / ``false`` / ``no`` → ``()`` (the
single full-width bucket — no intermediate rungs), a CSV of positive
KV-page counts → the sorted-dedup tuple.

Resolve ``--cudagraph-kv-pages-buckets`` precedence vs env / default.

Precedence (CLI wins, env fills the gap, config default last):

  1. ``cli_value`` is a non-empty, non-sentinel string → take it.
     ``off`` / ``none`` / ``()`` / ``0`` / ``false`` / ``no`` (any
     case) → ``()`` (the single full-width bucket).
  2. Otherwise the ``ARBI_CUDAGRAPH_KV_PAGES_BUCKETS`` env override
     (``None`` when unset).
  3. Otherwise the ``ServerConfig.cudagraph_kv_pages_buckets`` field
     default (``()`` — single full-width bucket).

Normalise bit widths from either a CLI comma-string ("2,4,8") OR an
already-parsed list/tuple (a recipe sets ``target_bit_widths`` as a list, so
both forms must be accepted; treating a list as a string raises ``'list'
object has no attribute 'split'``).

Resolve the serving task, auto-detecting from the model dir when asked.

Explicit ``--task generate|embed|rerank`` wins. ``"auto"`` inspects
the model directory: sentence-transformers pooling layout
(``config_sentence_transformers.json`` / ``modules.json``, shipped by
Qwen3-Embedding) → ``"embed"``; a ``"rerank"`` token in the dir name
→ ``"rerank"``; an ``"embed"`` token → ``"embed"``; otherwise
``"generate"``. The heuristic is a convenience — operators can always
pin it with ``--task``.

Resolve the per-step timing-debug tier from its three surfaces.

``--timing-debug``, the ``ARBI_TIMING_DEBUG`` env twin, and the
deprecated ``--detailed-timing`` spelling. Any one turns it on; the
deprecated spelling logs a one-line notice naming the replacement, since
the name is load-bearing in existing compose files and scripts.

Boot-the-server arg surface — shared between the default form and
forward-compat with future subcommands.

Delegates to themed helpers; the call order is the argparse add order
(and thus the --help layout), so it must not change.

argparse argument-group builders for the ``arbi-serve`` CLI.

Sibling of :mod:`arbi_serve.cli.parser`; holds cohesive ``_add_*_args``
group builders moved out to keep every module under the file-size limit.
Torch-free (only stdlib + the torch-free ``arbi_serve.config``) so
``arbi-serve --help`` never boots the engine.

The public parser entrypoints (``build_parser`` / ``_add_serve_args``)
remain importable from :mod:`arbi_serve.cli.parser`, which re-imports the
builders defined here.

Subset of the serve flags relevant to populating kernel caches.

The full serve parser is reused for ``cfg_from_args`` shape
compatibility — but the warmup path doesn't need or accept
``--host``/``--port`` (it never binds a socket).

Implement ``arbi-serve warmup ...``.

The lifecycle here mirrors the serve path: build a real
:class:`Engine`, run the same engine loop, submit one request
via the ``submit()`` sync API, drain it, then shut down. The Turbo prefill
CuTeDSL kernels JIT on first forward — populating
``$FLASH_ATTENTION_CUTE_DSL_CACHE_DIR`` — and Triton
autotunes populate ``$TRITON_HOME``. With both env vars pointed
at a persistent volume (the Dockerfile defaults to ``/cache/...``),
a subsequent serve boot skips the cold compile.

Idempotency is achieved at the cache-volume level: the kernel
JIT layers themselves probe the cache dir and short-circuit when
a hit is present, so a second warmup with the same flags does
no real work beyond model load.

Distinct per-layer ``(k_bits, v_bits)`` across the engine's TKV layers.

Reads ``eng.attn_ops`` (the live built attention ops); each
``TkvAttnOp`` carries the genuine per-layer bits the smart-mix bundle
assigned. Returns ``[]`` when no TKV layers are active (bf16 warmup) or
the ops aren't available — the report then just lists what's on disk.

Thin layer over PyTorch's compile primitives.

Every kernel arbi-serve calls is a ``torch.library.custom_op`` with
``register_fake`` (see ``arbi_serve/_custom_ops.py``), which Inductor
treats as a black box natively. The machinery here is:

  * a runtime-gated trampoline (``runtime.install_torch_compiled``) so
    the eager path stays bit-identical when compile-on is OFF,
  * a contextvar gate (``context.is_in_piecewise_cuda_graph``) the
    capture sweep flips on,
  * a tiny capture-side bridge (``capture_bridge.compile_capture_ctx``)
    that flips that gate iff the model is compile-eligible,
  * a constructor-time decorator (``decorator.support_torch_compile``)
    that installs the trampoline only inside an active compile context.

On-disk compile-cache persistence.

Two persisted artifacts share one versioned key (torch / GPU / compile
config):

  * the torch Mega-Cache (FxGraph + AOTAutograd Inductor backend),
  * the Inductor filesystem cache root (uid-namespaced under ``/cache``).

A hit on the first skips the cold-compile cost on warm boots. Every
loader/saver is best-effort and never raises into boot.

The Dynamo FRONTEND trace is deliberately NOT persisted. Persisting it
cost more than it saved: with the Mega-Cache warm the whole frontend trace
is ~0.6s, while loading and installing a per-callable guard package cost
more than that before anything else. See the removal in the git history if
the idea comes back.

Versioned cache key: custom-op ABI epoch + torch version + GPU
arch/capability + a hash of the compile-affecting config. Any of these
changing invalidates the blob
(a stale Mega-Cache from a different torch / GPU / config would either
miss-by-key here or — worse — load incompatible artifacts that fault at
runtime, so we key on every axis that changes codegen).

Pure and best-effort: a probe failure (no CUDA, odd torch) degrades to a
stable-but-coarse key rather than raising — the caller wraps load/save
in try/except regardless.

Resolve the on-disk Mega-Cache blob path for the current key.

``/cache/arbi-serve/compile-cache/<key>.bin``. Returns ``None`` when the
directory can't be created / written (read-only mount) — every caller
treats ``None`` as "no persistent compile cache, compile normally".
Best-effort, never raises.

Load the persisted torch Mega-Cache blob (if any) into this process.

Bundles FxGraph + AOTAutograd artifacts, so a hit skips the Inductor
BACKEND codegen on a warm boot (it does not skip the Dynamo frontend
trace — nothing does; see the module docstring). Returns True on a hit,
False on miss/error. A miss, corrupt blob, version mismatch, or any
error falls through to normal compilation and never raises into boot.

Number of artifacts a ``CacheInfo`` reports, or ``-1`` when unknown.

``CacheInfo`` groups artifacts by kind, and the groups are exposed as
properties (``inductor``, ``autotune``, …) rather than instance
attributes, so summing ``vars(info)`` counts nothing. Walk the public
sequence-valued attributes of the object AND its type.

Boot status of the compile blob arbi-serve persists itself.

``megacache-fx`` is the Inductor BACKEND bundle (FxGraph + AOTAutograd).
It is keyed on the torch / GPU / compile-config build key, so its path
can only be resolved once torch is importable — after the torch-free pin
has already stamped the env-var-based report. Registered into that report
via :func:`~arbi_serve._compile_cache_env.extend_cache_report` so the
boot log states its REUSE/COMPILE status instead of reporting a
fully-warm boot while it is missing.

Report-only (``asserted=False``): the blob key carries the GPU identity,
which the warm-cache baseline does not, so a first boot on another GPU of
the same torch build is indistinguishable from a lost cache.

Persist this process's torch Mega-Cache blob to the versioned path.

Called once after the warmup compile completes. Writes only when the
blob differs from what's on disk (first boot, or a content change), so
a warm-boot hit doesn't rewrite an identical file. Returns True if a
file was written. Any error is swallowed (never raises into
boot); an unwritable mount or missing API is a silent no-op.

Pick the on-disk Inductor cache root. Pure — no global mutation.

Precedence: explicit ``cache_dir`` arg → existing
``$TORCHINDUCTOR_CACHE_DIR`` → a uid-namespaced subtree under ``base``
(``/cache/inductor/u<uid>``) when writable → ``None`` (Inductor's own
``~/.cache/torch/inductor`` default).

The uid namespacing is the load-bearing part: ``/cache`` is a shared
mount, so a prior run by a *different* uid (e.g. a root-owned container
run on a bare box) leaves root-owned subdirs (``fxgraph/<hash>/…``)
that are unwritable here. A bare top-level probe would pass anyway, so
Inductor would then fail artifact writes mid-boot and — worse — load a
foreign/partially-written compiled artifact, which faults at runtime
with ``cudaErrorIllegalAddress``. Isolating by uid means one user's
cache can never poison another's; the current user owns its whole
subtree, so every nested write succeeds. If even the uid dir can't be
created (read-only mount, unwritable ``base``), we return ``None`` and
let Inductor fall back to the user-owned home cache — never crash.

Bridge between the cudagraph capture sweep and the compile trampoline.

The capture call sites in ``arbi_serve/runtime/capture/`` open
:func:`compile_capture_ctx` around each ``torch.cuda.graph(...)`` block
so the trampoline-installed ``model.forward(...)`` routes to the
Dynamo-compiled callable instead of eager. When the model isn't
compile-decorated the bridge is a no-op (eager forward captured into
the cudagraph as before).

True if THIS module's own ``forward`` body is Dynamo-traced.

Distinguishes a REAL whole-forward :func:`install_torch_compiled`
trampoline from the two installs that leave ``forward`` running
eager: the model-level marker
(:func:`install_torch_compiled_marker`, stamped
``_arbi_compile_marker_only``) and the split-attn per-method install
(:func:`install_torch_compiled_methods`, whose aggregate state
carries ``method_states`` and whose ``forward`` is a passthrough).

Asked of the model ROOT by anything that must run eager inside the
root's forward. The activation arena is the case that matters: it
allocates at the eager per-layer dispatch seam
(``LayerStackModelMixin._dispatch_layer_args``), which the root's
forward calls, and
:meth:`~arbi_serve.runtime.activation_arena.ActivationArena.alloc`
re-points a tensor with ``Tensor.set_`` — untraceable. A root whose
forward is compiled therefore swallows the seam into the trace and
the arena cannot be served. The compiled BLOCK forwards below it are
fine: they receive the buffer as an ordinary tensor argument.

The compile-decorated module a boot-warmup driver should gate on.

Most model classes carry ``@support_torch_compile`` directly, so
``model`` itself is the answer. A composite multi-modal wrapper
(``NemotronVoiceChatModel``: STT ``self.backbone`` +
``self.tts_backbone`` + perception/codec pieces) does NOT carry the
decorator on the outer class — only ``self.backbone`` does — so
``is_compile_eligible(model)`` reads False for it even though its
per-layer blocks are genuinely compile-decorated and compile
unconditionally on first call (:func:`install_torch_compiled`'s
trampoline). A caller that gates an entire warmup pass on
``is_compile_eligible(model)`` then skips it outright for such a
wrapper — every per-layer block's first-ever compile lands on a live
request instead of at boot. Falls through to ``model.backbone`` when
the top-level object itself isn't eligible; returns ``None`` when
neither is. Forward calls still go through the OUTER ``model`` (its
``forward`` delegates to the backbone), so this only changes the
eligibility gate, not what gets called.

Operator-tunable knobs for the compile pipeline.

Attributes
----------
compiler:
    ``"inductor"`` (default Inductor backend) or ``"eager"`` (Dynamo
    no-op — used by tests that exercise trampoline routing without
    paying compile cost).
enable_if:
    Operator master switch. ``False`` makes ``install_torch_compiled``
    a no-op so the model stays eager.
cache_dir:
    Override Inductor's compile cache root. ``None`` defers to
    Inductor's default (``$TORCHINDUCTOR_CACHE_DIR`` /
    ``~/.cache/torch/inductor``).
split_attn:
    When ``True``, decorated decoder layer classes that expose the
    triple ``_pre_attn`` / ``_attn_eager`` / ``_post_attn`` methods
    compile the two non-attn pieces independently (via
    :func:`install_torch_compiled_methods`) and the orchestrator
    captures them as two cuda graphs per layer per bucket, with the
    eager attn call running between replays. This keeps attention
    kernel scratch (FA varlen / paged_attention / FLA chunk fwd
    intermediates) out of the captured-graph mempool.

    Default ``False``: layer classes route through the
    whole-forward capture path. When ``True`` and a layer class
    only exposes ``forward`` (no triple), that class still falls
    through the whole-forward capture path; only layers wired for
    the split get the new capture path. This keeps split_attn a
    strict opt-in that scales arch-by-arch.

The trampoline gate. True when the compiled forward should run.

Compile-on integration: ``@torch._dynamo.assume_constant_result``
tells Dynamo to specialize the ContextVar read at trace time and
bake the value into the compiled graph as a constant. Dynamo can't
trace ``ContextVar.get()`` directly (graph break gb0156); the
decorator makes the call opaque-but-constant from the trace's POV.
Recompiles trigger if the constant changes between traces.

True when the boot-time piecewise per-layer capture sweep is running.

During ``_capture_one_bucket``, the engine drives ``model.forward``
through the dispatcher's pool branch so every eligible layer
captures into its own ``torch.cuda.CUDAGraph``. The layer-level
real trampolines see this flag set and route to their eager
``unbound_fwd`` for the duration — eager kernel sequences land in
the captured graph, no Inductor compile inside the cudagraph
capture region. Steady-state hot path (no contextvar set) keeps
using the compiled callable.

``@torch._dynamo.assume_constant_result`` keeps the read opaque-but-
constant under any Dynamo retrace; the engine only flips the
contextvar in eager paths, so the per-trace value is stable.

Engine-side: mark the with-block as inside the per-layer capture sweep.

Trampolines route to the eager ``unbound_fwd`` while this flag is
set so :func:`_capture_one_bucket`'s warmup + capture forwards
actually descend into the dispatcher's pool branch (where the per-
layer ``torch.cuda.CUDAGraph`` records happen). See
:func:`is_in_piecewise_capture_sweep` for the rationale.

``@support_torch_compile`` model-class decorator.

Marks an ``nn.Module`` subclass as compile-eligible. When the model is
constructed inside an active :func:`set_compile_context`, the wrapped
``__init__`` calls :func:`install_torch_compiled` to install the
runtime-gated trampoline; outside a context the decorator is a no-op
and the model stays eager.

Activate a compile context for the with-block (LIFO save/restore).

Opening a context resets the frame-cache purge ledger
(:func:`~arbi_serve.compile.trampoline.reset_frame_cache_purge_ledger`):
the first compile of each code object under this context purges Dynamo's
cache for it once, and every later instance of that code object shares
the entries the first trace produced.

True iff any direct or transitive child of ``module`` is a class
decorated with ``@support_torch_compile(level="block")``.

The block decorator stamps ``cls._arbi_compile_level = "block"``
on the class (``cls``, not the instance) — we walk
``module.modules()`` and inspect each one's class.

Decorate an ``nn.Module`` subclass as compile-eligible.

``level``:
  * ``"model"`` (default) — top-level model class. Installs a
    *marker* trampoline (eager passthrough) when block-decorated
    descendants exist, so the per-layer decorators do the real
    compile work. If the model has no block-decorated children
    (e.g. a future small model with no per-layer decoration),
    installs the real Inductor trampoline at the model level
    as a fallback.
  * ``"block"`` — per-layer decoder block class. Always installs
    the real Inductor trampoline.

``methods`` (block-only): when supplied, the decorator wraps each
named method (e.g. ``("_pre_attn_compiled", "_post_attn_compiled")``)
rather than ``forward``. Used by split-attn arches where attention
runs eager between two compiled pieces so its kernel scratch never
lands in the captured-graph mempool. The class's own ``forward``
remains the eager orchestrator that calls the compiled pieces in
order. Mutually exclusive with ``level="model"`` (split-attn is a
block-level concept).

Split-attn knob: when ``methods`` is supplied, the
per-method install path activates only when the active
:class:`CompilationConfig` has ``split_attn=True``. With the knob
OFF (the production default), the decorator transparently falls
through to :func:`install_torch_compiled` and compiles the
orchestrator ``forward`` as one whole-layer graph — bit-equivalent
to the whole-layer surface. Operators opt in via
``ServerConfig.split_attn=True`` (CLI ``--split-attn``); the model
layer's ``forward`` chains pre/eager/post in three Python calls
either way, so the only behaviour difference is which set of
callables Dynamo sees and where the orchestrator captures the
cuda graphs.

Context-local Dynamo recompile-limit / de-specialization knobs.

torch 2.12 stores ``torch._dynamo.config`` overrides in a
``contextvars.ContextVar`` (``torch/utils/_config_module.py``:
``user_override = ContextVar(...)``), so an override is scoped to the
``contextvars.Context`` that set it, not to the process or the
thread. Any context boundary (an ``asyncio`` task, which runs in a
``copy_context()`` taken at task creation; ``run_coroutine_threadsafe`` /
``call_soon_threadsafe``, which replay the caller's context on the engine
thread; ``run_in_executor``, which propagates no context at all) yields a
reader that sees the process default again. The helpers here re-apply the
limits in the calling context, on the forward hot path.

Raise Dynamo's per-code-object recompile / cache-size limits to ``limit``
on the calling thread. Idempotent; only ever raises (never lowers). Both the
new (``recompile_limit``) and legacy (``cache_size_limit``) names plus their
``accumulated_`` siblings are set defensively across torch versions.

Set ``torch._dynamo.config.allow_unspec_int_on_nn_module = True`` on the
calling thread. Idempotent; safe across torch versions (no-ops where the
knob is absent).

Like ``recompile_limit``, this knob is thread-local on torch 2.12: a value
set on the build/main thread is not observed by the engine run-loop thread
that drives serving compiles, so it must be re-applied per-thread next to
:func:`apply_dynamo_recompile_limits`. Without it, a per-layer compile on a
serving thread re-specializes the shared decoder forward on the nn.Module
int attribute ``self.layer_idx`` (guard ``self.layer_idx == N``) — one
recompile per layer, eventually tripping ``FailOnRecompileLimitHit`` and
surfacing as a CUDA illegal-memory-access. The flag makes Dynamo carry the
int as a SymInt graph input, so one compiled artifact serves every layer.

Re-assert the recompile limits and the ``allow_unspec_int_on_nn_module``
de-specialization flag in the calling context, so whichever context drives a
per-layer compile and its subsequent recompiles sees 256 (not the process
default 8) and leaves ``layer_idx`` unspecialized (no per-layer recompile
storm).

Deliberately not latched behind a once-per-thread flag. The knobs are
ContextVar-scoped (see the module docstring), so "this thread already did
it" does not imply "this context can see it": boot warmup sets 256 in the
boot context, and a latch would then no-op forever, leaving the first
batch>1 serving prefill — which runs in a different context — to recompile
under a limit of 8, trip ``FailOnRecompileLimitHit`` under fullgraph, and
surface as a CUDA illegal-memory-access at the next step's batch-build.
Concurrency 1 never exceeds 8 buckets, so it never triggers this path.

Both helpers are idempotent and only ever raise a limit, so re-asserting
costs a few attribute compares next to a model forward.

Apply ``ARBI_INDUCTOR_PROFILE`` config knobs onto
``torch._inductor.config``. Returns a dict of {attr: value} of what
was actually set (skipped attrs that don't exist on this torch
version are absent). Profile-time only; not on the hot path.

Profiles (each turns on a single audit knob, except ``all_on``):

  * ``default`` — no extra knobs (today's behaviour).
  * ``epilogue_fusion`` — ``epilogue_fusion=True``: fuse pointwise
    ops into the matmul epilogue (saves a kernel launch + a
    round-trip through HBM on every matmul-followed-by-pointwise
    chain — the elementwise budget on the decode hot path).
  * ``shape_padding`` — ``shape_padding=True``: pad to
    tensor-core-friendly multiples (8/16) so cuBLAS / Triton mm can
    select TC kernels without fallback.
  * ``combo_kernels`` — ``combo_kernels=True`` +
    ``benchmark_combo_kernel=True``: combine small kernels into one
    launch (saves launch overhead on the elementwise tail).
  * ``coord_desc`` — ``coordinate_descent_tuning=True``: block-size
    autotune via coordinate descent (slow at compile, fast at
    runtime).
  * ``reduce_overhead`` — equivalent to
    ``mode='reduce-overhead'``: skips kernel autotune, focuses on
    reducing CUDA launch overhead.
  * ``max_autotune`` — equivalent to
    ``mode='max-autotune-no-cudagraphs'`` (we manage cudagraphs
    ourselves): full kernel autotune. Slow to compile.
  * ``all_on`` — epilogue + shape_padding + combo_kernels +
    coord_desc stacked.

The mode-style profiles set the same knobs that
``torch.compile(mode=...)`` would; we set them on
``torch._inductor.config`` directly so per-call ``torch.compile``
invocations elsewhere in arbi-serve don't need to thread a ``mode``
arg through.

Turn on Inductor ``combo_kernels`` independently of the profile.

Horizontal kernel fusion combines many small independent
pointwise/reduction Triton kernels into one launch, cutting the
per-launch tax on the many tiny norm/RoPE/elementwise kernels a
decode step emits. It is layered on top of the active
``ARBI_INDUCTOR_PROFILE`` because those profiles are
mutually-exclusive (you cannot select both ``reduce_overhead`` and
``combo_kernels``), and on the cudagraph paths the active profile is
``reduce_overhead``.

Safe with arbi's compile/capture model: combo_kernels acts purely
inside Inductor's lowering — it emits one fused Triton kernel where
it would have emitted several, introduces no graph break (so
``fullgraph=True`` holds), and never touches cudagraph capture
(Inductor cudagraphs stay OFF; arbi owns capture). Returns the
{attr: value} actually set (absent on torch versions lacking the
knob).

Runtime-gated ``torch.compile`` wiring.

This module is the public entry point for arbi-serve's torch.compile
infrastructure. The implementation is split across sibling modules by
responsibility, and re-exported here so existing
``from arbi_serve.compile.runtime import X`` imports keep working:

  * :mod:`arbi_serve.compile.trampoline` — the runtime-gated trampoline
    installers (``install_torch_compiled`` and friends).
  * :mod:`arbi_serve.compile.cache` — on-disk Mega-Cache + Dynamo
    Mega-Cache persistence and the Inductor cache-dir resolver.
  * :mod:`arbi_serve.compile.inductor_profiles` — ``ARBI_INDUCTOR_PROFILE``
    config-knob bundles.
  * :mod:`arbi_serve.compile.dynamo_config` — thread-local Dynamo
    recompile-limit / de-specialization knobs.

:func:`configure_inductor_caches` lives here as the boot-time orchestrator
that wires those pieces together.

Wire up Inductor's filesystem caches so warm boots reuse compiled artifacts.

Sets:
  * ``torch._inductor.config.fx_graph_cache = True`` — caches the
    FX graph + Inductor lowerings keyed on a content hash of the
    graph + autograd config + active flags. Hits on warm boot
    skip the (multi-minute) pattern-matcher + reinplace passes
    that dominate compile time on hybrid Qwen3.5/3.6 graphs.
  * ``$TORCHINDUCTOR_CACHE_DIR`` — the on-disk root the cache
    persists into. Inductor reads this env var lazily; we set it
    early at boot so every subsequent ``torch.compile`` call
    sees the override. Idempotent — won't overwrite an existing
    explicit override.
  * ``$TRITON_CACHE_DIR`` — pinned to a ``triton`` subtree of the
    same persisted root (torch 2.12's ``triton_cache_dir()`` does
    not track ``$TORCHINDUCTOR_CACHE_DIR``, defaulting to an
    ephemeral ``/tmp`` path). Co-locating the cubins with the
    FX-graph bundle that references them by absolute path stops the
    ``use_static_cuda_launcher`` reload from missing + re-JITing the
    exl3 GEMM kernels on warm boots. Paired with
    ``bundle_triton_into_fx_graph_cache = True``.

``ARBI_INDUCTOR_PROFILE=<name>`` (env-only) layers an additional
set of Inductor config knobs on top of the baseline. Used for
Inductor-config sweep benchmarking; leave unset in production.
Recognized profiles: ``default``,
``epilogue_fusion``, ``shape_padding``, ``combo_kernels``,
``coord_desc``, ``reduce_overhead``, ``max_autotune``, ``all_on``
(see :func:`_apply_inductor_profile`). The selector is a strict
no-op when the env var is unset (the production default), so the
CPU gate / OFF path / default-ON path all stay byte-for-byte
identical with the selector absent.

The caller passes ``cache_dir`` (typically
:attr:`CompilationConfig.cache_dir`). When ``None``, we honor an
existing ``$TORCHINDUCTOR_CACHE_DIR`` if set; otherwise fall back
to ``/cache/inductor/u<uid>`` (the Dockerfile mounts ``/cache`` as a
VOLUME for exactly this purpose) when that path is writable, else
Inductor's default (``~/.cache/torch/inductor``). The subtree is
namespaced by uid so a shared ``/cache`` mount can't be cross-user
poisoned — a stale root-owned artifact from another run would
otherwise fault at runtime with ``cudaErrorIllegalAddress``.

Returns the resolved cache dir (or ``None`` if neither override
nor /cache writable — Inductor's default kicks in). Safe to call
multiple times; later calls just narrow the active dir.

Runtime-gated ``torch.compile`` trampoline installers.

Every arbi-serve kernel is a ``torch.library.custom_op`` with
``register_fake`` (see ``arbi_serve/_custom_ops.py``). Inductor treats
these as black boxes, so we don't need a custom Dynamo backend, a graph
splitter, or eager/Inductor adapters — a direct ``torch.compile(forward)``
on the unbound forward suffices. The installed trampoline gates per-call
on ``is_in_piecewise_cuda_graph()`` so the eager path stays bit-identical
when compile-on is OFF.

Dynamo's own cumulative compile wall-clock for this process, or
``None`` when torch does not publish it.

The trampoline's split straddles ONE ``torch.compile`` and ONE first call
per callable. A callable that recompiles at a later shape — a fixed-shape
forward driven at several capture buckets — pays that inside a plain
trampoline call the split never times, so the ledger reads the whole of
it off torch's accounting instead of guessing.

One line attributing compile cost to decoration vs first call.

Without ``since`` the line covers everything compiled so far in the
process. With it, the line covers only what compiled AFTER the snapshot
and names the ``phase`` it is charged to — the same accumulator, read as
a delta, so a phase line and the cumulative line can never disagree.
The dynamo wall-clock delta, when torch publishes it, carries the
recompiles the per-callable split cannot see.

Yield a context manager that routes allocations through
``graph_pool`` when one is supplied (NamedMemPool exposes ``.use()``).
A null ctx when graph_pool is None — same shape, no extra branching
at call sites.

Forget which code objects have been purged in this process.

Called when a compile context opens (one per ``Engine.build``), so the
first compile of each code object under the NEW context purges Dynamo's
frame cache once more instead of reusing entries traced under the
previous context's configuration.

Drop Dynamo's cached entries for ``code`` — at most once per context.

The FIRST compile of a code object must not reuse an entry left by an
earlier context: ``_ensure_compiled`` drives its trace inside
``enable_piecewise_cuda_graph()`` and the gate reader
:func:`~arbi_serve.compile.context.is_in_piecewise_cuda_graph` is
``@torch._dynamo.assume_constant_result``, so the gate value is baked
into the graph and no guard separates an entry traced under the other
gate state.

Every LATER instance of the same code object reuses what the first one
produced: Dynamo's guards decide whether an instance needs an entry of
its own, so per-layer compile costs one trace per genuine guard variant
rather than one per layer instance (the raised ``recompile_limit`` sizes
the shared cache for exactly that). Purging per instance also drops the
entries the already-compiled instances are running on, so their next
call re-traces.

Temporarily raise ``sys.getrecursionlimit()`` to ``target`` for the
duration of the with-block; restore on exit (success or exception).

No-op when the current limit is already at or above ``target`` — we
never lower it. Operators can override via
``ARBI_COMPILE_RECURSION_LIMIT=<int>``; values <=0 disable the bump
(useful for diagnosing whether a fresh failure is recursion-related
or something else).

True iff ``ann`` annotates a single ``torch.Tensor`` (optionally
``| None`` / ``Optional[...]``), but not a container of tensors.

Handles both real-type annotations and *string* annotations. Every
arbi-serve model module carries ``from __future__ import
annotations`` (PEP 563), so at runtime ``inspect.signature`` returns
the annotation as a **string** — ``"torch.Tensor"``,
``"torch.Tensor | None"``, ``"Optional[torch.Tensor]"`` — never the
type object. All three string forms must be recognized, including
the PEP 604 union form used by args like ``residual_buf`` /
``gate_flat`` on the compiled decoder-layer methods: missing one
leaves an ``(N_tokens, …)`` arg un-marked, so Dynamo specializes the
graph on its static ``shape[0]`` and recompiles on the next
token-count bucket even though a sibling tensor arg was already
marked dynamic. Marking dim 0 dynamic is a guard/codegen hint only:
identical kernels, identical numerics, fewer specializations.

A container annotation (``list[torch.Tensor]`` / ``dict[...,
torch.Tensor]`` / ``tuple[torch.Tensor, ...]``) must not match —
those args are not a single ``(N, …)`` tensor and marking dim 0 is
meaningless. We detect that by rejecting any string carrying a
bracketed container head.

Walk ``forward_fn``'s signature, pick tensor-typed args as dyn-dim 0.

A *scalar* (optionally-None) Tensor arg is marked dynamic on dim 0 —
the flat-token batch dim ``(N_tokens, …)`` every per-token activation
carries. See :func:`_ann_is_tensor_scalar` for why the PEP 563 string
forms (incl. ``"torch.Tensor | None"``) must all be recognized:
missing one leaves an ``(N, …)`` arg un-marked and forces a per-bucket
recompile of an otherwise shape-agnostic graph.

Install a *marker* trampoline that satisfies the compile-eligibility
contract without actually running ``torch.compile``.

Used in per-layer compile mode (the default under
``ARBI_COMPILE_ON=1`` on hybrid models, e.g. Qwen3.5/3.6) where the
real Inductor trampolines live on each decoder-layer subclass —
compiling many small per-layer FX graphs instead of one large
whole-model graph avoids Inductor's pattern-matcher / reinplace
passes scaling quadratically with graph size.

The model still needs to look compile-eligible to:
  * :func:`is_compile_eligible` (capture bridge gate),
  * :func:`precapture_compile_warmup` (pre-capture warmup driver),
  * the per-arch test suite that pins
    ``model.forward.__func__.__name__ == "trampoline"`` after
    construction inside :func:`set_compile_context`.

The installed trampoline:
  * runs the original eager ``forward`` body (which dispatches into
    per-layer trampolines that DO compile),
  * exposes ``_arbi_compile_state`` (with ``compiled=True``, since
    there is nothing to compile at the model level),
  * exposes ``force_compile()`` (idempotent; runs the forward once
    to drive every per-layer trampoline's lazy compile, then sets
    the cached "compiled callable" to the eager forward).

``ARBI_COMPILE_OFF=1`` still routes through the eager path verbatim.

Install a runtime-gated compiled forward on ``module``.

The replacement ``forward`` is a trampoline that calls the original
forward when ``is_in_piecewise_cuda_graph()`` is False, and the
Dynamo-compiled callable otherwise (lazily compiling on first
call inside the gate window).

When ``graph_pool`` is supplied (a :class:`NamedMemPool`), the
Inductor compile + first-call workspace allocations are routed
through that pool's ``.use()`` ctx so they show up in our VRAM
accounting instead of the default CUDA allocator.

``ARBI_COMPILE_PROBE=1`` (env-only, debug) forces ``fullgraph=False``
so Dynamo enumerates every graph break in one trace pass instead of
crashing on the first. Used by
``scripts/diagnostics/probe_compile_breaks.py``. Production paths
leave the env var unset; ``fullgraph=True`` stays on.

dict whose ``"compiled"`` key reports the live aggregate state.

Defined at module-global scope rather than as a local class inside
:func:`install_torch_compiled_methods`: Dynamo guards on the nn.Module's
attributes, so this type appears in guard machinery, and a locally
defined class is awkward for anything that inspects or serializes it.

The live ``"compiled"`` value is computed from the per-method states
stored under the ``"method_states"`` key (set by the installer), so
callers reading ``state["compiled"]`` always see the current
aggregate without a stale snapshot.

Install runtime-gated compiled trampolines on ``methods`` of ``module``.

The split-attn variant of :func:`install_torch_compiled`: instead of
replacing ``forward``, this wraps each named method (e.g.
``_pre_attn_compiled`` / ``_post_attn_compiled``) with its own
Dynamo-compiled callable. ``forward`` itself stays eager and serves
as the orchestrator that calls the compiled pieces with the eager
attention call sandwiched between them.

Each wrapped method gets:
  * a private ``state`` dict (``compiled``, ``compiled_callable``);
  * a per-method ``force_compile_<name>`` driver bound on ``module``;
  * an entry in ``module._arbi_compile_state["method_states"]`` keyed
    by method name.

The aggregate ``module._arbi_compile_state`` is constructed so the
standard test surface (``forward.__name__ == "trampoline"`` /
``_arbi_compile_state["compiled"]``) keeps working — ``forward`` is
swapped for a passthrough trampoline that calls the original
``forward`` (which itself orchestrates the now-compiled pieces).

Symmetric with :func:`install_torch_compiled` re. env-var gates
(``ARBI_COMPILE_OFF`` / ``ARBI_COMPILE_PROBE`` /
``ARBI_ALLOW_COMPILE_FALLBACK``); operators see one consistent
surface across whole-forward and split-attn compile modes.

Runtime-gated forward: eager under nesting/off-switches, compiled otherwise.

Runs the eager forward when traced inside an outer compile,
when ``compile_off`` is set, or during the piecewise capture
sweep; otherwise lazily compiles on first call and replays the
compiled callable every step.

Typed config for arbi-serve.

Multi-backend wiring speaks ``kind:name`` end-to-end: every attention
backend is registered as ``"<state-kind>:<name>"`` and selected via
:class:`StateKind`.

The config dataclass groups live in :mod:`arbi_serve.config_groups` and the
top-level aggregate in :mod:`arbi_serve.config_server`; both are re-exported
here so ``from arbi_serve.config import ...`` keeps working unchanged.

Single source of truth: are the named GPU pools cuMem-backed?

cuMem-backed pools are the DEFAULT (and only supported) GPU
allocator backend — sleep/wake, the phase-2 freeze, stable-VA
residency, model cycling and the VRAM accounting invariant all
require them; a native-allocator boot yields pools where
park/release is a silent NO-OP (frees nothing → OOM on the next
load). The fallback for a config object MISSING the field is
therefore ``True`` (cuMem), never ``False``: a stub / namespace /
``None`` cfg must not silently degrade to native pools. Only an
EXPLICIT ``cumem_pools=False`` (``--no-cumem-pools``, the
unsupported-for-sleep escape hatch) disables the backend.

Every call site must use this helper instead of a raw
``getattr(cfg, "cumem_pools", …)`` — enforced by
``tests/test_cumem_default_guard.py``.

Read the ``ARBI_COMPILE_ON`` env var as a boolean.

``"1"``, ``"true"``, ``"yes"``, ``"on"`` (any case) → ``True``.
``"0"``, ``"false"``, ``"no"``, ``"off"``, empty / unset → ``default``.

``default=None`` (the normal call) resolves to the **canonical**
default — ``RuntimeFlags.compile_on`` via
``_FIELD_DEFAULTS["compile_on"]`` — so this helper, the CLI
(:func:`arbi_serve.cli._resolve_compile_on`), the
:attr:`ServerConfig.compile_on` field, and the registry all agree
BY CONSTRUCTION (currently ``True``: compile-on is the perf-best
default). Callers may pass an explicit ``default`` to override the
env-empty/falsy case (e.g. tests forcing OFF). Do NOT hardcode the
default here — derive it, so it cannot drift from the registry.

Read the ``ARBI_PREFILL_BUCKETS`` env var as a bucket tuple.

Precedence rules mirror :func:`compile_on_from_env`:

- ``""`` / unset → ``default`` (perf-on default).
- ``"off"`` / ``"0"`` / ``"false"`` / ``"no"`` / ``"none"`` (any
  case) → empty tuple = disabled.
- Comma-separated integers (e.g. ``"256,512,1024"``) → sorted
  dedup tuple of positive ints.

Raises :class:`ValueError` on invalid integers or non-positive
bucket sizes — let it propagate so a misconfigured env var fails
boot loudly instead of silently disabling the capture path.

Used by :func:`arbi_serve.cli.cfg_from_args` when no explicit
``--prefill-cudagraph-buckets`` CLI flag is passed (= unset CLI
surface). The CLI flag wins when present.

Config dataclass groups for arbi-serve (relocated from ``config.py``).

These are the individual ``@dataclass`` config groups that compose into
:class:`arbi_serve.config.ServerConfig`. Import them from
:mod:`arbi_serve.config` (the canonical, widely-imported surface); this
module exists only to keep every first-party file under the size cap.

Paged KV cache + tkv codec config.

block_size: tokens per page (see the field default below) — the
    codec's page granularity.
max_context: per-request max context (prompt + generation).
num_pages: KV pool depth OVERRIDE; see the field below. ``None``
    auto-sizes from the profiled VRAM budget.

Calibration (mapping is explicit and fail-fast — no silent corruption):
  ``calibration_path``: optional path to a Lloyd-Max calibration bundle.
    - If multi-bit bundle (``centroids[layer][k_bits_N/v_bits_M]``),
      applies to every TKV backend whose (k,v) bit pair is present.
    - If single-bit (top-level ``k_bit_width``/``v_bit_width``),
      applies ONLY to TKV backends whose bits match exactly. Applying
      a single-bit JSON to a backend with different widths is a hard
      error — silent application would corrupt the codec (centroid
      counts differ: K4 has 16/channel, K8 has 256/channel).
  ``calibration_paths``: optional per-backend overrides keyed by
    backend name (``{"tkv-k4v4": "<path1>", "tkv-k8v8": "<path2>"}``).
    Overrides ``calibration_path`` for the named backend. Use this
    when serving multiple TKV bit widths concurrently.

Continuous-batching shape.

Attributes:
    max_batch: maximum number of rows per slate.
    max_batched_tokens: token budget per slate (prefill + decode).
        Caps the active-batch flat-tensor shape that ``ScheduledBatch``
        materializes. The MTP verify pass is NOT in this budget: it runs
        as its own forward off the boot-sized ``VerifyBuffers`` slabs,
        reserved separately by the serving floor's spec term, so a decode
        row costs one token here whatever ``n_draft`` is. A COMPUTE budget
        per forward step, decoupled from both
        ``max_batch`` and ``max_context``. It right-sizes the
        capture + activation buffers to the compute budget instead of
        the full context (a much larger step directly grows the
        capture pool at the KV pool's expense). Default ``2048`` =
        ``chunk_prefill``. A decode row already rides along with a
        prefill chunk at the default (the mixed step reserves budget for
        the decode rows before sizing the chunk); what a WIDER budget
        adds is co-batching several prefill ROWS into one step. It costs
        KV: the step's activation peak scales with the budget and the
        serving floor holds that peak out of the pool, so the KV loss
        grows with the ratio to ``chunk_prefill``. The piecewise-capture
        ladder tops out at this budget. The CLI seam
        (:func:`arbi_serve.cli.config_builder.cfg_from_args`) reads
        this field default when ``--max-batched-tokens`` is unset.
    chunk_prefill: per-step prefill chunk size. Long prompts are
        split across this many tokens per step so a single
        in-flight prefill doesn't starve decode. Default ``2048``;
        ``--chunk-prefill`` reads this field default.
    chunk_prefill_explicit: whether ``chunk_prefill`` is the operator's
        number rather than one the engine picked. Set at the CLI seam;
        read by the piecewise capture sweep to decide whether a rung it
        cannot capture is a substitution to record or a request to
        refuse.
    mtp_admission_starvation_steps: max consecutive steps a given
        MTP K-bucket may be skipped before the scheduler force-
        picks it. Admission-time bucketing
        (:meth:`Scheduler._bucket_mtp_k_uniform`) picks one K per
        step (the dominant bucket). Without a starvation guard, a
        heavy K=1 stream could permanently delay a single K=7
        request. Default ``4`` — the K=7 row waits at most 4 steps
        (a few ms at typical decode rates) before the scheduler
        forces a K=7 slate.

One additional resident model in a stable-VA multi-model pool.

The boot model is described by the top-level ``ServerConfig`` (its
``model`` / ``cache`` / backends). Each ``PoolMemberConfig`` declares an
EXTRA model that the startup hook prepares (built + captured under its own
tag namespace, then offloaded to host) so a request naming its
``served_name`` auto-switches the resident model recapture-free.

The axes that legitimately differ per member are exposed; everything else
(batch, cuda_graphs, device, …) is inherited from the boot ``ServerConfig``
so the members share one runtime profile. A member may point at a different
model ``path`` (a genuinely different model, possibly a different
ARCHITECTURE) or reuse the boot path with a different ``max_context`` /
calibration (the common "same weights, different context window / quant"
case).

Cross-architecture members (a different model CLASS, not just a different
config of the boot weights) generally also need their OWN dtype + attention
backends — e.g. a tkv-K4V4 boot model paired with a bf16 dense member, or
two models whose KV head_dim differs. Those are overridable here; ``None``
inherits the boot value so the same-weights case stays a one-liner.

``path`` / ``max_context`` / ``calibration_path`` are CONVENIENCE ALIASES
that map into the GENERAL overlay (see
:mod:`arbi_serve.config_overrides`) — they are not a parallel schema. Any
other overridable param goes in ``overrides`` (``{param_name: value}``), so
a pool member is just a named config variant prepared at boot. The same
overlay machinery drives the live ``POST /v1/admin/config_override`` path.

Attributes:
    served_name: the pool KEY — the value clients put in the request
        ``model`` field to select this member. Must be unique across the
        pool (and distinct from the boot model's served name).
    path: model directory. ``None`` inherits the boot model's path.
    max_context: per-request max context for this member. ``None``
        inherits the boot ``cache.max_context``.
    calibration_path: TKV calibration bundle for this member. ``None``
        inherits the boot ``cache.calibration_path``; the empty string
        ``""`` means EXPLICITLY UNCALIBRATED (a cross-arch tkv member has
        its own head_dim/layer layout — inheriting the boot model's bundle
        would apply wrong-shaped centroids and is refused loudly at boot).
    dtype: weight dtype for this member (e.g. ``"bfloat16"``). ``None``
        inherits the boot ``model.dtype``. A cross-arch bf16 member paired
        with a tkv boot model sets this so its weights load in the right
        precision.
    attention_backends: ``"kind:name"`` specs registered for this member
        (e.g. ``("paged_kv:tkv-bypass",)``). ``None`` inherits the boot
        ``attention_backends``. A bf16 member alongside a tkv-K4V4 boot
        model overrides this so it does NOT try to build a tkv backend it
        has no calibration for.
    default_backends: which registered spec is active per ``StateKind`` for
        this member. ``None`` inherits the boot ``default_backends``.
    overrides: general sparse delta of any other overridable param
        (``{param_name: value}``) applied on top of the boot config for
        this member — the same registry the live admin API uses.

Speculative-decoding (MTP) config.

Attributes:
    enabled: master switch. When False, ``SamplingParams.mtp_k`` is
        ignored and no draft head is built.
    n_draft: number of speculative tokens proposed per step.
        ``0`` means "read from the bundled config"; required when
        ``draft_model_path`` is set (external drafter cannot derive
        a default — it is a plain decode model with no MTP head).
    rejection_sampling_enabled: enable Leviathan-2023 / Chen-2023
        rejection sampling so MTP can run on stochastic
        (``temperature > 0``) requests with provably-correct output
        distribution. When ``False`` the verify pass is greedy-only
        and stochastic requests fall back to the K=1 decode path
        in :func:`arbi_serve.engine.submission._build_request` —
        correct distribution, no spec speedup. Default ``True``;
        an operator escape hatch if rejection sampling regresses.
        Greedy (``temperature == 0``) requests are unaffected and
        run the byte-identical greedy verify path either way.
    draft_model_path: HF-layout directory of a separate (smaller)
        draft model. When set, the engine wires
        :class:`arbi_serve.spec_decode.external_drafter.ExternalModelDrafter`
        instead of the bundled head. The drafter must share the
        main model's tokenizer (sha256 of ``tokenizer.json`` is
        asserted at boot) and vocab size; see the 8 boot asserts
        in ``ExternalModelDrafter.__init__``.
    draft_model_dtype: explicit drafter dtype override. ``None``
        (the default) inherits the main model's dtype; setting it
        opt-in unlocks cross-dtype drafter (e.g. fp16 drafter +
        bf16 main). Mismatched dtypes can drop accept rate by
        30-50% silently — opt-in only.
    draft_model_max_seqs: maximum concurrent rows the drafter pool
        holds. ``None`` falls back to ``cfg.batch.max_batch``. The
        drafter's KV-pool sizing is
        ``max_seqs × (max_context + max_k)`` worth of pages.

Coerce a ``max_batch`` input to the concrete concurrency, or refuse.

``max_batch`` is a KNOWN quantity from the start of boot — the
:class:`BatchConfig` default or the operator's explicit number. There is
no deferred sentinel: a value resolved from the profiled KV pool would
force every boot-time consumer that needs the served width (the
multi-shape activation profiler's decode/mixed shapes, the capture
ladder, the EXL3 reconstruct row bound) to run twice or to price itself
against a worst-case stand-in instead.

``"auto"`` is therefore refused by name rather than coerced: a silent
substitution would serve a concurrency the operator never chose.

True when an operator STATED ``savepoint_max_bytes``.

The repo's rule for every other configured value applies here too:
**exposing a value at its default is not pinning it away from the
default** (``tests/test_flag_contracts.py`` states it for the deploy
pins). So "stated" is "differs from the dataclass default", read off the
field rather than restated as a literal — a second copy of the default
would drift the first time the default moved, and the failure would be a
cap that silently stopped binding.

``0`` is stated too: it means the operator turned the savepoint store off,
which is a decision, not an absence.

Collapse the convenience aliases + ``overrides`` into ONE sparse
delta keyed by canonical param name (the general overlay).

The aliases map straight onto their registry param names;
``overrides`` is merged on top. ``None`` aliases
(inherit-the-boot-value) are omitted. ``calibration_path=""``
survives as-is — apply_overrides turns it into "explicitly
uncalibrated + allow_uncalibrated_tkv".

A copy with the per-model drafter-SOURCE fields reset to defaults.

Deriving a DISTINCT pool member (a boot ``--pool-member`` or a runtime
add-resident) starts from this so the member does NOT inherit the boot
model's assistant / draft-model / DFlash checkpoint (which is wrong for a
different model); the member's own overrides then layer its source on top.
``enabled`` / ``n_draft`` / the rest inherit — the member's enable and
depth still default from the boot config (mirrors the per-member enable).

The drafter this config loads, as a stable ``"field=value"`` string.

``""`` when MTP is off or the drafter is the checkpoint's OWN bundled
head — which is an identity too, and a different one from every
explicit source.

This is a CONFIG read, never a read of the loaded drafter, so it is
the same string at every point of a boot and reproduces across boots.
That is what lets the boot-manifest fingerprint fold it
(:func:`~arbi_serve.runtime.boot_manifest.build_fingerprint`): the
drafter attaches BEFORE the activation profile runs, and attaching one
changes what the profiled forwards allocate — a DFlash drafter arms the
target model's tap, and an EXL3 drafter re-runs
``reserve_exl3_reconstruct_scratch`` / ``pin_exl3_kernel_shapes``
against its own linear geometry, which can widen the shared reconstruct
scratch and move the pinned GEMM the profiled prefill dispatches. Two
boots that differ only in drafter therefore measure different peaks,
and before this they shared a manifest key.

Fail loud on a mutually-exclusive / incomplete drafter source.

A member may name at most ONE explicit source (external draft-model XOR
DFlash XOR EAGLE3 assistant); an external / DFlash drafter requires an
explicit ``n_draft >= 1`` (a plain decode model / block-diffusion head
carries no MTP K count); any explicit source requires ``enabled``. The
boot config is already validated at CLI build, so re-validating a
resolved member cfg is idempotent — it only catches a bad per-member
override combination (never silently drops the source).

Single-source param registry + sparse-delta overlay for live config override.

Each server config variable is both a CLI arg (at boot) and an admin-API param
(at runtime) — one registry, two interfaces. The admin API takes a sparse
delta against the currently-loaded config: ``{"max_batched_tokens": 4096}``
flips one thing; everything else inherits from what is live.

Each resulting config variant is prepared and parked once via the stable-VA
pool (``stable_va_controller.prepare_model_pool`` / ``aswitch_to``);
re-selecting a known variant is an instant swap, a new variant pays a
one-time prepare.

This module maps a canonical *param name* to where its value lives — a
dotted path into :class:`~arbi_serve.config.ServerConfig` or a
:class:`~arbi_serve.runtime_flags.RuntimeFlags` field — plus the metadata
the admin path needs:

  * ``capture_affecting`` — a delta touching this param requires a new member
    to be built and captured (cudagraph capture / VRAM-budget / pool sizing
    changes), then parked and swapped. A non-capture-affecting param either
    re-snapshots on a member's build (snapshotted-at-build runtime flags) or
    is read fresh per step via the active member's flags (the 3 fresh-read
    flags).
  * ``runtime_flag`` — True when the param maps to a RuntimeFlags field (vs a
    ServerConfig path).

The registry is keyed by the canonical param name (the same token a client puts
in the admin body). ``PoolMemberConfig``'s ``path`` / ``max_context`` /
``calibration_path`` survive as convenience aliases that map into the overlay
(see :data:`ALIASES`).

One overridable param: where its value lives + its rebuild ``scope``.

``target`` is either a dotted ServerConfig path (``"batch.max_batched_tokens"``,
``"cache.max_context"``, ``"gpu_memory_utilization"``) when ``runtime_flag``
is False, or a RuntimeFlags field name (``"split_mixed_decode_prefill"``)
when ``runtime_flag`` is True.

``parse`` coerces a JSON-decoded value (which may arrive as a str from a CLI
or a typed JSON value from the API) to the field's native type.

``capture_affecting`` and group-hash membership are derived from ``scope``;
they are not stored, so the two facts can never drift.

Coerce and refuse an unsupported ``async_output_depth`` on the live
/v1/admin/config_override path.

Only 0 (off) / 1 (default) are supported; depth>=2 (an extra-FIFO
pipeline that was measured, more than once, to buy no throughput and
cost TTFT — see :data:`arbi_serve.runtime_flags.MAX_SUPPORTED_ASYNC_OUTPUT_DEPTH`'s
docstring) was removed rather than kept as a never-recommended option.
Reject the live-swap the same way boot refuses it — the ``ValueError``
surfaces as a 400 refusal from the config_override endpoint.

Coerce and refuse an unusable split-mixed cost threshold on the live
``/v1/admin/config_override`` path.

Same guard boot runs, so the threshold cannot be entered out of range
through either door; the ``ValueError`` surfaces as a 400 refusal.

Coerce + validate a live ``watermark_key`` override.

``None`` / ``""`` / ``"none"`` / ``"null"`` clears the key → watermark OFF.
A non-empty value is parsed as a 64-bit integer key (decimal or REFUSED — the
same guard boot runs (``engine.build_phases_load``), so the key can be
entered through neither door. The original string form is stored (the flag
field is ``str | None`` and ``get_watermark_cache`` re-parses it on each
cache rebuild). The ``ValueError`` surfaces as a 400 refusal.

``full`` / ``half`` / ``quarter`` / an explicit N.

The three words are FRACTIONS of the loaded checkpoint's own lm_head, so
the cut moves with the model; an explicit N is an assertion that this
width is servable, and boot refuses it loudly when it is not (the binding
constraint is the 128-wide runtime Hadamard, not divisibility by 16).

``derive`` / an explicit N / ``0``.

``derive`` is the default and reads the window off the checkpoint's own
added token ids at boot; an explicit N is an assertion that this trailing
width is servable, and 0 asks for a plain prefix.

``0`` / ``1`` / ``auto`` — the tri-state VERIFY-class routing mode.

Folds the boolean spellings the flag carried while it was a plain bool, so
an existing override that says ``true`` keeps asking for the same leg.

Coerce an ``activation_quant`` override onto the ``awq_no_a8`` tri-state.

Mirrors the ``--activation-quant`` CLI resolver's discoverable vocabulary
(``arbi_serve.cli.parse_helpers._resolve_activation_quant``):

  * ``bf16`` (== ``w4a16``) → ``True``  — W4A16 everywhere,
  * ``fp8``  (== ``w4a8``)  → ``False`` — W4A8 on any fp8-MMA arch,
  * ``auto``                → ``None``  — arch default (W4A16 on sm89).

A bool / ``None`` passes through so the key round-trips its own resolved
signature value. Anything else FAILS LOUD (no silent substitution).

Registry param names whose rebuild ``scope`` is one of ``scopes``.

The single derivation point for :mod:`arbi_serve.engine.params_hash`: it
builds its model/backend/recipe sets from this, so a capture-affecting
registry param is automatically part of the member-identity group hash —
the value cannot reach ``cfg`` on a swap without also fragmenting the group.

Validate + split a sparse delta into (cfg_deltas, flag_deltas).

Each is ``{canonical_name: parsed_value}``. ``cfg_deltas`` covers
ServerConfig-path params (incl. the model-select trio); ``flag_deltas``
covers RuntimeFlags params. Raises ``ValueError`` on an unknown param.

Fail loud when THIS delta touches ``gdn_prefill_capture`` but the
resulting ``prefill_capture`` mode is ``"eager"``.

``gdn_prefill_capture`` is a GDN-hybrid opt-in *modifier* consumed by
TWO independent sweeps, not one:

  * the whole-forward sweep's model pre-flight
    (``runtime.capture.preflight._can_capture_prefill``), reached only
    when ``prefill_capture == "full"``;
  * the piecewise per-layer sweep
    (``engine.capture_admin.layer.precapture_layer_graphs``), reached
    only when ``prefill_capture == "piecewise"``.

So the flag is *not* redundant with ``prefill_capture == "full"``
specifically — collapsing the two would silently break (or force-escalate
the VRAM/capture-time cost of) an operator's legitimate
``prefill_capture="piecewise"`` + ``gdn_prefill_capture`` combination.
It genuinely IS orthogonal to mode selection: a modifier that applies
inside whichever capture mode is chosen.

What both sweeps agree on: with ``prefill_capture == "eager"`` NEITHER
sweep ever runs, at all — every prefill chunk goes eager regardless of
``gdn_prefill_capture``. In that mode the flag is unconditionally inert.
An operator who deliberately touches it (either direction) in the same
delta that leaves (or puts) the config in eager mode almost certainly
means to be seeing prefill-capture behavior change and would silently
get nothing — the exact trap an ablation harness hits when it
live-flips the flag against an eager-mode base config and then reads
a flat ``prefill_capture_sweep`` zero. Refuse it here, at
config-apply time —
covers both the live ``POST /v1/admin/config_override`` path and pool
members prepared at boot (``server/engine_boot.py`` calls this same
function), so the mismatch is caught the moment it's misconfigured
instead of discovered later via a silent zero (see
``engine.capture_admin._plans.prefill_sweep_dormancy_reason`` for the
complementary POST-hoc observability once a build already ran with a
dormant sweep — e.g. a raw CLI + env-var boot, which does not go through
this registry at all).

Strip the boot-sequence instructions a VARIANT config must not carry.

``inprocess_calibration`` says what the BOOT SEQUENCE does after it builds
— come up in ``tkv-bypass``, calibrate on the live engine, then swap. It
describes one pool: the one the calibration traffic runs against. It is not
a property of a model config, so a variant derived from that boot config
must not inherit it, and every member the deferred-calibration boot goes on
to build is derived from it — starting with the very member the swap builds
to SERVE.

Two sizers read it as a property of the pool in front of them:
:func:`~arbi_serve.engine.build_phases_kv._cap_pages_for_deferred_calibration`
caps the KV ceiling at the calibration basket's footprint, and
:func:`~arbi_serve.engine.serving_realization.realization_refusal` declines
to measure the serving floor. Both are correct for the calibration's pool
and wrong for a serving member — the first leaves it a pool sized for the
basket, the second sizes that pool against a floor nothing measured.
Clearing the instruction here is what makes either impossible on a variant,
rather than asking each sizer to work out whose pool it is holding.

Derived identically on every rank from a config every rank holds, so the
ranks stay in lockstep; and it is not a signature param
(:func:`_signature_payload`), so no variant key moves. Returns ``cfg``
itself when there is nothing to strip, which is every serving boot.

Apply a sparse delta against (cfg, flag_overlay).

Returns ``(new_cfg, new_flag_overlay, capture_affecting)``:

  * ``new_cfg`` — ``cfg`` with every ServerConfig-path delta applied
    (frozen-dataclass copies; ``cfg`` is untouched).
  * ``new_flag_overlay`` — the prior overlay merged with the new flag
    deltas and reduced to a DELTA against the boot flags
    (:func:`overlay_delta_vs_boot`). This is the per-member RuntimeFlags
    overlay the engine installs around the member's build and on switch.
  * ``capture_affecting`` — True iff any delta param is capture-affecting,
    i.e. a new member must be built+captured (vs an instant runtime-only
    re-snapshot / fresh-read).

Raises ``ValueError`` — same fail-loud discipline as an unknown param or
an invalid ``prefill_capture`` mode — when this delta touches
``gdn_prefill_capture`` while the resulting ``prefill_capture`` is
``"eager"``: that combination can never do anything (see
:func:`_check_gdn_prefill_capture_has_an_active_sweep`).

``value`` as this param's OWN parser spells it.

The registry parser is the single definition of a param's value, so it is
also the only thing that can say whether two spellings are the same value.
Both spellings reach the same field: the CLI stores what it was handed
(``config_builder`` keeps ``--draft-vocab-prefix N`` as text) while the
override endpoint stores what :func:`split_overrides` parsed. Comparing
them raw makes a re-statement of the live value read as a delta, and a
capture-affecting delta is a member build plus a recapture for no change.

A value the parser refuses is handed back unchanged: a signature must
always be computable, and a field the registry cannot parse still keys
itself.

``flag_overlay`` with every entry that restates the BOOT value dropped.

A member's flag overlay is a DELTA against the env-derived boot flags: an
absent key resolves through ``RuntimeFlags.from_env()``, which is also what
:func:`_signature_payload` signs it as. An entry equal to that baseline
therefore changes neither the resolved value nor the signature — but it is
still a key the member record carries, a park/wake re-installs and a
derived member inherits, so it outlives the request that wrote it as an
override of a value nobody is overriding. Dropping it is what makes an
override that returns a flag TO its boot value clear the flag instead of
pinning it there.

Compared through :func:`canonical_value`, so the two sides are the same
value however each of them was spelled.

The (small, stable) set of values that define a variant for keying.

Only the registry-covered params contribute — two configs that differ only
in an un-overridable field still serve the same variant. The signature is
over the current value of every registry param (resolved against cfg /
overlay), so an explicit override that happens to equal the base value
re-selects the same member (idempotent). Every value goes through
:func:`canonical_value` first, so that promise is about the parsed value
and not about how the caller spelled it.

A param the engine resolves for itself contributes its DECLARATION rather
than the resolved value (see :data:`_ENGINE_RESOLVED`), because the resolved
value depends on the realized pool and so on which members are already
resident -- signing it would make the same config sign differently each time
it was built.

A runtime-flag param resolves to its effective value: the overlay entry when
set, else the engine's boot baseline (``RuntimeFlags.from_env()`` — env/CLI
derived, not merely the dataclass default). This is what makes the
idempotence promise hold for RuntimeFlags params too: flipping a
capture-affecting flag and flipping it back to its baseline yields the boot
signature again, so the swap-back re-selects the boot member (an instant
switch) instead of building a redundant identical variant — which, under TP
at gmu-0.99, would host-OOM parking the sibling to build the third member.

Resolve the active backend spec per StateKind for the variant key,
mirroring the active-selection in :mod:`arbi_serve.engine.build`:
``default_backends`` win their kind's slot; remaining kinds fall to the
first matching ``attention_backends`` spec. Returns ``{kind: spec}`` sorted
for a stable signature.

A short stable hash keying a variant for the stable-VA pool.

Same (registry-param) values ⇒ same key ⇒ an instant re-select swap; a
changed value ⇒ a new key ⇒ a one-time prepare. Deterministic across
processes (sorted-keys JSON → sha1, first 12 hex).

Top-level config aggregate. Parsed from CLI args.

Multi-backend wiring:
  - ``attention_backends``: ``"kind:name"`` specs registered on the
    engine (e.g. ``("paged_kv:tkv-bypass", "paged_kv:tkv-k4v4")``).
  - ``default_backends``: which spec is active at boot per
    :class:`StateKind`. At most one entry per kind; if a kind is
    omitted, the engine picks the first registered spec for it.

Distributed:
  - ``parallel_cfg``: TP / EP topology. Single-process default
    (TP=1, EP=1) runs without ``torch.distributed`` initialization.

Distributed-runtime primitives for arbi-serve.

Two surfaces:

  - :mod:`arbi_serve.distributed.parallel_state` — :class:`ParallelConfig`
    plus the ``init_distributed_environment`` /
    ``destroy_distributed_environment`` lifecycle and the
    ``get_tp_group()`` / ``get_ep_group()`` /
    ``get_attention_tp_group()`` / ``get_attention_dp_group()``
    accessors. Every TP=1 / EP=1 path is a no-op: no NCCL handshake, no
    extra collective.
  - :mod:`arbi_serve.distributed.dp_attention` — the attention-DP
    boundary: gather the per-rank token blocks before the FFN, take the
    local slice after.
  - :mod:`arbi_serve.distributed.launch` — helpers for code paths that
    need to know whether the process was launched under torchrun.

Async-TP: the TP collective decomposed into the matmul that hides it.

WHY THIS EXISTS. A TP>1 block's row-parallel exit all-reduces its partial
sums, and that all-reduce is a KERNEL — it occupies the same SMs as the GEMM
you would like to hide it behind, so a second stream queues the two rather
than overlapping them (#2170, #709). Async-TP does not run the collective
beside the GEMM at all: it decomposes the GEMM along the sharded dimension and
moves each chunk as the mainloop finishes producing it, so the transfer lives
inside the compute rather than in front of it.

Measured on the two peer-capable link regimes we have, one decoder layer's
linear+collective structure at the 27B's per-rank TP2 shapes, against a
collective-free floor of the same four matmuls (#2172):

* NVLink pair: the layer's exposed collective goes from 0.692 ms to 0.002 ms;
* PCIe peer-capable pair: from 1.546 ms to 1.038 ms.

Both bit-identical to the NCCL reference on the fused legs. This is a
SCHEDULING change and it must stay one: the arithmetic is the same matmul and
the same summation, differently ordered in time, not differently rounded.

WHAT IT REQUIRES, AND WHAT HAPPENS WHERE THAT IS ABSENT. The fused ops take
their workspace from ``_SymmetricMemory.empty_strided_p2p`` and dereference
the PEER's buffer directly, so the group must have peer access. That is the
common case on multi-GPU serving hardware and it is the mainline path here.
Where it is absent the engine takes another tier and SAYS SO BY NAME -- the
copy-engine host-staged exchange, or NCCL. A path that quietly does nothing is
the failure mode this seam exists to avoid, so every decline is counted with a
reason and the reasons are distinct.

WHAT IT IS NOT. It is not a win at every shape. The fused decomposition costs
a chunked mainloop, and past some contraction width the chunking costs more
than the collective it hides -- K=24576 regressed on all three pairings
measured. So the transport is admitted by a PREDICATE, and the predicate is
MEASURED at boot rather than written down as a constant: a threshold that is
not measured is a knob, and a knob is not a fix.

What this boot measured, and therefore what it will route.

``admits_k`` is the set of row-parallel contraction widths the boot probe
found the fused leg at or under the exposed all-reduce for. A SET and not
an interval because the crossover is not guaranteed monotone -- the fused
op's chunking interacts with the GEMM's tiling, and asserting an interval
would extrapolate between two measured points into a shape nobody timed.

``group_name`` is the process group the symmetric-memory workspace was
rendezvoused on. Held so a leg cannot be issued against a different group
than the one the plan was measured for.

``all_gather(a_shard) @ b``, gather pipelined into the mainloop.

Returns BOTH the gathered activation and the product, because a
sequence-parallel region has more than one column-parallel projection
reading the same gathered rows (q, k, v and the attention gate all read
it). Gathering once and handing the rows on is the difference between one
all-gather per region and one per projection.

Is the fused row-parallel leg BIT-IDENTICAL to, and at or under, what it replaces?

Two independent questions and both are measured, because the answer to
neither generalises across shapes. Measured on the two peer-capable links
available (#2172): every width from 3072 to 16384 came out bit-identical
and K=24576 did NOT -- the fused op chunks the contraction differently at
that width, so the partial sums are added in a different order and bf16
rounds differently. That is the same width that regressed on time, but the
two facts are independent and a boot that found only one of them must
still refuse.

The speed denominator is what the engine DOES today -- ``mm`` then an
exposed all-reduce -- not the unfused reduce-scatter pair. Comparing
against the pair prices the sequence-parallel rewrite and the fusion
together and cannot say which one paid.

Does symmetric memory actually rendezvous on this group?

Peer access is necessary and NOT sufficient: the driver can report that two
devices may address each other and the symmetric-memory allocator still
refuse, and the refusal that matters is the real one rather than an
inference from the capability bit. Its own function because it is its own
question, and because a test that has to stub ``torch.distributed`` to ask
it ends up asserting on the stub.

Resolve, MEASURE and arm the async-TP transport. BOOT ONLY.

Order matters and each step declines by name:

1. the operator has to have asked (``enabled``);
2. the group has to be non-trivial and CUDA;
3. the group has to have PEER ACCESS -- read off the boot topology probe,
   never re-derived here, because two transports resolving the same
   topology separately is how they come to disagree;
4. ``symm_mem`` has to rendezvous -- peer access is necessary and not
   sufficient, and the refusal has to be the real one rather than an
   inference from step 3;
5. and then each shape is TIMED against what it would replace. A shape
   that does not win is not admitted. This is the step that makes the
   predicate a measurement rather than a constant: the crossover moves
   with the card, the link and the checkpoint, so writing it down as a
   number would be right on the box it was written on and nowhere else.

Returns the plan, or ``None`` with a counted reason.

May a row-parallel leg contracting ``k`` take the fused path?

Every term is a property of the WEIGHT and of the boot measurement.
Nothing rank-local may enter here: the two ranks would disagree about
which transport is in use, and that failure is a hang, not a wrong
answer -- the same rule ``copy_engine_takes`` is written under.

Bit-identity FIRST, and it is not negotiable against speed.

This transport is a scheduling change: the same matmul and the same
summation, differently ordered in time. A width where the fused
decomposition sums the chunks in a different order is a different
arithmetic, and no throughput number buys that.

A two-rank all-reduce that runs on the GPUs' COPY ENGINES, not their SMs.

Why this exists at all. On a box where the two cards sit under a PCIe host
bridge with peer-to-peer disabled, every all-reduced byte is staged through
host memory, and NCCL's shared-memory all-reduce does that staging from
*inside a kernel*. A kernel occupies SMs; so does the GEMM you would like to
hide the transfer behind. A second CUDA stream therefore does not overlap the
two — it queues them, and the collective is fully exposed on the prefill
critical path.

A ``copy_(non_blocking=True)`` between page-locked host memory and the device
runs on the **copy engines**, which are separate hardware from the SMs. The
same bytes then move underneath compute instead of in front of it. That
difference — the engine the transport occupies — is the whole reason this
module exists, and it is the one property that must survive every change here.

What that costs, and the shape it forces. A copy-engine exchange needs a peer
handshake ("my partial has landed in the shared buffer"). For the transfer to
stay off the SMs the handshake must not be a kernel: a spin-waiting kernel
holds SMs for the whole peer latency and reproduces exactly the behaviour this
module exists to avoid. The kernel-free device-side primitive for it
(``cuStreamWaitValue32``) is unsupported on GeForce Ada — the attribute
``CAN_USE_STREAM_MEM_OPS`` reads 0, and the enqueue is *accepted* and then
never satisfied. So the handshake is on the HOST, split across
:meth:`Exchange.start` and :meth:`Exchange.wait`:

* ``start`` enqueues the device->host copy of this rank's partial and returns
  immediately — the caller then enqueues independent compute;
* ``wait`` blocks the launch thread until this rank's copy has landed and the
  peer has published its own, then enqueues the host->device copy of the peer's
  partial and the add.

The launch thread is inside ``wait`` for most of a chunk. It must therefore
release the GIL while it is there, or it stalls its own rank's event loop for
the duration — see :class:`_PeerFlags`.

Scope. Two ranks only, prefill-width partials, on a host-staged link. It is
not capture-safe and must never run under cudagraph capture: the handshake is
host-side control flow, and a captured graph replays device work without ever
re-running it. See :meth:`Exchange.assert_not_capturing`. That makes it viable
only while prefill is executed rather than replayed —
``ServerConfig.prefill_capture == "eager"`` — and
:meth:`arbi_serve.distributed.parallel_state.GroupCoordinator.arm_copy_engine`
refuses to arm it under any other mode, at boot, so a replayed prefill can
never be priced on this transport by way of a silent NCCL fallback.

A handshake wait exceeded its deadline.

Always raised rather than waited out: a wedged peer on a leased GPU costs
the whole card, so every wait in this module is bounded and every bound
fails loud.

How the launch thread waits on the peer's flag.

Both release the GIL, which is the requirement — the launch thread is
inside the wait for nearly all of a prefill chunk, and a GIL-holding spin
would stall the rank's own event loop for that whole time (the engine has
already been bitten once by a host-side loop holding the interpreter).
They differ in how they pay for it:

``FUTEX``
    Blocks in the kernel on the shared flag word and is woken by the peer's
    ``FUTEX_WAKE``. Costs no CPU while waiting; pays a wake-up latency.
``SPIN``
    Polls the flag, yielding the interpreter each iteration. Lowest
    latency, burns a core.

Which one is right is a measurement, not a preference, so both are
reachable and the transport records which one produced a number.

The shared 32-bit words the two ranks hand each other.

Every word is monotonic and never reset. A rank that has already moved on
must not be able to strand a peer still finishing an earlier exchange, and
a monotonic word makes "has the peer reached ``n`` yet" a comparison rather
than a protocol.

Two words per rank, and both are needed:

``pub[r]``
    exchanges rank ``r`` has LANDED in its own host slot — the peer may
    read slot ``(pub-1) % _SLOTS``.
``ack[r]``
    exchanges rank ``r`` has finished DRAINING out of the PEER's slot — the
    peer may overwrite slot ``(ack-1) % _SLOTS``.

``pub`` alone is not enough. It says the peer published, not that the peer
has read what this rank published, so a producer running ahead could
overwrite a slot whose bytes the peer had not yet copied out. That is a
silent wrong-value failure, not a hang, which is why the ack is a
structural part of the protocol rather than a safety margin.

One rank's half of a two-rank copy-engine all-reduce.

Construction is a BOOT step: it maps and page-locks the shared host slots
and allocates the device landing buffers, so a reduce on the serving path
allocates nothing and registers nothing. Both ranks must construct it, and
the shared-memory names come from a value broadcast over the group, so the
pairing is derived rather than assumed.

The two halves are deliberately separate. Calling :meth:`start` and then
:meth:`wait` back to back is a correct, synchronous all-reduce and is what
:meth:`all_reduce` does; it recovers nothing, because nothing is enqueued
in between for the transfer to hide under. The win is entirely in a caller
that puts independent compute between them.

Bounded wait on a CUDA event that does not hold the GIL.

``Event.query`` is polled from Python only to bound the wait; the blocking
is done by ``synchronize``, which releases the interpreter lock. A bare
``synchronize`` would be unbounded, and an unbounded wait on a leased card
is how a wedge costs minutes instead of seconds.

Store ``value`` and wake anyone blocked on the word.

The store must be visible before the wake, which it is: the wake is a
syscall and the store is to a shared mapping the syscall follows.

Block until the peer's word reaches ``target``. Releases the GIL.

``ctypes`` foreign calls and ``time.sleep`` both drop the interpreter
lock, which is the property this method exists to guarantee; a bare
Python ``while`` over the flag would not.

Refuse to run inside cudagraph capture.

The handshake is host control flow: a captured graph records device
work and replays it without re-running any of this, so a captured
exchange would replay stale flags and copy whatever the host slots
happened to hold. There is no way to make that safe, so it is refused
rather than documented.

Enqueue this rank's partial to its shared host slot.

Returns a ticket for :meth:`wait`. The caller should enqueue
independent compute before waiting on it — that is the entire point of
the split.

Complete exchange ``ticket``, summing the peer's partial into ``out``.

``out`` is the same tensor whose contents were handed to :meth:`start`;
the sum lands in place, which is what every ``RowParallelLinear`` caller
wants and what makes this a drop-in for the NCCL reduce.

Synchronous all-reduce — correct, and no faster than NCCL.

Nothing is enqueued between the two halves, so nothing hides the
transfer. This is the correctness entry point and the fallback for a
caller with no independent work; the win lives in :meth:`start` /
:meth:`wait` used apart.

The attention-DP boundary: gather before the FFN, slice after.

Under ``ParallelConfig.attn_dp_size > 1`` an attention-TP set of
``attn_tp_size`` ranks owns whole REQUESTS — their tokens, their KV — and
shards only the heads among its own members. The MoE / dense FFN stays
parallel over the full TP group, so the two halves of a decoder layer
disagree about what a rank's rows mean:

* through attention, a rank's rows are ITS requests' tokens, and its
  row-parallel exit reduces over the attention-TP group only;
* through the FFN, every rank must see the SAME rows — the whole step's
  tokens — because the expert / intermediate shards it holds are a slice
  of one computation over all of them.

So the layer gathers the attention-DP peers' token blocks into one buffer
before the FFN and takes its own contiguous slice back out after. The
residual add stays on the local slice, which is why only the FFN input and
output cross the boundary.

RAGGED TOKEN COUNTS ARE THE NORMAL CASE. The peers hold different
requests at different lengths, and a rank with no admitted request
contributes zero rows. The gather therefore pads each rank's block up to
the step's widest block, all-gathers the fixed-size padded blocks, and
concatenates the valid prefixes in attention-DP rank order. Padding costs
one copy of ``(max_count - local_count)`` rows and NO arithmetic: every
value that reaches the FFN is a byte copy of the value the corresponding
single-rank run computed, so the boundary contributes no numerical
difference of its own.

The per-step token layout is exchanged ONCE, in :func:`attn_dp_step`, not
per layer: the counts are a step property, and re-deriving them per layer
would multiply a host-syncing collective by the layer count. Every layer
inside the context reads the installed layout, and a layer that reaches
the boundary with no layout installed raises rather than inventing one —
a rank that silently disagreed with its peers about the row layout would
gather garbage instead of failing.

How one step's tokens are distributed over the attention-DP ranks.

``counts[d]`` is the number of rows attention-DP rank ``d`` holds;
the gathered buffer lays the blocks out contiguously in ``d`` order,
so ``offsets[d]`` is the prefix sum. ``local_rank`` is this process's
index into both.

All-gather each attention-DP rank's token count for this step.

Trivial group (``attn_dp_size == 1``) short-circuits to the single-block
layout with no collective. The gate reads ``attn_dp_size``, which IS
the attention-DP coordinator's ``size`` — the group is constructed
from it, and a test pins the two equal on every rank. Reading the
number rather than the coordinator keeps the per-layer boundary from
materializing a group object on a path that will not use it.

Install this step's attention-DP token layout for the model forward.

Entered once per forward, around the whole layer stack. Stacked entry
is allowed and restores the outer layout on exit (a draft-head forward
nested inside a verify forward is a different step's row layout, and
the inner one must not leak).

Gather the attention-DP peers' row blocks into one buffer.

Returns ``local`` unchanged when the attention-DP group is trivial, so
the ``attn_dp_size == 1`` path allocates nothing and copies nothing.
See :func:`exchange_dp_token_layout` for why the gate reads the size
rather than the coordinator.

Rows land in attention-DP rank order. Ragged blocks are padded to the
step's widest block for the fixed-size all-gather and the pad rows are
dropped on the way out, so the result is a pure re-arrangement of the
peers' values.

Take this rank's contiguous row block back out of a gathered buffer.

The inverse of :func:`attn_dp_gather` for the row axis: pure indexing,
no collective, no arithmetic. Returns ``gathered`` unchanged when the
attention-DP group is trivial.

TP-aware cudagraph capture context.

Capturing a forward whose layers issue ``all_reduce`` collectives is
hairy: NCCL allocates internal communication buffers lazily (first
call), spawns its own internal CUDA streams, and PyTorch's
``torch.cuda.graph`` rejects any ``cudaMalloc`` that happens during
capture. The :func:`graph_capture` context manager:

  - hands out a single process-lifetime side :class:`torch.cuda.Stream`
    per device (the shared capture stream — see below);
  - calls ``torch.cuda.synchronize`` so any in-flight transient state
    is settled;
  - runs a warm-up ``all_reduce`` on the side stream so NCCL's
    per-stream comm buffers are allocated before capture begins
    (``torch.cuda.graph`` rejects ``cudaMalloc`` during capture);
  - yields a :class:`GraphCaptureContext` whose ``stream`` field the
    caller binds via ``with torch.cuda.stream(ctx.stream)`` and
    passes to ``torch.cuda.graph(graph, stream=ctx.stream)`` so AR
    issued on the side stream is captured into ``graph``.

A single shared capture stream is used across the whole sweep. NCCL
allocates its internal per-collective communication buffers lazily,
once per ``(communicator, CUDA-stream)`` pair, on the first
``all_reduce`` that runs on that pair. Those buffers land in the
default torch caching allocator's heap (a ``cudaMalloc`` on the warm-up
call, outside the ``torch.cuda.graph`` window) — invisible to the cuMem
``mapped_bytes`` accounting and never freed (NCCL holds them for the
comm's life). The boot capture sweep captures many shapes (decode ×
kv-page buckets + drafter-chain B×K + the MTP-verify ladder). Minting a
fresh ``torch.cuda.Stream`` per captured shape would make NCCL allocate
a fresh per-stream buffer set for every shape — a cuMem-invisible
residual that the post-capture KV grow (which sizes against
driver-free VRAM) can never reclaim, starving the KV pool. Reusing one
stream across the entire sweep makes NCCL allocate its per-stream
buffers exactly once; the freed VRAM flows straight into the
post-capture KV grow.

This is safe because the capture-time side stream is not referenced at
replay time: ``CapturedGraph.replay`` calls ``self.graph.replay()`` (a
plain ``torch.cuda.CUDAGraph.replay``), which re-issues the captured
kernels — including the captured NCCL AR — on whatever stream is
current at replay, not on the capture stream. The captured graph is
self-contained once captured; the side stream's only job is to host
the issue order and the pre-warmed NCCL state during capture. Captures
are strictly sequential (the coordinator forbids nested capture
regions), so the shared stream is never concurrently driven by two
captures.

Every non-trivial coordinator is warmed and pinned — the region
iterates
:func:`arbi_serve.distributed.parallel_state.nontrivial_coordinator_groups`
(populated at ``GroupCoordinator`` construction), never a named TP/EP
pair. A rank's collectives are split across coordinators and a
topology can make any of them trivial: ``--ep-size 2`` on two GPUs
gives ``tp_size == 1`` and ``ep_size == 2``, so the only collective in
the captured decode forward is the per-MoE-layer EP all-reduce, and a
TP-keyed region would capture it off the un-warmed, un-pinned
live-eager path (a staged round trip per layer, plus NCCL's lazy
per-(comm, stream) alloc landing inside the capture window). Any
additional coordinator (e.g. a PP/DP coordinator) is covered by
construction — it registers itself when built. Single-rank deployments
(every group trivial) still get the shared side stream so the
call-site structure is uniform.

The overlap side stream is minted and warmed here too, for the same
reason, and bound onto each non-trivial coordinator for the region —
see :meth:`GroupCoordinator.all_reduce_overlapped`. A collective forked
onto it must find NCCL's per-(comm, stream) buffers already allocated,
and the only place that can be arranged is outside the capture window.

Replay itself needs no special handling: ``CapturedGraph.replay`` runs
``self.graph.replay()``, the captured AR operations re-issue on the
current stream, and the per-rank result lands in the captured
``logits_out`` buffer the same as TP=1.

The distinct coordinators whose collectives a capture must cover.

Delegates to the process-wide registry
(:func:`~arbi_serve.distributed.parallel_state.nontrivial_coordinator_groups`)
— every coordinator registers itself at construction, so this region
covers any group a forward can issue a collective on, with no named
TP/EP enumeration to extend.

Force ``group``'s lazy NCCL per-(comm, stream) buffer alloc on ``stream``.

``torch.cuda.graph`` rejects any ``cudaMalloc`` inside the capture window,
and NCCL allocates its internal comm buffers on the first collective it
sees for a given (communicator, stream) pair. Running one tiny all-reduce
here — outside capture — moves that allocation into the ordinary caching
allocator. Fires at most once per (device, comm, stream role).

Side-stream + group-coordinator state for one capture region.

The :class:`GroupCoordinator` reads ``stream`` while the context is
active to decide whether it's running inside a captured forward;
callers MUST use ``stream`` as the cuda stream for both
``torch.cuda.stream(...)`` and ``torch.cuda.graph(graph,
stream=...)`` so the captured kernels (regular kernels + NCCL AR)
all enqueue on the same stream the graph records against.

Enter a parallelism-aware cudagraph capture region.

Use:

.. code-block:: python

    with graph_capture(device) as ctx:
        ctx.stream.wait_stream(torch.cuda.current_stream(device))
        with (
            torch.cuda.stream(ctx.stream),
            torch.cuda.graph(graph, stream=ctx.stream, pool=mem_pool),
        ):
            model.forward(...)
        torch.cuda.current_stream(device).wait_stream(ctx.stream)

The manager flips every non-trivial :class:`GroupCoordinator` (TP and
EP alike) into "graph-safe" mode for the duration of the with-block:
collectives issued by ``GroupCoordinator.all_reduce`` enqueue on
``ctx.stream`` (the same stream the graph records against), and each
communicator's NCCL per-stream scratch is pre-warmed via a tiny dummy
AR before the capture begins — on the capture stream and on the
overlap stream :meth:`GroupCoordinator.all_reduce_overlapped` forks
onto.

On exit, the prior coordinator state is restored — a subsequent
live-path forward issues ``dist.all_reduce`` on the default
stream as before.

torchrun launch helpers.

Introspects the standard ``torchrun``-set environment variables
(``WORLD_SIZE``, ``RANK``, ``LOCAL_RANK``) and surfaces them as a
:class:`ParallelConfig`. This is the single place the engine
translates env-var conventions into typed config; every other surface
reads :func:`get_parallel_config`.

Build a :class:`ParallelConfig` from the current process environment.

When launched under torchrun, ``WORLD_SIZE``, ``RANK``, and
``LOCAL_RANK`` are populated; otherwise we fall back to the
single-process default.

The user supplies ``tp_size`` and ``ep_size`` on the CLI;
``world_size`` is derived from the launch environment and must
equal ``tp_size * ep_size`` (validated by :class:`ParallelConfig`).
``moe_ep_size`` is not a factor of the world — it divides ``tp_size``
and selects how the MoE experts shard across the TP ranks.
``attn_dp_size`` likewise divides ``tp_size``: it splits the TP group
into that many sets, each owning whole requests and their whole KV.

Process-group state for arbi-serve TP / EP.

A :class:`ParallelConfig` carries ``(tp_size, ep_size, moe_ep_size,
world_size, rank, local_rank)`` and is the single source of truth for
"where am I" inside the distributed runtime. Two singleton
:class:`GroupCoordinator` instances (``_TP_GROUP`` and ``_EP_GROUP``)
wrap ``torch.distributed.ProcessGroup`` objects for TP and EP
collectives.

Two separate notions, deliberately not conflated:

* ``tp_size`` / ``ep_size`` are orthogonal factors of the world —
  ``tp_size * ep_size == world_size``. A rank belongs to exactly one TP
  group and exactly one EP group, and the two groups form a 2-D grid
  over the world. This is a genuine topology dimension (independent
  model replicas each holding a distinct expert partition) and the
  invariant is kept.
* ``moe_ep_size`` is a sharding mode for the MoE sublayer, not a
  dimension of the world. It says how the ranks of one TP group divide
  the expert set among themselves, and it divides ``tp_size``. This is
  the ``--enable-expert-parallel`` topology: attention and the dense
  projections shard TP-wise over N ranks while the MoE experts shard
  EP-wise over those same N ranks.

* ``attn_dp_size`` is the third notion: a sharding mode for ATTENTION,
  and the only one of the three that needs its own process groups. It
  divides ``tp_size`` and says how the ranks of one TP group split the
  REQUESTS among themselves. A latent-attention model (MLA) has exactly
  one KV head, so TP cannot shard its KV cache — it replicates it, and
  aggregate KV grows with ``tp_size``. Under ``attn_dp_size > 1`` each
  attention-TP set of ``attn_tp_size = tp_size // attn_dp_size`` ranks
  owns whole requests and their whole KV, while the MoE / dense FFN stays
  parallel over the full TP group. Two groups follow from that: the
  attention-TP group (ranks sharing a request's heads, over which an
  attention row-parallel partial is reduced) and the attention-DP group
  (the peers holding the other requests, over which hidden states are
  gathered before the FFN and sliced after). See
  :mod:`arbi_serve.distributed.dp_attention` for that boundary.

The second notion needs no new process group, which is why it is
expressed this way rather than by relaxing the ``tp * ep == world``
invariant. The MoE combine is a sum over every rank that holds a piece
of a token's mixture, and the pieces (this rank's expert subset ×
this rank's intermediate-dim slice) partition the world grid however
``moe_ep_size`` cuts it. The existing pair of collectives — one
all-reduce over the TP group, one over the EP group — already sums the
whole grid, so the forward's collectives are bit-for-bit the same
sequence at every ``moe_ep_size``. Only the per-rank weight-shard
assignment changes. See :meth:`ParallelConfig.expert_shard_rank`.

Single-process (``world_size==1``) is a working default. ``init_*`` is a
no-op in that mode — no NCCL handshake, no group creation, no extra
state to tear down. Every parallel module checks
``get_parallel_config().tp_size == 1`` and short-circuits its
collectives.

Distributed mode: launched under ``torchrun --nproc-per-node N``.
``init_distributed_environment(cfg)`` calls
``dist.init_process_group("nccl")`` and constructs the TP / EP
:class:`ProcessGroup` objects matching the topology in ``cfg``. The
constraint ``tp_size * ep_size == world_size`` is validated at parse
time in :mod:`arbi_serve.cli` and re-asserted here defensively.

The pair ``init_distributed_environment`` /
``destroy_distributed_environment`` lands together — the singletons are
fully reset on destroy so a test process can exercise multiple
``ParallelConfig`` variants in sequence.

Topology declaration for the engine.

Defaults are the single-process working default — every callsite
that takes a ``parallel_cfg`` argument can supply
``ParallelConfig()`` and inherit TP=1 / EP=1 / world=1 behaviour.

The constraint ``tp_size * ep_size == world_size`` is the binding
invariant — a model rank participates in exactly one TP group and
exactly one EP group. The constructor validates it.

``moe_ep_size`` is not a factor of the world: it divides ``tp_size``
and says how the ranks of one TP group split the expert set among
themselves (``--enable-expert-parallel``). See the module docstring
for why this is a sharding mode rather than a topology dimension.

Thin wrapper over a :class:`torch.distributed.ProcessGroup`.

Methods are out-of-place (return a fresh tensor). TP=1 / EP=1 paths
short-circuit to the identity (no comm issued) so single-process
call sites pay zero cost.

Cudagraph-capture mode. While
:meth:`arbi_serve.distributed.graph_capture.graph_capture` is
active for this coordinator, ``_capture_stream`` is set to the
side-stream the captured graph records against. Inside that
context, :meth:`all_reduce` enqueues the collective on the
capture stream rather than letting NCCL pick the current default;
that pin is what allows ``torch.cuda.graph(g, stream=s)`` to
capture the cross-rank op into ``g``.

Apply the bounded NCCL workspace defaults (see
:data:`_BOUNDED_NCCL_ENV_DEFAULTS`) before the communicator is created.

NCCL reads ``NCCL_*`` from the environment when a communicator is built
(``init_process_group`` / ``new_group``), so this must run before either.
Each var is set only when unset in the environment — an explicit operator
override always wins, and is left untouched (no silent clobber of a
deliberate tuning choice). No-op at TP1 (no NCCL communicator). Returns the
effective ``{var: value}`` map (operator value where overridden, bounded
default where this applied it) for the boot log so the chosen workspace
sizing is observable.

Pure environment mutation — no CUDA, no torch import. Idempotent.

Bring up the distributed runtime for ``cfg``.

Single-process (``cfg.world_size == 1``) — install the singletons
with trivial groups; no NCCL init. Multi-process — call
:func:`torch.distributed.init_process_group` with the requested
backend, set the cuda device to ``cfg.local_rank``, then construct
one TP and one EP :class:`ProcessGroup` per the topology.

Before any communicator is created the bounded NCCL workspace env
(:func:`_configure_bounded_nccl_env`) is applied so the libnccl-internal
per-channel comm buffers — resident physical that is invisible to both
torch's allocator and the cuMem accounting — are sized for the small
no-P2P TP group instead of NCCL's larger stock defaults. That invisible
workspace directly steals from the post-capture KV grow; bounding it is a
KV-reclaim lever.

Idempotent: a second call with the same ``cfg`` is a no-op; a call
with a different ``cfg`` while one is already active raises.

Return this process's global rank from the ``RANK`` env var.

``RANK`` is set by ``torchrun`` on multi-rank launches and absent
(defaulting to ``0``) on single-process runs — the same convention the
launcher uses in :func:`arbi_serve.distributed.launch.parallel_config_from_env`.

Unlike :func:`get_parallel_config`, this reads the launcher environment
directly, so it is available before :func:`init_distributed_environment`
has installed the typed config — e.g. module-import-time rank gating in the
calibration drivers and the server metric ``rank`` labels.

Return the active :class:`ParallelConfig`.

Pre-init, returns the single-process default. This matches the
"callsite can read unconditionally" contract — modules that need
TP-awareness can call this without first having to check whether
init has run.

Return ``(start, length)`` of ``rank``'s contiguous shard of ``total``.

The single source of truth for the per-rank slice arithmetic that
column-/row-parallel linears (and every quant backend's copy of them)
re-derive: ``length = total // size; start = rank * length``.

``total`` must be divisible by ``size`` — TP sharding requires even
splits, and callers validate divisibility at construction time, so a
non-divisible ``total`` here is a bug, not a recoverable condition.

Block-packed formats (nvfp4's 2-elem-per-byte / 16-elem scale blocks,
exl3's trellis blocks) derive their block range by integer-dividing
this element range: ``start // block, length // block``. That is exactly
equal to a per-backend inline derivation because ``length % block == 0``
is guaranteed by the backend's own block-alignment validation, so
``(rank * length) // block == rank * (length // block)``.

Return ``(start, length)`` of ``rank``'s shard, both multiples of ``align``.

Unlike :func:`tp_shard_range` this does NOT require ``total % size == 0``.
It splits ``total // align`` whole units as evenly as possible and hands
the remainder to the lowest-numbered ranks, so shards may differ in size
by one unit. ``total`` must be a multiple of ``align``.

Backends whose sharded axis carries a transform with a fixed block width
(EXL3's 128-wide Hadamard) are exact under an uneven split as long as
every boundary is a whole block, and are wrong under an even split that
lands mid-block. Such a backend wants this, not the even split.

Return the attention-TP :class:`GroupCoordinator` — the ranks this
rank shards attention HEADS with.

At ``attn_dp_size == 1`` this is the TP coordinator itself (the same
object), so an attention collective routed here is the TP collective it
already was. Same "read unconditionally" contract as
:func:`get_tp_group`.

Return the attention-DP :class:`GroupCoordinator` — the peer ranks
holding the OTHER requests.

Trivial (size 1) at ``attn_dp_size == 1``: one request shard, nobody to
gather from. See :func:`get_tp_group` for the pre-init contract.

Return ``(size, rank)`` of the per-rank split named ``axis``.

``cfg`` defaults to the installed :class:`ParallelConfig`. An unknown
axis raises — a typo must never silently degrade to the whole TP group,
which loads and serves without error while handing the rank the wrong
slice.

Return the :class:`GroupCoordinator` whose ranks index ``axis``.

The collective that reduces / gathers a weight cut by ``axis`` must run
over EXACTLY these ranks: reducing an attention-TP partial over the full
TP group sums activations computed from DIFFERENT tokens (the
attention-DP peers hold other requests) and is silently wrong. The
caller gates on the returned group's own ``size``, never on ``tp_size``.

Re-arm every process group's collective timeout for SERVING.

``ARBI_DIST_INIT_TIMEOUT_S`` (default 1800 s) is a COLD-BOOT bound: the
first collective on a fast rank sits behind the slow rank's whole
checkpoint load, so the rendezvous timeout has to be sized to the load.
``torch.OutOfMemoryError`` inside the MTP verify ``lm_head``
all-gather left the peer parked in ``_ALLGATHER_BASE`` for the full 1800 s.

This is the STRUCTURAL backstop for that whole class. The SPMD loop's
per-tick step-outcome agreement turns an asymmetric step failure into a
clean symmetric refuse, but ONLY when the surviving rank can get out of the
step to reach it — a rank blocked inside a collective its peer never issued
cannot. Nothing at the Python level can rescue that; the only bound is the
collective timeout itself, so past boot it must be a serving-sized number.

Returns the timeout applied in seconds, or ``0`` when nothing was re-armed
(single-rank, no process group, ``ARBI_SERVE_COLLECTIVE_TIMEOUT_S=0``, or a
torch build without the setter). Never raises.

Paired with :func:`arm_boot_collective_timeout`, which the admin
:class:`~arbi_serve.engine.critical.CriticalSection` uses to widen the bound
back to the boot value for the duration of a sanctioned mutation (a model
reload or backend swap legitimately re-loads weights while a peer waits, and
that is the COLD-BOOT rendezvous the 1800 s exists for — not a wedge).

Widen the collective timeout back to the COLD-BOOT bound.

Called on entry to a sanctioned admin mutation window
(:class:`~arbi_serve.engine.critical.CriticalSection`: backend swap, model
reload, in-process calibration). Those legitimately re-load weights or
re-run a capture sweep while a peer rank waits at the next collective —
exactly the load-bound rendezvous ``dist_init_timeout_s`` is sized for, and
exactly what the serving-sized bound would false-trip on. The section's
``__aexit__`` restores the serving bound.

No-op (returns 0) when the serving re-arm itself is disabled — then the boot
timeout was never lowered and there is nothing to widen.

Every live :class:`GroupCoordinator`, in creation order.

Dead weak refs are pruned in place. See :data:`_LIVE_GROUPS` for why
the registry is populated at construction and why creation order is
the only rank-symmetric (deadlock-free) iteration order.

Live coordinators whose collectives are real: ``size > 1``, deduped
by identity, in creation order.

The selector for any path that must cover "every collective this
process can issue" — the cudagraph capture warm/pin region and the
piecewise capture cross-rank consensus gate both key on it. Keying on
"this coordinator has size > 1", never on a parallelism flavor, is
load-bearing: ``--ep-size 2`` at ``tp_size == 1`` makes TP trivial and
the EP reduce the only collective in the forward, and any TP-keyed
selection leaves it uncovered. Identity dedup matters because
single-rank processes hand back the same object for several roles.

Route the TP + EP coordinators' ``_staging`` buffers into ``pool``.

Call once right after :func:`init_distributed_environment` (engine
build, before any collective runs / before capture warms NCCL state).
A no-op for trivial (size-1) groups — they never stage a collective —
but harmless to call unconditionally; the engine calls it whenever a
``nccl_staging_pool`` exists so the multi-process path always pools its
staging strand. See :meth:`GroupCoordinator.attach_staging_pool`.

Iterates every live coordinator (not a named TP/EP pair), so any
additional group's staging buffers pool by construction; attaching to
a trivial group is a no-op (it never stages a collective).

Drop the TP + EP coordinators' cached ``_staging`` buffers.

Called on a hot-swap reload before re-attaching the freshly-created
``nccl_staging_pool`` (the old buffers point at pages the old pool's
teardown freed). No-op for trivial groups / pre-init. See
:meth:`GroupCoordinator.reset_staging_buffers`. Iterates every live
coordinator for the same reason as :func:`attach_staging_pool`.

Rank within the EP group this process belongs to.

EP groups stride over TP groups: rank ``r`` belongs to EP group
``r // tp_size`` and TP group ``r % tp_size``. This matches the
``--tp-size N --ep-size M`` topology where the first ``N`` ranks
form one TP group, the next ``N`` form the next, etc.

Which of the :attr:`attn_dp_size` request shards this rank owns.

At ``attn_dp_size == 1`` this is ``0`` on every rank — one shard
holding every request, which is what TP-only attention is.

Ranks sharing this rank's expert shard, i.e. the width of the
MoE intermediate-dim split.

At ``moe_ep_size == 1`` this is ``tp_size``: every TP rank holds
every (local) expert and cuts only the intermediate dim.

Which of the :attr:`expert_shard_count` expert partitions this
rank owns.

Consecutive tp_ranks share an expert shard (``tp_rank //
moe_tp_size``), matching how every other TP-sharded weight in the
engine cuts by consecutive ``tp_rank`` — so a rank's MoE
intermediate slice is the same slice index its dense projections
use, and the two agree by construction.

At ``moe_ep_size == 1`` this reduces to ``ep_rank`` exactly
(``ep_rank * 1 + tp_rank // tp_size == ep_rank``), so every
``moe_ep_size == 1`` topology gets a byte-identical shard.

Sum ``x`` across the group, in place.

Every caller passes a fresh per-rank partial (a matmul/GEMM or
elementwise output produced earlier in the same forward) whose
pre-reduce value is never read again, so the reduce runs in
place — no defensive copy. ``.contiguous()`` stays: it is a
no-op when already contiguous and guards the rare strided input
NCCL would reject.

Capture-aware dispatch — the single invariant is whether the
collective's buffer address is pinned for the life of the
(captured) execution:

* Address-pinned: reduce in place on the activation, zero
  staging copies. Two ways an execution is address-pinned, and
  arbi's TP2 decode hot path hits both:
    - Compiled (``torch.compiler.is_compiling()``): the compiled
      decode/prefill artifact is recorded once and cudagraph-
      replayed verbatim — the target-model :class:`RowParallelLinear`
      collectives (``@support_torch_compile(level="block")`` blocks,
      ``fullgraph=True``) live here.
    - Raw cudagraph capture (``_capture_stream is not None`` outside
      a Dynamo trace): ``torch.cuda.graph`` records the ncclAllReduce
      against the captured-pool address and replays it verbatim —
      the eager MTP drafter chain (``ARBI_MTP_HEAD_COMPILE`` off by
      default, so the head is not Dynamo-compiled and its per-step
      ``o_proj`` all-reduce is captured raw, not under
      ``is_compiling()``) lives here.
  In both, every replay pins ``sendbuff == recvbuff`` to the same
  device address, so NCCL's address/alignment-dependent reduction
  gives the identical result each request (bit-stable, strictly
  stronger than the shared staging buffer), and SPMD ranks agree
  bit-for-bit (all-reduce broadcasts identical bytes; a world-2 sum
  has no reduction-order freedom).

* Live eager (``_capture_stream is None`` and not compiling — a
  non-captured prefill first-token / first-seen runtime shape):
  the activation is an allocator-fresh address that varies per
  request, so NCCL's protocol/chunking choice (address-keyed) would
  make the TP reduction drift between same-prompt requests. Only
  here, the reduce is staged through the boot-stable per-shape
  buffer (:meth:`_staging`) to pin the address.

Keying on the capture context makes the hot-path collective in
place by construction regardless of compiled vs raw capture, and
confines the staged round-trip to the one place its
address-pinning is genuinely needed.

Route the persistent ``_staging`` buffers into ``pool``.

``pool`` is a :class:`arbi_serve.runtime.named_pool.NamedMemPool`
(typed loosely to keep this module import-light — torch is the
only heavy dep it pulls). Called once right after
:func:`init_distributed_environment` so every subsequent staging
``torch.empty`` lands in a named, cuMem-backed, accounted pool
instead of the cuMem-invisible default heap. Idempotent re-attach
of the same pool is allowed; re-attaching a different pool after
buffers already exist would orphan the existing (already-pooled)
buffers, so it is rejected — the buffers' addresses are baked into
captured graphs and must not move.

Drop the cached ``_staging`` buffers so the next collective
re-allocates them into the (re-)attached staging pool.

Called on a hot-swap reload: the engine recreates ``named_pools``
(a fresh ``nccl_staging_pool``) and the old staging tensors point at
cuMem pages the old pool's teardown freed. The reload tears down the
old captured graphs first, so no graph references these addresses any
more — clearing them here is safe, and lets ``attach_staging_pool``
re-point to the new pool without tripping the orphan guard.

Prefix view of the grow-once ``all_reduce`` staging arena.

Returns a CONTIGUOUS ``numel``-element view whose base address is
constant for the process life, so NCCL's address-keyed protocol /
chunking choice is pinned exactly as the per-shape buffers pinned it —
and pinned ACROSS shapes too, which the per-numel dict never achieved.

The arena only ever GROWS, and :meth:`reserve_reduce_arena` sizes it at
boot from the config's widest eager reduce (``max_batched_tokens x
hidden``) so a growth can never happen on the serving path. A growth
past :data:`_REDUCE_ARENA_WARN_BYTES` is logged LOUD (once per dtype /
device): it moves the base address, which is only safe because the
reduce staging buffer is never referenced by a captured graph — every
address-pinned context returns from :meth:`all_reduce` BEFORE the
staging call (the ``torch.compiler.is_compiling()`` tier reduces in
place on the activation, and the raw-capture tier reduces in place on
the capture stream). :meth:`all_reduce_overlapped` likewise never
reaches staging: it either forks in place under capture or delegates
to :meth:`all_reduce` synchronously, so there is never more than one
in-flight staged reduce to alias against this single arena.

Pre-allocate the reduce arena for ``max_numel`` elements.

Called at BOOT (before the deferred KV resize measures free VRAM) for
the same two reasons :meth:`reserve_gather_arena` is: the serving path
must never allocate inside a collective region, and the bytes must be
resident when the KV sizer reads ``mem_get_info`` so the reservation is
accounted by construction rather than silently eating the post-boot
serving headroom.

Returns the bytes reserved (0 when the arena is disabled or the group
is trivial).

Build this group's copy-engine transport. BOOT ONLY.

Reserved at boot for the same two reasons every other staging strand
here is: the serving path must allocate and page-lock NOTHING inside a
collective region, and the bytes must be resident before the KV sizer
reads ``mem_get_info``, so they are accounted rather than quietly eaten
out of post-boot serving headroom.

``cfg`` is the :class:`~arbi_serve.config.ServerConfig` this boot
resolved. It is required, and the mode is read off it rather than out
of the environment, because the value that governs the tier is the one
the config RESOLVED — an env read would miss a CLI flag, a recipe, and
every live override, and would answer a question nobody asked.

``allow_on_p2p`` is the operator override for the topology refusal
below. It exists because an A/B of this transport against NCCL on a
P2P box is a legitimate thing to want; it is opt-in, never a default,
and the arm says loudly that it took it.

Returns the device bytes reserved — 0 when the tier is not armed, which
is every configuration except a two-rank CUDA group with ``min_rows``
set under a prefill capture mode this transport can serve, on a link
that has no peer access.

Can EVERY ordered pair in this group read every other's memory?

The boot-time topology fact, resolved ONCE and agreed across ranks.
``True`` peer-capable, ``False`` not, ``None`` the group could not
agree — and ``None`` is not ``False``: "we disagree about what the
device indices mean" is a different state from "this link has no peer
access", and a caller that conflates them arms a transport on a group
whose topology was never established.

WHY A FACT AND NOT A POLICY. Two transports now read this and they want
OPPOSITE answers from it: the copy-engine tier is a no-peer-access
answer and refuses when this is ``True``; async-TP's rendezvous
dereferences peer pointers and refuses when it is ``False``. Resolving
it in either one of them and letting the other re-derive it is how the
two come to disagree — so it is resolved here, once, and both read the
same value out of the boot configuration.

RANK-SYMMETRY, which is the whole difficulty, and it is a CORRECTNESS
property rather than a nicety. Both ranks must reach the SAME verdict
or the pair deadlocks with one in one transport and one in another.
:meth:`copy_engine_takes`'s docstring states the rule this obeys:
nothing rank-local may decide a transport, because the failure is a
HANG, not a wrong answer. Two collectives buy the agreement:

  1. all_gather the device index each rank holds, so every rank has the
     same ordered device list and evaluates the same ordered pairs.
  2. all_gather the per-rank verdict, and require UNANIMITY. Ranks that
     disagree do not get a majority vote — they get ``None``, because a
     disagreement means the indices did not mean the same thing in both
     processes (a per-rank ``CUDA_VISIBLE_DEVICES`` is how), and an
     unresolvable topology is exactly the case no transport may assume.

Never called per-collective. A driver query on the hot path is a cost
on every reduce AND a value that could in principle answer differently
on the two ranks — which is the hang again, arrived at from the other
direction.

The resolved topology fact, or ``None`` if it was never resolved.

Read by every transport that routes on topology. Deliberately does NOT
resolve on demand: a resolution is two collectives, so a lazy one on a
path only some ranks reach is a hang. Boot resolves it; this reads it.

Is this link one the copy engines are the right answer for?

The copy-engine tier is a NO-PEER-ACCESS answer. It never touches the
GPU-to-GPU link at all: both partials go out over the host, so its
floor is the same on every pair regardless of what connects them. That
is a win on a link the GPUs cannot cross directly and a loss on one
they can, and the tier had no way to tell those apart — the predicate
that routes each reduce (:meth:`copy_engine_takes`) may read only the
tensor and the boot configuration, so the topology has to be resolved
HERE, once, and left in the boot configuration for it. Peer capability
is also not a driver call that belongs on a per-reduce path: it would
be a query on the hot path whose answer, if it ever differed between
the ranks, turns a wrong transport into a hang rather than a wrong
answer.

RANK-SYMMETRY, which is the whole difficulty. Both ranks must reach the
SAME verdict or the pair deadlocks with one in the host exchange and
one in NCCL. Two collectives buy that:

  1. all_gather the device index each rank holds, so every rank has the
     same ordered device list and evaluates the same ordered pairs.
  2. all_gather the per-rank verdict, and require UNANIMITY. Ranks that
     disagree do not get a majority vote — they get a refusal, because
     a disagreement means the indices did not mean the same thing in
     both processes (a per-rank ``CUDA_VISIBLE_DEVICES`` is how), and
     an unresolvable topology is exactly the case that must not arm.

The refusal direction is the safe one throughout: not arming leaves
every reduce on NCCL, which is correct on every link. Arming wrongly is
the one that costs, so any doubt refuses.

May this boot's prefill capture mode carry a copy-engine reduce?

Only ``"eager"`` may. Everything else — including a mode that cannot be
read at all — is refused LOUD and the tier stays unarmed, because the
alternative is a transport that replays stale bytes and reports
success (see the tier comment above).

The mode is asked of the config object that owns it, not of the
environment: ``ServerConfig`` is where the CLI flag, the recipe, and
any live override have already been folded into one resolved value,
and it is the only place that value is validated against
``PREFILL_CAPTURE_MODES``. A caller that cannot produce that object
cannot prove the tier is safe, so it does not get it.

Does this reduce go to the copy engines? RANK-SYMMETRIC.

Every term is a property of the TENSOR and of the boot configuration,
both of which are identical on both ranks by the time a reduce is
issued — an SPMD pair reduces the same shape and dtype in the same
order, and a pair that did not would already be deadlocked in NCCL.
Nothing rank-local (free memory, a timing, a local queue depth) may
ever enter this predicate: the two ranks would disagree about which
transport is in use, and the failure is a hang, not a wrong answer.

The NCCL fallback, in place. Called by the copy-engine ops.

It lives here because a collective scoped to a named parallelism group
may only be issued from this module — the op that needs it is generic
over groups and holds no process group of its own.

Begin a copy-engine reduce of ``x``, or do the whole NCCL one.

Pair every call with :meth:`all_reduce_wait` on the same tensor. The
caller should enqueue independent compute between the two — that is
the only thing that makes this transport worth having.

Prefix view of the grow-once gather ``src`` / ``dst`` arena.

Returns a CONTIGUOUS ``numel``-element view whose base address is
constant for the process life, so NCCL's address-keyed protocol /
chunking choice is pinned exactly as the per-shape buffers pinned it —
and pinned across shapes too, which the per-shape dict never achieved.

The arena only ever GROWS, and :meth:`reserve_gather_arena` sizes it at
boot from the config's maximum so a growth can never happen on the
serving path for the one strand that matters. A growth past
:data:`_GATHER_ARENA_WARN_BYTES` is logged LOUD (once per dtype /
direction): it moves the base address, which is only safe because the
gather staging buffers are never referenced by a captured graph (every
address-pinned capture context takes the ``pinned_capture`` branch in
:meth:`all_gather`, which allocates inside the graph pool instead).

The threshold is deliberate: small metadata gathers (a handful of fp32
scalars — plan headers, per-rank counters) legitimately walk a few
element counts and their arena tops out in the tens of BYTES. Warning
on those would bury the one message that matters, which is a
VOCAB-scale strand allocating on the hot path.

Pre-allocate the gather arenas for ``max_src_numel`` source elements.

Called at BOOT (before the deferred KV resize measures free VRAM) so
that (a) the serving path never allocates inside a collective region
and (b) the bytes are already resident when the KV sizer reads
``mem_get_info`` — the reservation is accounted by construction instead
of silently eating the post-boot serving headroom.

Returns the total bytes reserved (0 when the arena is disabled or the
group is trivial).

Boot-stable staging buffer (see ``__init__``): one exactly-sized
buffer per (dtype, device, numel), allocated on first use and never
reallocated, so NCCL and captured graphs always see the same base
address for a given collective shape.

The first allocation for each shape is routed through the attached
cuMem-backed :class:`NamedMemPool` (:meth:`attach_staging_pool`) so
the persistent collective buffers are a NAMED, accounted pool — not
an unaccounted strand on the default heap. Routing only governs
WHICH heap segment the one-time ``torch.empty`` lands in; the buffer
is still allocated once per shape and held for the process life, so
its address is unchanged across collectives + cudagraph replays.

Context: pin collectives to ``stream`` for cudagraph capture.

Stacked entry is forbidden — re-entering raises so
accidentally-nested capture regions fail loudly rather than
silently picking up the inner stream.

Caller (typically
:func:`arbi_serve.distributed.graph_capture.graph_capture`) is
responsible for warming any per-stream NCCL state before
entering this context, so capture itself sees no
``cudaMalloc``.

Context: make :meth:`all_reduce_overlapped` fork onto ``stream``.

``None`` is a legal (inert) binding — the overlap degrades to the
ordinary synchronous :meth:`all_reduce` for the duration. Stacked
entry is forbidden for the same reason :meth:`graph_capture` forbids
it: a nested region would silently inherit the inner stream.

Sum ``x`` across the group with the collective in flight inside
the ``with`` body.

The body must compute something independent of ``x``; the yielded
tensor is only valid to read after the body exits (the join is the
``__exit__``). Small collectives on a P2P-less multi-GPU box are
latency-bound, not bandwidth-bound: a per-layer decode all-reduce
costs its wire latency no matter how it is tuned, so the only
lever is running other work underneath it.

Address-pinning is the gate. The fork reduces in place on the
activation, which is only legitimate where :meth:`all_reduce`'s
address-pinned tiers already reduce in place — inside a raw cudagraph
capture (``_capture_stream`` set), where every replay pins the same
device address. Outside that (live eager, allocator-fresh addresses;
or a Dynamo trace, where ``torch.cuda.stream`` is untraceable) this
delegates to :meth:`all_reduce` and simply runs the body afterwards
— same values, no overlap. So the bit-stability contract
:meth:`all_reduce` documents holds on every path.

Capture safety: the side stream forks off the capturing stream via
``wait_stream`` (an event record + wait) and rejoins it before the
body's caller reads the result, which is exactly the fork/join
pattern CUDA stream capture records as a graph branch. NCCL's
per-(comm, stream) buffers must already be warm on the side stream —
:func:`arbi_serve.distributed.graph_capture.graph_capture` warms them
before it binds it.

All-gather into a ``(size, *x.shape)`` VIEW — no concatenation.

The gather itself, without the ``torch.cat`` :meth:`all_gather` adds
to lay the ranks out along one axis. A caller that reads only part of
each rank's contribution — the drafter's vocabulary window, whose
per-rank widths are unequal and gather at their maximum
(:class:`~arbi_serve.weight_quant.exl3.linear
.EXL3OutputPrefixVocabParallelLinear`) — composes its own narrower
result in ONE copy instead of paying for the full-width one first.

The return is a view of the staging buffer, valid until the next
gather on this group reuses it: read it, or copy out of it, before
the next collective.

Capture-aware, mirroring :meth:`all_reduce`. In an address-pinned
capture context (compiled trace, or raw cudagraph capture with
``_capture_stream`` set — e.g. the eager MTP drafter's per-step
``lm_head`` pair gather) the gather reads straight from the
activation ``x``, no ``src`` staging copy; the graph pins ``x``'s
address across replays. Only a live-eager (non-captured) gather
stages ``src`` through the boot-stable buffer.

CPU shared-memory control-plane queue for the TP hot path.

Replaces the per-step pickle-then-NCCL-broadcast control plane (rank 0
→ rank 1) with a lock-free shared-memory ring buffer. NCCL is reserved
for tensor collectives only.

Topology: one writer (rank 0), one local reader (rank 1), zero remote.
Under load the ring is lock-free and syscall-free (busy-spin); only
when idle does the reader fall back to a zmq poll. Payloads larger than
``max_chunk_bytes`` ship over a zmq socket (the overflow path).

Reader busy-spin window from the env, falling back to the default.

A malformed or negative value falls back LOUD-ly to the default rather
than silently disabling the spin (kill-footguns: no silent fallback).

Reader wake-up that busy-spins under load, blocks when idle.

Pure busy-spin burns a core; a blocking wait adds a syscall per
message. This does both: for ``busy_loop_s`` after the last read,
``wait`` only yields the scheduler (hot path, no syscalls); after
that it polls a zmq socket for a notify / cancel / timeout.

A single shared-memory block split into data + metadata regions.

``max_chunks`` data slots of ``max_chunk_bytes`` each, followed by
``max_chunks`` metadata records of ``1 + n_reader`` bytes. Per-chunk
metadata byte 0 is the written flag; bytes ``1..n_reader`` are
per-reader read flags::

    0???...???  not written; writer may write, no read
    1000...000  just written; readable, not writable
    1???...???  written, read by some; readable by those who haven't
    1111...111  read by all; writer may overwrite

Picklable via ``handle()`` — the receiver attaches by shm name.

Single-writer / N-local-reader shared-memory broadcast queue.

Rank 0 writes (``enqueue``); every worker rank reads the same slot
(``dequeue``). FIFO order is guaranteed by the single sequential
``current_idx`` cursor on each side over the same ring; the writer
paces to the slowest reader (it cannot overwrite a slot until all
readers have acked it via their per-reader flag byte).

All readers are local (single-node TP), so ``n_reader`` must equal
``n_local_reader``. The per-slot metadata reserves one read-flag byte
per reader (``metadata[1 .. n_reader]``), so the protocol is N-general
— TP=2 is just the ``n_reader == 1`` case.

Yield within the busy window, else poll for a notify/cancel.

A context ``term()`` from another thread (clean shutdown) makes
the poll return ETERM; we swallow it so the caller re-checks the
``shutting_down`` flag and exits.

Collective: rendezvous writer and reader before any message.

Defeats the zmq slow-joiner — the writer recvs a subscription
from the reader, then sends ``b"READY"``; the reader blocks on
it. Both sides must call this.

Wake an idle reader so it can exit (raises in ``dequeue``).

Sends the cancel ping (and, for a writer, also closes its sockets
+ context); a reader only signals, leaving its own sockets for
the dequeuing thread to unwind via :meth:`close`.

Close every zmq socket (LINGER=0) and terminate the context.

Must be called only once no thread is polling this queue's
sockets (e.g. after the reader's dequeue loop has exited). Closing
with LINGER=0 + terminating keeps context teardown from blocking
at GC time. Idempotent.

SPMD tensor-parallel derivation engine.

True SPMD TP replaces the current driver/worker per-step batch handoff
(:mod:`arbi_serve.engine.distributed_driver` +
:mod:`arbi_serve.distributed.worker_bridge`) with a model where every TP
rank derives the identical per-step batch from once-broadcast scheduler
state, instead of receiving a materialized batch every forward.

This package is the derivation core, deliberately decoupled from the
engine / torch.distributed plumbing so the control flow is CPU-testable
under a 2-rank gloo harness:

  - :class:`SlateDelta` — the compact, picklable scheduler decision rank
    0 broadcasts once per tick (admits / slate / evicts / preempts /
    spec mode). The only large field is a new request's prompt token ids,
    which ride once on admission — not per step.
  - :class:`RankSlateMirror` — the per-rank request-state mirror every
    rank advances from the delta (prompt ids, prompt_consumed, committed
    output tokens, sampling digest). Fed the identical delta sequence,
    every rank's mirror is bit-identical to rank 0's.
  - :func:`apply_slate_delta` — advance a mirror by one tick's delta.
  - :func:`derive_step_tensors` — the rank-local twin of ``_gather_slate``'s
    per-step batch walk (:mod:`arbi_serve.runtime._batch_build_materialize`).
    Run against a rank-local page table fed the identical alloc/free order,
    it produces ``input_ids`` / ``positions`` / ``slot_mapping`` /
    ``seq_lens`` / ``cu_seqlens_q`` / ``cu_seqlens_k`` / ``block_table``
    tensor-for-tensor identical to rank 0's — with no per-step plan
    broadcast. Row shape (prefill / inject / decode) is resolved through
    :func:`arbi_serve.engine.row_shape.resolve_row_shape`, the same pure,
    torch-free function ``_gather_slate`` calls — sharing it costs nothing
    against the "CPU-testable, decoupled from engine/torch.distributed
    plumbing" goal above, since it has no dependency on either.

The engine wiring (the SPMD run loop, gated on ``ARBI_SPMD_TP``) is
separate; this module + its tests prove the derivation contract.

The derivation core is split across sibling modules by concern — the
broadcast wire types (:mod:`.spmd_delta`), the per-rank mirror +
delta/commit advancers (:mod:`.spmd_mirror`), the rank-symmetric sampling
views (:mod:`.spmd_sampling`), the scheduling-tensor derivation
(:mod:`.spmd_derive`), and the rank-0 delta/AdmitRow builders
(:mod:`.spmd_builder`). Every public name is re-exported here, so
``arbi_serve.distributed.spmd`` remains the single import surface. This
module hosts the per-rank :class:`SpmdRankLoop` driver that binds them.

Per-rank holder binding the slate mirror to a rank-local page table.

Every TP rank (including rank 0) owns one. It applies each
once-per-tick :class:`SlateDelta`, derives the per-step batch tensors
against its rank-local page table, and advances its mirror by the
step result — so two ranks fed the identical delta + result sequence
produce byte-identical tensors with no per-step batch broadcast.

This is the derivation core the flag-gated ``_run_loop_spmd`` driver
wraps; kept engine-free so the 2-rank gloo harness exercises the full
contract on CPU. The page table is supplied by the caller
(:class:`arbi_serve.cache.pagetable.FlatPageTable` today; radix in a
later phase) and must be fed the identical admit/free order on every
rank — :meth:`apply_delta` registers admits + frees evicts in that
order.

Advance the mirror + page table by one tick's delta.

Order matches rank 0's scheduler exactly: free evicted requests'
pages first, then register admits — so the rank-local pool
free-list advances in lockstep with rank 0's. A delta may evict a
request the rank-local page table never registered (e.g. it
finished mid-prefill before any decode page was allocated); the
mirror drop is authoritative, so the page-table free is
best-effort (KeyError suppressed).

When ``drive_page_table`` is False (rank 0, sharing its table with
the scheduler) only the mirror advances; the scheduler already
applied this tick's admits/evicts to the shared table.

The recurrent-state pool (``state_pool``) is driven in the same
order and under the same ``drive_page_table`` gate as the page
table: a worker frees evicted/preempted rows first, then allocs
admitted rows — mirroring rank 0's ``Scheduler.add`` (alloc) and
finish (free) call order. Fed that identical order, the pool's
LIFO free-list assigns the same slab row to the same request on
both ranks, so ``RecurrentStatePool.row_for`` resolves identically
for plain decode and the MTP verify pass. ``free_recurrent_state``
/ ``alloc_recurrent_state`` are idempotent (unknown id is a no-op
free), so a row the pool never registered frees cleanly.

Refuse a credited prefix this rank cannot restore the state for.

The worker's half of the hybrid invariant rank 0 asserts at both of
its admission sites: a credited match skips prefill for the
recurrent layers as well as the attention ones, so it is sound only
where the recurrent state at that boundary is restored — and this
rank has no savepoint store to restore it from. Raising is what
makes the failure loud instead of an answer computed from a state
that existed on one rank only.

Promote a finishing request's completed full pages into the tree.

Sources the committed token sequence from the mirror — the one
state that is byte-identical on every rank at every point (rank 0's
real ``Request`` objects are not: the verify stream stops appending
at a mid-accept finish, and a failed step never commits its tokens
to them at all, while the mirror advanced on every rank). This is
therefore the single finish-commit implementation for all ranks:

  - the worker calls it from :meth:`apply_delta`'s evict replay;
  - rank 0's SPMD driver calls it from its deferred-release drain
    (``Scheduler.set_spmd_deferred_release`` queues both rank-0
    finish paths — natural finish and cancel/timeout ``remove`` —
    and the drain runs at the top of the next tick, the byte-exact
    point the worker's replay runs). Committing from the Request's
    own (possibly shorter) token list, skipping the commit on a
    cancel, or committing/freeing mid-tick all desync the per-rank
    trees or free-lists — a later broadcast ``evict_spmd`` victim
    or ``add_request_spmd`` match then fails to resolve on the
    worker.

No-op on the flat table or when the mirror never admitted ``rid``
(then no rank ever commits it — still symmetric).

Replay rank 0's preempt-time commit of the victim's full pages.

The worker twin of
:meth:`arbi_serve.scheduler.preemption.Preemptor._commit_written_pages`,
and it must run for the same reason the finish-time commit does:
rank 0 promotes the victim's decode-completed pages into its tree
just before the preempt drops the request's refs, so a worker that
skipped it would hold those pages private where rank 0 holds them
shared — a structural tree divergence, and the broadcast
``prefix_page_ids`` on this very row would then fail to resolve.

Truncated to the page table's own ``length`` for the same reason
rank 0 truncates: the last sampled token is in the mirrored
sequence but has not been fed, so its slot holds no KV.

TEST-ONLY: derive this tick's MTP-verify scheduling tensors.

Thin wrapper over :func:`derive_verify_tensors` — see that
function's docstring. NOT on the live SPMD verify path (the real
path runs ``build_verify_plan`` through the ``_WorkerVerifyRequest``
shim); only the GLOO derivation harnesses and
``tests/test_spmd_derivation.py`` call this method.

Promote completed prompt pages after a successful prefill step.

Mirrors :meth:`Scheduler.commit_state`'s ``commit_full_pages(rid)``
on the prompt-completing chunk — the cross-batch sharing seed, so a
later request's broadcast match resolves the same nodes on the
worker. Fires only for rows that left prefill on this commit. No-op
on the flat table / rank 0 (its scheduler drives the real commit).

Record a request's head-driven drafts for its next verify step.

Mirrors the engine's ``req.mtp_next_drafts = ...`` write after the
seed-drafter chain. Under SPMD every rank calls this with its own
(identical) drafts.

Rank-0 SlateDelta / AdmitRow builders (the broadcast source).

:class:`RankZeroSlateTracker` turns a real scheduler slate into the
minimal non-derivable :class:`SlateDelta`; :func:`admit_row_from_request`
/ :func:`sampling_digest_from_params` build the once-per-request
:class:`AdmitRow` + :class:`SamplingDigest` from a live engine
:class:`Request`. Re-exported from :mod:`arbi_serve.distributed.spmd`.

Rank-0-only bookkeeping that turns a real scheduler slate into a
:class:`SlateDelta`.

The scheduler owns admission / eviction; this tracker observes its
decisions across ticks and emits the minimal non-derivable delta:

  - ``admit``: request ids appearing in a slate for the first time
    (their prompt ids + sampling digest + prefix-match ride once).
  - ``evict``: request ids the tracker had seen running but that are
    gone from this tick's slate and reported finished/removed by the
    caller (pages return to the rank-local pool in the same order).
  - ``slate`` / ``spec_mode`` / ``step_K``: this tick's run set.

The tracker holds no tensors and no page table — it is the thin
rank-0 shim between the scheduler and the broadcast wire.

Build an :class:`AdmitRow` from a
:class:`arbi_serve.engine.request.Request`.

The radix-cache match rank 0's scheduler applied rides the wire so the
worker applies it (seeds ``prompt_consumed`` + aliases the same shared
pages) rather than re-running the LRU match. ``prefix_match_len`` comes
from ``req.cached_prefix_tokens`` (the scheduler pinned it at admission)
and the matched page ids from ``page_table.shared_page_ids(req_id)``;
both are 0 / empty on the flat table (no prefix sharing) and on a cache
miss.

Structured output / tool-call grammar, per-request ``logit_bias`` and the
LoRA selection are supported rank-symmetrically: the grammar spec, the bias
map and the ``lora_id`` each ride the ``AdmitRow`` once, and every rank
reproduces them locally — the matcher via
:meth:`xgrammar_processor.XGrammarLogitsProcessor.build_grammar_state`
(advanced per committed token), the bias via the sampler's order-independent
scatter-add, the LoRA via each rank's TP-sharded adapter shard (only the
name crosses; weights load per-rank from disk). Still refuses loud on the
one sub-case the delta genuinely cannot represent: multimodal / vision input
(no image features / M-RoPE plan), since a worker cannot reproduce the
vision forward. Rank 0 raises here before the delta is built, so such a
request never enters the SPMD slate.

Build the :class:`SlateDelta` for one scheduler tick.

``slate`` is the scheduler emit reduced to ``(request_id,
n_tokens)``. ``admit_info`` maps a request id to its
:class:`AdmitRow` — the caller (rank 0, which holds the request
map) supplies it for every id the tracker has not seen yet.
``finished`` is the list of request ids the scheduler retired
this tick (their pages free). ``preempt`` is this tick's
:class:`PreemptRow` decisions; ``evict_pages`` is the radix LRU
victim plan rank 0 reclaimed up front. A preempted request stays
known (it re-prefills on a later slate), so it is not discarded
from ``_known`` — only ``finished`` rows are.

SPMD broadcast wire types — the once-per-tick scheduler decision.

The compact, picklable payload rank 0 broadcasts to every TP rank
(:class:`SlateDelta` and its rows) plus the plain-int :class:`SpecMode`
mirror, kept engine-import-free so the pure-CPU gloo derivation harness
carries no engine dependency. Re-exported from
:mod:`arbi_serve.distributed.spmd`.

The per-request sampling fields a worker rank needs to reproduce
rank 0's sampling decision.

Tokenization + admission stay rank-0-only; this digest carries only
what the on-device sampler reads so every rank runs the identical
argmax (greedy) / Gumbel (stochastic) draw. The published sampler
(temperature / top_k / top_p / min_p), the three penalties, and the
per-request ``seed`` are all carried so the stochastic draw is
rank-symmetric (the Gumbel noise is derived from ``(seed, position)``;
penalties are a deterministic function of the replicated token
history). ``logit_bias`` and the grammar spec are constant per request,
so they ride the :class:`AdmitRow` once (not in this per-step digest) —
see :func:`admit_row_from_request`.

One newly admitted request, broadcast once when it first appears.

``prompt_token_ids`` is the only large field and rides exactly once
per request, never per step. ``prefix_match_len`` is the prefix-cache
match rank 0's scheduler applied — the worker applies it (seeds
``prompt_consumed``) rather than re-running the radix match, so the
rank-local page table stays in lockstep without owning the matching
policy. ``prefix_page_ids`` is the matched cache pages rank 0's radix
lookup returned; the worker aliases exactly these (same physical page
ids, since the pool free-list is mirrored) instead of re-walking its
own tree — the radix-match decision is the only non-derivable bit, so
it rides the wire. Empty on the flat table / a cache miss.

Caller-supplied context tokens fed into a running request this tick.

The tool-call resume path: ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>``
appended to an already-decoding sequence WITHOUT being sampled (see
``Request.pending_context_token_ids`` and ``_gather_slate``'s
injection row). The ids are arbitrary caller data — no rank can
derive them — so they ride the wire once, on the tick whose slate
serves the wide row.

One running request rank 0's scheduler preempted back to waiting.

Preemption (``Preemptor.preempt_for_space``) frees the victim's pages
and rebuilds it as a request whose prompt is its whole prompt plus
what it has generated, re-prefilling from whatever the radix cache
still holds of that sequence. The worker reproduces the
same page free + re-add + state reset off this row, so its pool
free-list + recurrent pool stay in lockstep — the alloc that re-grows
the request rides the next tick's per-step ``allocate_slots`` (mirrored
like any decode). ``new_prompt_len`` pins how many of the mirrored
tokens become the restarted prompt.

``prefix_match_len`` / ``prefix_page_ids`` are the radix re-match rank
0's admission pass applied when it picked the victim back up, carried
for exactly the reason :class:`AdmitRow`'s are: the tree is in lockstep
but the MATCH is not derivable — rank 0 caps it by the page headroom
its step has left and by its recurrent-savepoint coverage, and a worker
has neither number. The worker aliases exactly these pages and seeds
``prompt_consumed`` from this length. Empty / 0 on the flat table, on a
genuine miss, and whenever the victim was not re-admitted on the pass
its own preempt made room for.

``cache_enabled`` is an OVERRIDE, and ``None`` — keep whatever the
entry being dropped carried — is the default on purpose: the re-admit
seam reads the victim's namespace and cache-write flag off its own
table entry precisely so no caller can forget them (#1818). Rank 0
sets it only where its resolution DIFFERS from that carried value:
the savepoint admission force-disables the cache when a hybrid model
matched pages with no covering savepoint, and a worker that missed
that promotes pages rank 0 did not — the same per-rank free-list split
:class:`AdmitRow` carries its field to avoid.

The compact per-tick scheduler decision broadcast to every rank.

Steady-state decode ticks carry empty ``admit`` / ``evict`` /
``preempt`` so the delta is just ``slate`` (B small int pairs) plus
two scalars — tens of bytes. New-request ticks additionally carry the
admitted prompt ids (once).

One tick's scheduler decision for EVERY attention-DP set.

At ``attn_dp_size > 1`` each set forwards its own rows, so rank 0
broadcasts one of these instead of a bare :class:`SlateDelta` and each
rank reads ``per_dp[attn_dp_rank]``. Sending all the sets' deltas in a
single message — rather than one message per set — is what makes the
tick's control flow a function of the message alone: every rank sees
the same object and therefore reaches the same
:attr:`forwards` verdict.

THE FORWARD DECISION IS GLOBAL, NOT LOCAL. The FFN exit reduces over
the whole TP group, so a rank that skipped the forward because its own
set had no row would leave every other rank blocked in that reduce.
:attr:`forwards` is the one predicate every rank gates on, and a set
with an empty slate still runs the forward, contributing zero rows to
the attention-DP gather.

Per-step scheduling-tensor derivation from the rank-local mirror.

:func:`derive_step_tensors` is the rank-local twin of ``_gather_slate``'s
scheduling-tensor build (:mod:`arbi_serve.runtime._batch_build_materialize`,
the real per-step batch walk) and IS on the live SPMD path (reached via
:meth:`arbi_serve.distributed.spmd.SpmdRankLoop.derive`).

:func:`derive_verify_tensors` is TEST-ONLY. It is DEAD on the live SPMD
verify path and must never gain a caller under ``arbi_serve/`` — the
live path runs ``build_verify_plan`` itself through the
``_WorkerVerifyRequest`` shim in :mod:`arbi_serve.spec_decode.mtp_verify_spmd`
(see that module's docstring), NOT this function. Confirmed via
``grep -rn "\.derive_verify(" arbi_serve/ tests/``: its only
callers are ``tests/test_spmd_derivation.py`` and the 2-rank GLOO
derivation harnesses (``tests/_spmd_gloo_harness.py``,
``tests/_spmd_mixed_slate_gloo_harness.py``). Kept rather than deleted
because those harnesses need a fast, model-free, CPU-only way to
regression-test both the tensor arithmetic AND its 2-PROCESS GLOO
parity (independent processes deriving byte-identical tensors from a
tiny broadcast delta); repointing them at the real ``build_verify_plan``
path would require a live model forward under GLOO, which defeats the
harnesses' whole reason to exist.
``tests/test_spmd_derivation.py::test_derive_verify_tensors_has_no_production_caller``
greps ``arbi_serve/`` for stray calls so this stays true.

Run against a rank-local page table fed the identical alloc/free order,
both functions produce tensors bit-identical to rank 0's with no per-step
plan broadcast. Re-exported from :mod:`arbi_serve.distributed.spmd`.

``derive_step_tensors`` decides each row's shape (prefill / inject /
decode) and start position through :func:`arbi_serve.engine.row_shape.
resolve_row_shape` — the SAME function ``_gather_slate`` calls — rather
than re-deriving that arithmetic independently. See that module's
docstring for why: two hand-written copies of that arithmetic let one
silently fall behind the other.

The per-step scheduling tensors every rank derives identically.

All numpy (host) — the engine wiring moves them to the rank-local
device with the same dtypes :class:`ScheduledBatch` uses
(int32 for ids/positions/seq_lens/cu_*/block_table, int64 for
slot_mapping). Keeping them host-side here makes the derivation
CPU-testable and the cross-rank parity assertion exact.

Derive one step's scheduling tensors from the rank-local mirror.

The rank-local twin of ``_gather_slate``'s per-step batch walk
(:mod:`arbi_serve.runtime._batch_build_materialize`), reduced to the
fields the SPMD parity contract covers. ``page_table`` is the
rank-local :class:`arbi_serve.cache.pagetable.FlatPageTable` (or
radix, later) — fed the identical admit/alloc/free sequence as rank
0, ``allocate_slots`` / ``block_table_row`` / ``length`` /
``page_ids`` return identical values, so the derived tensors are
bit-identical across ranks without receiving a materialized batch.

Row shape (prefill / inject / decode) and each row's start position
are resolved through :func:`arbi_serve.engine.row_shape.
resolve_row_shape` — the single function What stays
local to this derivation:
  - prefill row: input_ids = ``prompt_token_ids[start_pos:+n]``.
  - inject / decode row: input_ids = ``[last_token_id, *pending_ctx]``
    (numpy source, no GPU D2D overwrite — that optimization is
    rank-0-only, see ``_gather_slate``).
  - ``allocate_slots(req_id, n)`` extends the page table and returns
    the per-token slot indices (``page_id*bs + off``).
  - ``seq_lens[i] = page_table.length(req_id)`` after the extend.
  - ``cu_seqlens_q`` = prefix-sum of per-row ``n``.
  - ``cu_seqlens_k`` = ``[0, cumsum(seq_lens)]`` int32.
  - ``block_table`` row = ``page_table.block_table_row(req_id,
    max_pages)`` padded to the batch-wide ``max_pages``.

The per-step MTP-verify scheduling tensors every rank derives.

Superset of :class:`DerivedStepTensors` with the verify-pass meta the
forward + accept path reads. Mirrors the
:func:`arbi_serve.spec_decode.mtp.build_verify_plan` flat-batch
assembly: per row ``i`` contributes ``K_i + 1`` flat tokens at slots
``[tail_slot_i, draft_slots_i...]`` and positions
``[pre_len_i - 1, pre_len_i, ...]``. Under SPMD the draft tokens come
from the rank-local mirror (every rank ran the same seed-drafter
chain) and the draft slots from the rank-local pool (same alloc
order), so the whole tensor set is derived — not broadcast.

TEST-ONLY. Derive one MTP-verify step's tensors from the rank-local mirror.

NOT on the live SPMD verify path — ``DistributedEngineDriver._spmd_run_verify``
calls :func:`arbi_serve.spec_decode.mtp.run_verify_step_spmd_worker`, which
runs the REAL :func:`build_verify_plan` directly through the
:class:`arbi_serve.spec_decode.mtp_verify_spmd._WorkerVerifyRequest` shim
(a ``Request``-shaped adapter over a mirror row) rather than
re-deriving the flat batch here. That shim IS this function's
unification story — genuine code reuse, not a second implementation
to keep in
lockstep by hand. This clone is exercised only by the CPU-only 2-rank
GLOO derivation harness (``tests/_spmd_gloo_harness.py`` and friends),
a fast, no-model-forward way to
regression-test the arithmetic in isolation. Kept in sync with
``build_verify_plan`` by inspection, not by construction — unlike
:func:`derive_step_tensors` (see its docstring), which shares
:func:`arbi_serve.engine.row_shape.resolve_row_shape` with its rank-0
twin. A change to this function must be double-checked against
``build_verify_plan`` / ``_assemble_verify_flat_batch``
(``spec_decode/mtp_verify_plan.py``) by hand.

The rank-local clone of
:func:`arbi_serve.spec_decode.mtp.build_verify_plan` reduced to the
scheduling-tensor + MTP-meta fields the SPMD parity contract covers.
``page_table`` is the rank-local table and exposes the draft-slot
lifecycle (``allocate_draft_slots`` / ``finalize_draft_slots``) the
flat + radix tables both implement.

Lockstep-identical to ``build_verify_plan``:
  - uniform-K invariant (the scheduler bucketed at admission); the
    per-row draft cache must hold ``>= step_k`` drafts or the whole
    step collapses to ``eff_k == 0`` (one tail token per row).
  - slot allocation order: tail slot first (``allocate_slots(.,1)``)
    for every row, then draft slots (``allocate_draft_slots(.,K)``)
    for every row — the same two-pass order rank 0 uses, so the pool
    free-list advances identically on every rank.
  - flat layout per row: ``[last_committed, drafts...]`` at positions
    ``[pre_len-1, pre_len, ..., pre_len+K-1]``.

Per-rank request-state mirror + the delta/commit advancers.

:class:`RankSlateMirror` is every rank's local copy of the running
request set. :func:`apply_slate_delta` advances it by one broadcast
tick; :func:`commit_step_tokens` / :func:`commit_verify_tokens` advance
it by a step's sampled / accepted tokens. Fed the identical delta +
result sequence, two mirrors are bit-identical — the foundation of the
SPMD lockstep. Re-exported from :mod:`arbi_serve.distributed.spmd`.

A worker rank's local copy of one request's decode state.

Mirrors the subset of :class:`arbi_serve.engine.request.Request` the
batch derivation reads: prompt ids, how much prompt is consumed,
committed output tokens, and the sampling digest. Advanced purely
from the broadcast delta + the rank-local sampler — never received as
a materialized batch.

Per-rank mirror of the running request set, advanced by deltas.

Every TP rank (including rank 0) owns one. :func:`apply_slate_delta`
advances it; :func:`derive_step_tensors` reads it. Fed the identical
delta sequence, two mirrors are bit-identical — the foundation of the
SPMD lockstep.

The restarted prompt a :class:`PreemptRow` describes, off the mirror.

The whole prompt plus whatever has been generated from it — the same
rebuild rank 0's :class:`Preemptor` performs, sourced from the one
state that is byte-identical on every rank. NOT
``prompt[:prompt_consumed] + output``: a victim caught mid-prefill has
no output yet, so that expression drops its unconsumed prompt tail.
``PreemptRow.new_prompt_len`` is the cross-check: a mismatch means the
worker's mirror diverged from rank 0's pre-preempt state, which is a
lockstep break, not a recoverable condition.

Both consumers of the preempt row need this — the page table's
re-admit replay (which keys the radix walk on it) and the mirror
reset — and they run at different points of ``apply_delta``, so it is
one function rather than two copies that can drift apart.

Advance ``mirror`` by one tick's :class:`SlateDelta`.

Order matters and must match rank 0's scheduler exactly so the
rank-local page-table alloc/free order stays in lockstep:

  1. evicts — drop the request from the mirror (its pages return to the
     rank-local pool when the caller frees them).
  1b. preempts — reset the request in place (rebuild the restarted
     prompt, zero ``prompt_consumed`` / output), mirroring rank 0's
     ``Preemptor``; the request stays known and re-prefills later.
  2. admits — register new requests; seed ``prompt_consumed`` from the
     broadcast prefix match.

Output-token advancement is not done here — it happens via
:meth:`RankSlateMirror`-side commit after each step's sampler runs
(every rank samples identically), so the mirror reflects committed
state before the next delta arrives.

Rank-local reasoning-budget accounting for one committed token.

Mirrors the engine's post-token ``count_reasoning_token``
(run_step/terminal.py) so the reasoning count + force-close state are
DERIVED identically on the mirror as on rank 0's real Request — no
broadcast, exactly like grammar advancement. A cheap no-op unless the
model has a ``</think>`` marker AND this request opened in reasoning and
has not closed it yet; the ThinkingBudgetGuard reads the resulting
``force_close_*`` off this request's sampling view on the next draw.

Advance each row's mirrored state by this step's result.

Mirrors the engine's post-step commit (scheduler.commit + the
prompt_consumed advance in the scheduler): a prefill chunk advances
``prompt_consumed`` by ``n``; once the prompt is fully consumed the
sampled token is appended to ``output_token_ids``. Every rank calls
this with its own (identical) sampled tokens so the mirror is current
before the next delta. ``None`` entries are mid-prefill rows with no
output token yet.

Advance the mirror by one MTP-verify step's accepted tokens.

Mirrors the per-row commit in
:func:`arbi_serve.spec_decode.mtp.run_verify_step`: each row appends
its ``n_accepted + 1`` accepted tokens (the verify pass returns the
accepted prefix plus the bonus token). Every rank calls this with its
own (identical) accept result so the mirror's ``total_length`` —
which the next step's slot allocation depends on — stays in lockstep.

Rank-symmetric sampling read-views over the SPMD mirror.

:class:`_DigestSamplingView` exposes a :class:`SamplingParams`-shaped
read view over a :class:`SamplingDigest`; :class:`_MirrorSamplingRequest`
wraps a :class:`_MirroredRequest` in the :class:`Request`-shaped surface
the sampler + xgrammar processor read, so every rank runs the identical
draw. Re-exported from :mod:`arbi_serve.distributed.spmd`.

A :class:`SamplingParams`-shaped read view over a digest.

The one sampling surface every SPMD rank samples through — both the K=1
sampler chain (via :class:`_MirrorSamplingRequest`) and the MTP verify /
drafter chain (via ``mtp_verify_spmd._WorkerVerifyRequest``). One view over
the digest, not one per call site: a knob added for one chain is a knob the
other chain gets, so the two cannot drift into disagreeing about what a
request asked for.

Exposes the whole digest — the explicit properties below coerce or override
(``logit_bias`` comes from the :class:`AdmitRow`, constant per request and
not in the per-step digest; ``response_format`` is always None because the
grammar rides ``grammar_state`` on the mirror request, not this view), and
``__getattr__`` forwards every remaining digest field verbatim. That
forwarding is what makes "the digest carries it but the view drops it"
unrepresentable: a processor reading a knob no property lists still reads
rank 0's value rather than raising mid-step.

A :class:`Request`-shaped view the sampler can read.

Exposes ``sampling`` (the digest view), ``prompt_token_ids`` and
``output_token_ids`` as plain int lists (the penalty processor reads
their concatenation) — sourced from the replicated mirror, so the
penalty inputs are identical on every rank.

The two token lists are lazy (materialized on first read, then cached
for the life of the view). One view is built per slate row per step,
and the mirror holds ``prompt_token_ids`` as an ndarray, so an eager
``.tolist()`` here would be an O(context) host copy on every rank on
every tick. Only the penalty
(:func:`arbi_serve.sampler.processors.apply_penalties`) and
no-repeat-ngram processors read them, and only for rows that asked
for those knobs; every other row now pays nothing. The materialized
values are byte-identical to the eager form.

Carries every ``Request`` attribute the sampling chain reads, mirroring
the real Request's value/default — not just the few the penalty path
needs. The engine's sampler always has the xgrammar processor attached
(:meth:`arbi_serve.engine.build` calls
``set_logits_processors([eng.xgrammar])``), and that processor reads
``req.grammar_state`` on every row in the slate. For a constrained
(structured-output / tool-call) request the grammar spec rides the
``AdmitRow`` once and every rank builds its own identical matcher; the
matcher lives on the mirror (``_MirroredRequest.grammar_state``) and is
surfaced here so the xgrammar processor masks this row rank-symmetrically.
``None`` for unconstrained rows (the processor short-circuits them).
Listing the full set here rather than enumerating per-access-site keeps
the shim from AttributeError-ing the next time a processor reads a new
Request field.

The mirrored committed output as a plain int list.

A snapshot taken at first read, not a live alias of the mirror's
list: the view is consumed entirely inside one step's sampling
chain, which runs before the step's mirror commit, so the snapshot
is the same sequence the eager copy captured.

Cross-rank work hand-off for MTP-at-TP>1.

MTP at ``tp_size > 1`` needs cross-rank coordination. The verify-pass
forward, the drafter chain, and the rejection sampler all run inside
rank-0 code paths (:mod:`arbi_serve.spec_decode.mtp_verify_driver` /
:mod:`arbi_serve.spec_decode.mtp` / :mod:`arbi_serve.engine.run_step`)
that don't issue their own broadcasts; under TP>1 those forwards
would deadlock — workers wait at ``broadcast_object_list`` while rank
0 runs the verify forward locally and never sends the plan.

This module is the bridge that lets rank-0 code paths broadcast a
"please run this forward / drafter chain / seed update" hand-off to
worker ranks without reaching into the engine's run-loop layer. The
:class:`WorkerBridge` is set on the engine (``eng.worker_bridge``) by
the :class:`DistributedEngineDriver` constructor on every rank; rank-0
code paths consult it before running a forward / drafter and broadcast
the inputs, workers sit in :meth:`WorkerBridge.run_loop` and execute
each broadcast op against their local engine in lockstep.

Pattern 1 (rank-0-only MTP)
----------------------------

The MTP verify pass + drafter chain run on rank 0 (its scheduler / page
table / sampler / accept-reject logic is the source of truth), but the
head's ``RowParallelLinear.o_proj`` issues an ``all_reduce`` collective
that every rank in the TP group must participate in. So workers re-run
the same forward + drafter-chain calls in lockstep, discarding their
local logits/tokens. The bridge marshals just enough state across the
wire:

  - **batch broadcast** — the verify-pass / legacy forward's
    materialized :class:`ScheduledBatch`, encoded as :class:`StepPlan`
    so workers reconstruct + run forward.
  - **drafter broadcast** — ``(B, K, hidden_in, last_token_in)`` plus
    optional sampling-params for the stochastic chain so workers replay
    or live-run the same drafter chain.
  - **seed broadcast** — the engine's ``mtp_seed`` step counter so all
    ranks read identical Gumbel-max draws.
  - **tick-done** — rank 0 signals "this slate is finished, return to
    the outer recv loop"; the worker dispatch unblocks and waits for
    the next slate.

Why a bridge instead of inlining broadcasts in the strategy / runner?
Each rank-0 forward / drafter call wants to consult exactly one
boolean — "am I in TP>1 mode" — and the corresponding broadcast call.
Threading that through every call site (scheduler → strategy →
verify-driver → model_runner → mtp_driver → rejection_sampler) is much
more invasive than this single hook. The bridge is None for TP=1 (no
broadcast issued), which is the ``world_size <= 1`` short-circuit on
single-process runs, and for the rank-0 driver loop on multi-rank runs
the bridge is set to a coordinator that calls ``self.broadcast_*``;
workers' bridge is set to a coordinator that runs each broadcast op
locally. The same engine code path (the rank-0 run-loop body) ships
to every rank — only the bridge implementation differs.

Engine ↔ TP-driver hand-off for MTP-at-TP>1.

Subclasses:

  - :class:`RankZeroBridge` — broadcast each op to workers, then
    let the rank-0 path run its own local forward / drafter / seed.
  - :class:`WorkerRankBridge` — sit in a recv loop, execute each
    broadcast op against the local engine.
  - ``None`` — single-process / TP=1; engine code paths short-
    circuit the broadcast check.

The interface is sync (not async): the broadcasts are CPU-side
``broadcast_object_list`` calls that complete in microseconds. The
forward / drafter calls themselves are sync GPU work the engine
already runs synchronously inside ``inference_mode``.

Broadcast ``batch`` to worker ranks before a TP forward.

The single owner of the broadcast-before-forward obligation. Every
forward path calls this instead of re-deriving the
``world_size > 1 and rank-0`` guard, so the invariant "a forward that
isn't mirrored to the workers deadlocks under TP (rank 0 enters the
row-parallel collectives alone)" lives in one place. No-op at TP=1
(``bridge is None``) and on a worker rank (``is_rank_zero()`` False —
workers receive the batch via their dispatch loop, they don't emit).

Bridge installed on rank 0: every method broadcasts when active.

The driver constructor passes a ``broadcast_obj`` callable that
pickles a payload and sends it on the TP group; the bridge is
decoupled from torch.distributed plumbing so unit tests can pass a
fake.

Slate-gated broadcasts. ``broadcast_object_list`` is a TP-group
collective — it blocks rank 0 until every worker rank also calls
it. Workers consume broadcasts only inside :meth:`run_worker_dispatch`,
which the rank-0 driver enters via the ``"slate"`` outer-loop control
op. Boot-time forwards (``profile_activation_peak``, captured-graph
capture sweeps) and post-shutdown teardown forwards run before / after
the run loop; if those fired ``broadcast_batch`` while workers were
still in the outer recv loop, rank 0 would deadlock waiting for a
receiver that never matches.

The driver toggles :attr:`_slate_active` around each per-slate
dispatch (set in ``_run_loop_rank_zero`` after the ``"slate"`` outer
op, cleared after ``broadcast_tick_done``). When it's False, every
broadcast method is a no-op — the rank-0 engine code path runs its
local forward / drafter / seed without any cross-rank hand-off,
which is exactly what boot-time profiling needs (workers run their
own profile in lockstep via ``RowParallelLinear.all_reduce``).

Driver hook: slate finished; broadcasts disabled.

Symmetric with :meth:`begin_slate`. After this returns, any
rank-0 forward / drafter / seed call (e.g. a post-shutdown
teardown forward, or a boot-time admin sweep) skips the
bridge — the worker dispatch is not in scope to receive.

Rank 0: send a drafter-chain op to workers.

For real-attention heads (Qwen3.5/3.6), the optional
kwargs carry the rank-0 slot allocation + page-table rows so
workers reconstruct identical per-step ``attn_meta``s. For
stub heads / the legacy no-op path they stay ``None``.

``replay_captured`` tells the worker rank 0 took the captured-
graph replay so it should replay its own captured chain too.

``frontier_repair_meta`` is the device-gathered override meta on
the ASYNC optimistic-advance path (``None`` on every sync chain);
when set, the rank-0 bridge D2H's its frontier slots / positions /
seq-lens into the op so the worker runs the identical eager
``_repair_kv_frontier`` collectives in lockstep.

Broadcast the partial-accept rollback; no-op outside a slate.

Sent right before rank 0 applies its own
``pool.rollback_partial_accept`` so every rank's recurrent slab
shard drops the rejected drafts' state in the same step (see
:class:`RollbackOp`).

Broadcast the DFlash embed leg so workers join ``embed_tokens``.

``blk_ids`` is the ``(B, block)`` masked-block id tensor rank 0
is about to embed; workers run the same lookup to rendezvous on
the vocab-parallel all_reduce. No-op outside a slate.

Broadcast the DFlash lm_head leg so workers join ``lm_head``.

``hidden`` is the ``(rows, hidden_size)`` denoised hidden rank 0
is about to project; ships as contiguous CPU fp32 (shm OOB path)
so workers compute the same local logit shard and rendezvous on
the all_gather. No-op outside a slate.

Broadcast a full sharded DFlash draft step; no-op outside a slate.

Every worker rebuilds the masked block ids + its own shard context
and runs the identical ``embed → forward_block_fixed → lm_head``
train so the shard all_reduces / all_gather rendezvous 1-to-1.

Receive + execute ops until rank 0 broadcasts tick-done.

Errors inside an op are logged but do not exit the loop —
rank 0 may still broadcast more ops on its slate (the engine's
own per-step error guard catches forward failures); we have to
stay rendezvoused on the broadcast collective even after a
worker-side failure or the next slate's broadcast deadlocks.

Returns the :class:`ControlOp`s that arrived mid-slate, deferred
for the caller (:meth:`DistributedEngineDriver._run_loop_worker`)
to apply between ticks. Rank 0's step loop and its admin HTTP
handlers share one event loop, and ``profile_action`` does not
drain in-flight requests (profiling wants live traffic) — so its
``ControlOp(op="profile")`` broadcast can be enqueued between an
active slate's ``(kind, payload)`` tuples. Without this guard the
tuple unpack below would raise ``TypeError`` on the dataclass, the
worker would fall out of the dispatch mid-slate, and the slate's
remaining tuples would then poison ``_recv_control`` (expected a
``ControlOp``, got a tuple), desyncing the TP group. Deferring —
rather than applying inline — keeps every admin op on the
"applied between ticks, never racing a half-run slate" contract
the drained ops (swap / config_override) already rely on.

Per-op payload schemas for the TP worker-bridge hand-off.

The dataclasses here are the wire payloads rank 0 broadcasts to worker
ranks and the workers receive + execute (see
:mod:`arbi_serve.distributed.worker_bridge`). They are split out so both
the bridge classes and the worker-side executors can import them without
pulling in the heavier engine/spec-decode machinery.

One forward over a materialized batch.

Carries the :class:`StepPlan` (existing wire format) so workers
reconstruct the :class:`ScheduledBatch` and run forward. The flavor
selects the worker-side dispatch — it must reproduce rank 0's
collective sequence 1-to-1:

* ``"plain"`` — rank 0 ran the rung-aware ``execute`` dispatch; the
  worker runs the rung-aware ``forward(batch)`` (both sides resolve
  the same batch-keyed replay rung, or both run eager).
* ``"verify"`` — MTP verify pass (``max_query_len = K+1``); the worker
  mirrors the verify forward + the per-token ``lm_head`` parity call.
* ``"seed"`` — rank 0 ran the raw eager ``_run_model_forward(
  return_hidden_state=True)`` (MTP drafter-seed / embed-rerank steps
  that need per-token hidden, which no captured graph retains). The
  worker takes the identical raw-eager dispatch: a rung-aware
  ``forward(batch)`` here could replay a padded prefill-bucket graph
  whose baked collectives run at the bucket size against rank 0's
  live-size eager collectives — a cross-rank NCCL size mismatch that
  wedges the TP group.

One drafter-chain forward.

Workers reconstruct ``hidden_in`` / ``last_token_in`` from the
broadcast lists, rebuild K per-step :class:`AttnPagedKVMeta`s from
the realized rank-0 slot allocation + page-table rows, and run the
head's K-step chain against their own KV slab + RoPE cache —
replaying the captured chain when ``replay_captured`` is set (rank 0
did too) else the live eager chain. Either way each step's
:class:`RowParallelLinear.o_proj` / MLP ``all_reduce`` matches
rank 0's so the forward stays in lockstep.

Slot allocation lives on rank 0's pool only; the worker writes
KVs into the same slot indices in its own per-rank slab (the slab
geometry is identical; only the head dim is sharded). This ships
``per_step_slots`` (B × K) plus the page-table row per request so
the worker rebuilds a verbatim ``attn_meta``.

One leg of a DFlash draft step's vocab-parallel collective.

Unlike the bundled MTP head, the DFlash drafter runs entirely on
rank 0: it owns the separate 5-layer draft model + all per-request
draft-KV state, and those dense draft layers carry no TP collectives
of their own (the draft model is replicated, never sharded). The only
cross-rank work is the two vocab-parallel collectives the draft
borrows from the target — ``embed_tokens`` (all_reduce) on the masked
noise block, and ``lm_head`` (all_gather) on the denoised hidden.
Both are unconditional inside the parallel module forwards, so rank 0
deadlocks unless every rank enters them in lockstep.

Rank 0 therefore brackets its local draft with two broadcasts:

  1. ``phase="embed"`` carrying the ``(B, block)`` masked block ids —
     every rank runs ``embed_tokens(blk_ids)`` (all_reduce); workers
     discard the result. Rank 0 then runs its draft model locally.
  2. ``phase="lmhead"`` carrying the denoised hidden ``(rows,
     hidden)`` — every rank runs ``lm_head(hidden)`` (all_gather);
     workers discard. Rank 0 keeps the gathered logits and argmaxes.

Both legs are batched across all draft rows, so it is exactly two
collectives per draft step at a fixed shape regardless of batch size
or disabled rows. Workers never load the draft model and hold no
draft state — they only rendezvous. Mirrors the verify pass's
bare-head collective parity call.

One full DFlash draft step on a tensor-sharded drafter.

Unlike :class:`DFlashDraftOp` (the replicated build, where the draft
model runs on rank 0 alone and workers only join the two vocab-parallel
collectives), the sharded build runs the whole draft forward on every
rank against that rank's head/inner shard. Each draft layer's
``o_proj`` / ``down_proj`` is a :class:`RowParallelLinear` carrying an
``all_reduce`` — so the per-layer draft forward now emits ``2 * L``
collectives that rank 0 cannot run alone. Rank 0 broadcasts this one op
at the top of :meth:`DFlashDrafter.draft`; every worker rebuilds the
same masked block ids + its own shard context and runs the identical
``embed → forward_block_fixed → lm_head`` sequence, so the whole
collective train (embed all_reduce, 2·L row-parallel all_reduces,
lm_head all_gather) rendezvous 1-to-1. Workers discard every output —
only rank 0 samples; the workers exist to make the shard all_reduces
sum the full head space.

Slot ids are rank-0-allocated but valid on every rank: the per-rank
slabs share geometry (only the head dim is sharded), and each worker
populated the same slot ids via :class:`DFlashObserveOp`. The worker
reads its own ``slots.lengths[slot]`` for the host start_pos, so no
start position is shipped on the host path; ``max_full_len`` carries
the sync-free width bound rank 0 computed for the device-slots path.

Populate a tensor-sharded drafter's per-rank shard KV.

The observe / context-projection (:meth:`DFlashDraftModel.project_context_stacked`)
carries no collective (K/V projections are column-parallel), so this op
never risks a deadlock — but a worker whose shard KV is empty makes its
``o_proj`` all_reduce contribute a wrong partial, so the summed hidden
(and thus rank 0's draft) drifts and accept degrades. Rank 0 broadcasts
the resolved per-row stash spec after each verify / seed forward; every
worker gathers the same tap features from its own (replicated) tap
capture, projects its head/inner K/V shard, and appends at the same
slot + absolute position. Only the committed prefix is shipped, so the
worker slab stays byte-aligned with rank 0's.

Seed-buffer fill.

Rank 0 advances ``mtp_seed_counter`` and broadcasts the new value;
every rank's ``Engine.mtp_seed`` is filled with that value before
any sampling kernel runs. Ensures captured Gumbel-max draws are
identical across ranks.

Partial-accept recurrent rollback (GDN/conv state).

The verify forward mutates each rank's local recurrent slab shard
(the kernel writes its half of the heads in lockstep), but the
post-verify partial-accept rollback is a pool op the rank-0 driver
issues outside the forward. Without this broadcast the worker's
shard keeps the rejected draft tokens' state every partial-accept
step — the worker half of every GDN layer drifts and greedy
spec-decode output diverges at TP>1. Workers apply the same
``rollback_partial_accept(slab_rows, n_accepted)`` against their own
pool; the snapshot rows were written by their own verify forward at
the same ``state_indices``.

Worker-side executors for the TP worker-bridge hand-off.

These are the functions worker ranks run to mirror a rank-0 forward /
drafter / DFlash draft step so the cross-rank collective sequence stays
in lockstep (see :mod:`arbi_serve.distributed.worker_bridge`). They live
in this module for file-size hygiene; ``worker_bridge`` re-imports them
so ``worker_bridge.<name>`` stays importable (tests reach them through
that module).

Stage the verify step into the worker's persistent VerifyBuffers.

Mirrors the rank-0 ``load_step`` call in
:func:`spec_decode.mtp.build_verify_plan`. The captured verify graph
is ``shares_verify_buffers=True``, so ``CapturedGraph.replay`` does
not copy the per-step inputs into the graph's buffers — it assumes
``load_step`` already wrote the current step into the same storages
the captured kernels reference by ``data_ptr``. The worker reaches
the replay through a freshly reconstructed batch and would otherwise
leave VerifyBuffers holding stale ``block_table`` / ``seq_lens``,
corrupting the TKV page-metadata / split-K KV-page gather.

Returns the batch rebound to the loaded view's slices (so its
tensors alias the VerifyBuffers storage the captured graph reads)
and otherwise no-ops — returning the input batch unchanged — when
the engine has no VerifyBuffers (MTP disabled) or the captured
verify graph for this shape is not shared-buffer (e.g. capture
skipped at boot, replay falls through to the live forward which
reads ``batch`` directly).

Join one DFlash draft collective leg on a worker rank.

The DFlash drafter runs in full on rank 0 (it owns the 5-layer draft
model + per-request draft-KV state); workers only enter the two
vocab-parallel collectives so rank 0's ``embed_tokens`` all_reduce /
``lm_head`` all_gather rendezvous 1-to-1. The result is discarded —
mirrors the verify pass's bare-head collective parity
call. The draft model's own dense layers carry no collectives and run
on rank 0 alone; workers never see them.

The head mirrored here is the DRAFTER's, not ``model.lm_head``: under
``--draft-vocab-prefix`` rank 0 projects through a vocab-parallel WINDOW
whose all-gather runs at the window width, while the full head's runs at
the shard width. Two widths on the same collective is not a wrong answer,
it is a wedged group. (The verify mirror in ``worker_bridge`` keeps the
full head, which is the head verify actually uses.)

Run one full sharded DFlash draft step on a worker rank.

The sharded drafter runs its own head/inner shard on every rank, so a
worker does not merely join the two vocab-parallel collectives — it
re-runs the whole ``embed → forward_block_fixed → lm_head`` train
against its shard. Each draft layer's ``o_proj`` / ``down_proj`` is a
:class:`RowParallelLinear` whose ``all_reduce`` sums the head space
across ranks, and the two vocab-parallel legs (``embed_tokens``
all_reduce, ``lm_head`` all_gather) bracket it — so rank 0's whole
collective sequence rendezvous 1-to-1. Every output is discarded (only
rank 0 samples); the worker exists to make the shard all_reduces sum
the full head space so rank 0's draft is byte-correct.

Slot ids come from rank 0 but are valid here (shared slab geometry +
:func:`_run_worker_dflash_observe` populated the same ids). The host
start position is read from this rank's own ``slots.lengths`` (kept in
step by observe); ``max_full_len`` carries the device-slots width bound.

Populate this rank's DFlash shard KV from the local tap.

Projection carries no collective, so this never deadlocks — it exists
so the worker's shard ``o_proj`` contributes a correct partial to the
draft all_reduce. The worker gathers the same tap features from its own
(replicated) tap capture, projects its head/inner K/V shard, and appends
at the rank-0-allocated slot + absolute position. The slot's length is
forced to ``first_pos`` before the append so a reused slot id (rank 0
released + re-acquired) is repositioned correctly on this rank, which
never runs the rank-0 acquire/release path.

Run the drafter chain on a worker rank.

Mirrors the live-chain branch of :meth:`MtpDriver.draft` against
the worker's local engine. Each chain step runs the head directly
(no ``draft`` re-entry) so we don't invoke the rank-0-only slot
allocator + page-table lookups; rank 0 ships the realized slot
indices + page-table rows in :class:`DrafterOp` and the worker
rebuilds identical per-step :class:`AttnPagedKVMeta`s. Each step's
:class:`RowParallelLinear` ``all_reduce`` keeps the per-rank
kernels in lockstep with rank 0.

Tokens and probabilities are discarded on workers — only the
collective participation matters; rank 0 owns sampling decisions
and request state. ``op.K`` matches the rank-0 ``k`` so the chain
length / autoregressive carry sequence is identical.

Replay the worker's captured drafter chain (SPMD parity).

Mirrors :meth:`MtpDriver._replay_captured_chain` on rank 0 — same
captured ``(B, K)`` bucket, same ``replay`` call — so the K
all_reduce + K all_gather collectives baked into rank 1's graph
rendezvous 1-to-1 with rank 0's. Tokens / probs are discarded; only the
collective participation matters.

``sampling_params`` picks the pool:

  - ``None`` (greedy): rank 0 replayed its ``drafter_graphs`` (argmax)
    pool, so the worker replays the same greedy bucket.
  - non-``None`` (true-stochastic): rank 0 replayed its
    ``drafter_graphs_stoch`` pool (sampled carry), so the worker replays
    the same stochastic bucket, encoding the per-row params + the
    rank-agreed ``seed_buf`` into its persistent buffers. Its
    counter-keyed draws are byte-identical to rank 0's (same seed, same
    ``OFFSET_DRAFTER + 17 * step`` lanes), so its per-step carry tokens —
    and therefore its K/V writes — match rank 0's with no token exchange.

Either pool's captured graph bakes the identical per-step o_proj
all_reduce + lm_head all_gather, so the collective count matches
regardless of which pool is replayed.

Fails loud on a pool miss: rank 0 only set ``replay_captured`` because
its pool hit, and both ranks capture identical buckets at boot, so a
worker miss is a real asymmetry that would desync the collective count.

Drafter-scoped vocabulary-PREFIX lm_head (``cfg.draft_vocab_prefix``).

The drafter re-reads the lm_head vocab shard once per chain step, so on a
248,320-wide head that projection is the largest single cost in a B=1
decode step. This lever gives the DRAFTER a narrower VIEW of the lm_head —
token ids ``[0, N)`` — while the verifier keeps scoring through the whole
of it, at the cost of the drafter being unable to propose any id ``>= N``.

**The win is BANDWIDTH, and only bandwidth.** What shrinks is the bytes
moved per draft step: the drafter's weight read, its logit row, and (under
sampled drafting) the softmax, the top_k/top_p mask and the dense ``q``
write that ride that row. It is not a memory optimisation either — the
window head allocates nothing, so the lever is VRAM-NEUTRAL, not
VRAM-negative. ``exl3_gemm`` takes ``N`` as an output-width bound and reads
only the leading trellis blocks of the verify head's own tensors; there is
no second head.

**Default ``"full"``.** The lever is OFF unless asked for.

  ``full``               the drafter keeps the whole lm_head — THE DEFAULT
  ``half`` / ``quarter`` that fraction of THIS checkpoint's lm_head, rounded
                         DOWN to a whole :data:`PREFIX_ALIGN` block
  ``N``                  an explicit width; every ineligible combination
                         fails LOUD at boot

The fractions are fractions rather than constants on purpose. "Half the
vocabulary" is a statement about the model in front of you; a fixed rung is
a number that happens to be half of one particular head and 86% of another.
See ``docs/draft-vocab-prefix.md``.

What a boot can and cannot answer
---------------------------------
Two risks look alike and are not:

* **Checkpoint portability** — *is this vocabulary frequency-ordered?* This
  IS a property of the checkpoint, so a requested cut is measured against it
  at boot (:func:`probe_prefix_order`) and WARNED about when the answer is
  no.
* **Traffic language** — *will this deployment serve CJK / Cyrillic / Arabic?*
  This is NOT a property of the checkpoint. A boot sample knows nothing about
  future traffic, so nothing here tries to answer it. The per-script table is
  printed at the rung actually chosen so the operator can.

Because the second question is the operator's, no width is ever chosen for
them — which is why the default is ``full`` and every cut is named.

Why a PREFIX
------------
The prefix is what makes the install free of a remap: the argmax index over
the sliced row IS the global token id, so nothing is added to the captured
drafter graph downstream of the projection. It is also, on this
vocabulary, already close to a frequency ordering — BPE ids are emitted in
merge order, so a low id is a token the tokenizer's own corpus saw often.

It is not the ONLY valid slice. The binding constraint is the 128-wide
Hadamard below, not the trellis codebook: a scatter of whole 128-blocks
reproduces the full head just as exactly as a prefix does (measured on
Qwen3.5-0.8B-exl3 and Qwen3.8-27B-exl3 lm_heads, max |diff| 0.002 on logits
of scale 1). A frequency-ordered shortlist at 128-block granularity is
therefore constructible — it is simply not built here, because it buys
almost nothing over the merge order it would replace and it would cost the
remap this design does without.

Under tensor parallelism
------------------------
The lm_head is vocab-sharded at TP>1, so no rank holds ``[0, N)`` — rank
``r`` holds global columns ``[start_r, start_r + width_r)``. The window is
therefore taken PER RANK, as the intersection of the global window with that
rank's own columns, and the ranks compose it back through the all-gather the
column-parallel head already fires
(:class:`~arbi_serve.weight_quant.exl3.linear
.EXL3OutputPrefixVocabParallelLinear`). Because the global window is a
prefix and the shards are contiguous in rank order, that intersection is
always a LEADING window of the rank's shard — the same output bound the
replicated window applies — so nothing is copied and nothing moves between
ranks.

The point is that the DRAFTER IS THE SAME DRAFTER at every rank count. Its
proposal set is ``[0, N)`` at TP=1 and at TP=2, so an accept rate measured
at one rank count is comparable with the other. Without this the lever is
simply absent at TP>1 and the two deployments run different drafters.

What the lever still BUYS varies with ``N``, and the boot log says which:

* ``N`` at or below a shard's width — every rank's GEMM and the gather
  itself shrink with ``N``, the same relative cut as at TP=1;
* ``N`` above it — the low ranks' windows ARE their whole shards, so their
  GEMM and the gather are exactly what they were; the narrowing lands on
  the straddling rank alone. TP has already narrowed each rank's read
  further than the request, and what remains is the drafter's REACH.

Losslessness
------------
Emitted output is unchanged, and that is a property of the accept test
rather than a hope. Rejection sampling is lossless for ANY proposal ``q``
(Leviathan-2023 Thm 1), and a ``q`` supported on ``[0, N)`` is a valid
proposal. Verify still scores against the FULL-vocab target ``p`` — it
projects through ``engine.model.lm_head``, which this lever never touches —
so a token outside the prefix is still EMITTABLE, through the residual
branch. It is simply never DRAFTED. A narrower drafter row can therefore
only change ACCEPTANCE, never the emitted distribution.

That holds in both drafting regimes. Under a greedy draft the proposal is
the ``POINT_MASS_Q`` marker (:mod:`arbi_serve.spec_decode.drafter`) and the
slice only moves WHICH token the mass sits on. Under sampled drafting
(``ARBI_TRUE_STOCHASTIC_DRAFT``, default ``"auto"``) the drafter builds a
dense ``q`` that is ``N`` wide, and the verify seam copies it into the
prefix of a ZEROED full-width residual buffer — so ``(p - q)⁺`` is exactly
``p`` above ``N``. Restricting ``q`` to the prefix RENORMALISES it, which
raises ``q(x)`` toward ``p(x)`` wherever the drafter under-weighted the
target — and because the accept rate is the overlap ``Σ min(p, q)`` rather
than a conditional accept probability, that recovers most of the coverage
the slice gives up. Measured paired at the sampler seam on the reference
checkpoint at half the vocabulary: coverage -0.0920% against renormalisation
+0.0832%, net **-0.0088%** at temperature 0.7, and **-0.1897%** at 1.0.

The alignment refusal
---------------------
EXL3 packs the vocab axis in 16-element trellis blocks but applies a
128-wide Hadamard along it at runtime. The Hadamard is block-diagonal in
128, so a slice reproduces the unsliced result exactly when — and only
when — its offset and width are whole 128-blocks; this is the same
constraint
:class:`~arbi_serve.weight_quant.exl3.linear._EXL3HadAlignedShardMixin`
enforces on a TP shard, here at offset 0. Measured on the same two heads: a
16-aligned slice at offset 16 is off by |5| on logits of scale 1, and a
scatter of bare 16-blocks by |4|. 128 is therefore the binding alignment
and a 16-aligned-but-not-128-aligned N is REFUSED, not clamped.

Import discipline
-----------------
This module lives at the package ROOT, not under
:mod:`arbi_serve.spec_decode`, and imports nothing heavy at module scope
(the EXL3 builder and the runtime-flag read are function-local). The
per-emitted-token counters are bumped from
:func:`arbi_serve.engine.run_step.terminal.post_token`, which must bind them
ONCE at import; importing them out of ``spec_decode`` would run that
package's ``__init__`` — which reaches back into ``engine.run_step`` — and
close an import cycle. Same discipline, and the same reason, as
:mod:`arbi_serve.flag_truth`.

The map between a drafter ROW INDEX and a TOKEN ID.

The drafter's projection emits one row of ``width`` logits. Which token
each column of that row MEANS is this object and nothing else. For the
shipped prefix window the two coincide — column ``i`` is token ``i`` —
and that identity is exactly why a prefix needs no map. It is also why
every consumer wrote its own arithmetic: with one window there is no
arithmetic to write.

A reach with a second window breaks the identity, and the failure mode is
not an exception — it is a drafter that proposes token 131,072 when it
meant ``</think>``, silently, on the accept path. So the two spaces get
two names and one owner: :meth:`token_ids` converts, nothing else may,
and :attr:`windows` is the only place the ranges are written down.

``windows`` are ``(token_id_start, width)`` in ROW order, ascending and
disjoint. Both bounds are whole :data:`PREFIX_ALIGN` blocks because the
EXL3 output Hadamard is block-diagonal in 128 — the constraint that makes
a window reproduce the source head exactly, and the reason a reach cannot
simply be a list of token ids.

``[0, n_keep) u [vocab - tail, vocab)`` — a prefix plus the head's tail.

The tail exists for one reason: a checkpoint's ADDED tokens are appended
ABOVE every BPE id, so no prefix short of the whole head reaches them —
``</think>`` and ``<|im_end|>`` among them, which a thinking-ON request
emits once each. What that costs a prefix-only drafter is reported at
``draft_vocab_prefix_missed``; :func:`derive_draft_vocab_tail` chooses the
window from the ids themselves so it does not have to be typed.

Process-global row-index -> token-id TABLE for the installed reach.

A held tensor rather than a :class:`DraftReach` for the same reason
:class:`_EmitWindow` holds an int: the read sits inside the drafter block
walk, once per position, and must cost one attribute load and one branch.
``table is None`` — the default, and what an IDENTITY reach installs — is
the whole cost when the lever is off or is a plain prefix.

Process-global, like every path counter and like the emit window, and for
the same reason: two co-resident engines at different reaches would share
one table. A reach is read per process, not per model — which is safe
today because the drafter is refused outright when a second override is
installed, and which is stated here rather than discovered.

The installed reach's TRAILING window as ``(row_start, id_start,
width)``, or ``None``. Held beside the table for the same reason: the
``q`` relocation runs once per drafted slate and must cost one load and a
branch when there is nothing to relocate.

The installed reach's row-index -> token-id table, or ``None``.

``None`` means the drafter's row index IS its token id — no lever, or a
prefix — so a consumer's conversion is a no-op it can skip entirely.

Publish the installed reach to the drafter walks (``None`` clears).

Called once at install, beside :func:`set_emit_reach`, so the two cannot
describe different reaches: a counter that says ``</think>`` is reachable
while the walk still maps its row to 131,072 would report the lever
working in the arm where it is not.

Move a union row's TRAILING window from where it was WRITTEN to where
its token ids are, in a full-vocab ``q`` buffer.

The drafter walk writes its row into ``dense_q``'s leading columns because
that costs nothing — the destination is a view, not a copy. Under a prefix
those columns already ARE their token ids and there is nothing to do. Under
a union the leading window still coincides, and only the trailing window is
misplaced: it was written at rows ``[N, N + T)`` and belongs at ids
``[V - T, V)``. Moving that one slice is two static ops on ``T`` columns,
against a row-space scratch plus a full copy — which is why the walk is
left writing where it does.

The vacated band must be CLEARED, not left: it holds the tail's
probabilities at ids the drafter never proposed, and the verify seam reads
``(p - q)+`` over the whole row.

No-op when no reach is installed or the reach is a prefix.

Convert drafter ROW INDICES to TOKEN IDS under the installed reach.

The one conversion every drafter walk calls. Two loads and a branch when
the reach is the identity, which is every deployment that does not ask for
a non-prefix reach.

Process-global ``N`` for the per-emitted-token miss counter.

A slotted attribute rather than a config read: the check sits on the
per-token commit path (:func:`arbi_serve.engine.run_step.terminal
.post_token`), where it must cost one attribute load and one integer
compare. ``n == 0`` (the default) short-circuits before anything else,
so a server without the flag pays a single predictable branch.

Arm the per-emitted-token counters from the installed reach.

The counters answer "could the drafter have proposed this token", which is
a question about the REACH. Deriving it from a bare width was correct only
while the reach was a prefix; a union whose counter still tested
``token >= n`` would report every ``</think>`` as unreachable in the very
arm that fixed it.

Arm (``n_keep > 0``) or disarm (``0``) the per-emitted-token counters.

Process-global, like every :class:`~arbi_serve.flag_truth.PathCounter`:
under multi-model residency the miss rate is reported for the process, not
per engine. Two co-resident engines at DIFFERENT prefixes would therefore
share one window — read the counters per process, not per model.

Count one draft dispatch through the sliced head.

Returns whether ``head`` IS a sliced prefix head, so the caller can
attribute the step to exactly one counter.

A dispatch through a VOCAB-PARALLEL window is counted separately. Under
TP the window only means anything if the ranks actually composed it, and
the difference between "composed" and "each rank quietly drafted from its
own shard" is invisible in the emitted tokens — so it is read off a
counter rather than inferred. Called from the eager driver body, never
from inside a captured region, so one fire is one SERVED step.

Refuse a structurally impossible setting before any weight is loaded.

Pure arithmetic — the model-dependent refusals (no MTP, no EXL3 head,
``N`` wider than the vocabulary) live in :func:`apply_draft_vocab_prefix`.
Both refuse BY NAME; neither clamps. A silently clamped N would serve a
drafter whose reach the operator cannot read off their own config.

``full`` and the fractions validate trivially: neither names a width
here. A fraction's width comes from the checkpoint (:func:`fraction_width`
rounds DOWN to a whole :data:`PREFIX_ALIGN` block), so it satisfies the
alignment by construction and has nothing left to check.

Per-script counts of the vocabulary, and how many fall below each rung.

``decoded[i]`` is the text of token id ``i``. A token counts once for
EVERY class it contains, so ``sum(totals) >= len(decoded)`` — the
question each row answers is "how many tokens containing this script
can the drafter still propose", which is what a mixed token (a CJK
character with an ASCII space) has to be counted in both rows for.

A pure function of the checkpoint: no traffic, no sample, no model.

Fraction of :data:`VOCAB_ORDER_PROBE`'s tokens below each rung.

USAGE-weighted, which is the whole point: "frequency-ordered" is a claim
about which ids the tokens people actually emit land on, and a token
COUNT cannot see it (see :func:`probe_prefix_order`).

Measure whether an id PREFIX is a shortlist on this checkpoint.

A GUARD on a cut the operator asked for, never a chooser. It answers the
one question a boot sample CAN answer — *are the frequent tokens at low
ids?* — which is a property of the checkpoint and of nothing else. The
question it cannot answer is *which languages will this deployment
serve?*; that is traffic, so the per-script table is REPORTED for the
operator rather than voted on here.

Both halves of that split are measured facts:

* **Why the probe, and not the table.** A token COUNT is dominated by
  each script's rare long tail, so it cannot see frequency ordering. On
  the reference checkpoint at half the vocabulary, ``ascii_punct`` sits
  at ``100.00%`` of the
  probe's emitted tokens. Gemma-4 and gpt-oss-20b are likewise
  indistinguishable in the table (every class within ``0.9833``).
* **Why the table cannot be a gate.** At half the vocabulary the
  reference checkpoint already strands six of its own script classes —
  Cyrillic Arabic Thai Devanagari
  Greek Hebrew A gate protecting them
  would refuse the very checkpoint this lever was measured on.

A failing probe WARNS rather than refuses. Every other refusal in this
module names a thing that cannot be built — no MTP head, no EXL3 head, a
TP-sharded head, an unaligned ``N``. An unordered vocabulary builds fine
and serves correct output; it only accepts worse, and that is a
magnitude the deployment reads off ``draft_vocab_prefix_missed`` on its
own traffic. Failing a boot on a sampled statistic would also make boot
success depend on a corpus file, which no structural refusal here does.

Width of the GLOBAL lm_head a window would cut, or ``0`` if unreadable.

The fractions are taken against THIS number rather than the tokenizer's
vocabulary because it is the axis actually being windowed: a head padded
past the tokenizer's ids is still a head of that width, and a fraction of
the tokenizer's count could land inside the padding or, on a head padded
down, above it.

Under TP the head is vocab-sharded and ``svh`` carries only THIS rank's
columns, so the width comes from the column-parallel head's global output
size. A fraction resolved against the shard would name a different
vocabulary on every rank AND a different one per rank count — "half" would
mean half the model at TP=1 and a quarter of it at TP=2.

``vocab_size / divisor`` rounded DOWN to a :data:`PREFIX_ALIGN` block.

Rounding DOWN, not to nearest: the alignment exists because the runtime
Hadamard is block-diagonal in 128, so a width that is not a whole number
of blocks admits a partial block and decodes different weights. Rounding
up could also exceed the vocabulary on a head whose width is not itself
a multiple of ``divisor * PREFIX_ALIGN``.

Normalise ``cfg.draft_vocab_prefix`` into ``(mode, n)``.

``("full", 0)`` — the DEFAULT: no slice; ``("half", 0)`` /
``("quarter", 0)`` — a fraction of the loaded checkpoint's own lm_head,
resolved at boot by :func:`fraction_width`; ``("explicit", N)`` — this
width or a boot failure.

What :func:`derive_draft_vocab_tail` chose, and why.

``refusal`` is ``None`` on a derivation and a machine-readable reason
otherwise; a refusal always carries ``tail == 0``, which is a plain
prefix. :attr:`message` is the operator-facing sentence for either case —
a derivation that cannot say which ids it covers, or a refusal that
cannot say what it could not cover, is not reportable.

The trailing window that reaches this checkpoint's ADDED token ids.

The width is not a number an operator should have to know. A checkpoint
appends its added tokens ABOVE every BPE id, so WHICH ids they are is a
property of the tokenizer being served; the window that reaches them is
the whole ``align``-blocks spanning them, because the EXL3 output
Hadamard is block-diagonal in ``align`` and a window that splits a block
decodes different weights.

Only the BOTTOM end is free. ``draft_vocab_tail`` names a window anchored
at the head's top, so the derivation rounds the low end OUT — DOWN to a
block boundary — and takes everything above it. Rounding it the other way
would produce a window that silently excludes the very tokens it was
chosen to reach, which is the property :func:`validate_draft_vocab_prefix`
protects for an explicit width and this must not lose.

Every shape the derivation cannot describe is REFUSED by name and falls
back to a plain prefix, because a window that reaches none of what it was
chosen for costs drafter row width and buys nothing, and no counter on the
served path would say so. The shapes it CAN describe are exactly: a
contiguous band of ids, above every base BPE id, above the prefix, inside
a head whose top is itself a block boundary.

Ids the tokenizer reports as ADDED, ascending; ``()`` when unreadable.

Read off :attr:`arbi_serve.tokenizer.Tokenizer.added_token_ids`: the class
holding the vocabulary is the one that can name its added tokens. An
object that does not publish that attribute is unreadable, and the caller
refuses by name on an empty result rather than deriving a window.

Count of BASE (non-added) ids the tokenizer holds, ``0`` if unreadable.

Read off :attr:`arbi_serve.tokenizer.Tokenizer.base_vocab_size`, and off
nothing else. ``vocab_size`` names a different quantity on that class —
the id space WITH the added tokens — and a base count taken from it puts
the added band inside the base range, where the derivation refuses it. An
object that does not publish ``base_vocab_size`` is unreadable and returns
0, which the caller reports as its reason.

Normalise ``cfg.draft_vocab_tail`` into ``(mode, n)``.

``(DERIVE, 0)`` — the DEFAULT: the window comes from the checkpoint's own
added token ids at boot (:func:`derive_draft_vocab_tail`), and a shape the
derivation cannot describe falls back to a plain prefix with its reason
logged. ``("explicit", 0)`` — a plain prefix, which is the control arm and
stays expressible. ``("explicit", N)`` — this width or a boot failure.

This model / flag combination cannot serve a sliced drafter head.

A subclass of ``RuntimeError``: every width that reaches the install was
named by the operator, so a refusal is a failed ASSERTION and fails the
boot. The distinct type exists so the caller can record :attr:`reason` on
``draft_vocab_prefix_refused`` and re-raise, rather than catching a bare
``Exception`` and mislabelling a real fault — an allocator failure while
building the head, a broken checkpoint — as an ineligible config.

Which drafter this boot will install, from config alone.

Mirrors the source dispatch in
:func:`~arbi_serve.engine.mtp_attach.build_mtp_driver` (a named external
source wins over the checkpoint's bundled head) and the identical order in
:func:`~arbi_serve.flag_reachability.drafter_facts`. Config-only because
the install phase runs BEFORE the drafter is built — see
:func:`draft_lm_head_owner`.

The subset of :func:`apply_draft_vocab_prefix`'s refusals that is
decidable from CONFIG ALONE — no model, no weights, no shared state.

Split out so the SAME refusal can be raised twice from one source of
truth: once by :func:`assert_config_eligible` as a pre-flight, before a
build is allowed to touch anything a still-serving member shares, and once
by :func:`apply_draft_vocab_prefix` at the install seam, which is reached
on paths that never ran the pre-flight.

``request`` is the operator's setting as written (a width, or a fraction
rung), so the message names what was asked rather than a width the
fraction has not resolved to yet.

Returns the refusal to raise, or ``None`` when nothing visible in the
config objects. A ``None`` is NOT an eligibility verdict: every remaining
refusal needs the built model (head class, TP sharding, bound vocabulary).

Raise this lever's config-only refusals against ``cfg``.

Reads nothing but the config, so it is callable before a member build has
bound a weight or mutated anything the ACTIVE member shares. The install
phase runs deep inside the build (after the donor-share bind and the
compaction release hooks, which deregister EXL3 kernel descriptors the
donor's own linears are addressed by), so a refusal decided there lands
after the mutation and there is no eligibility left to recover — see
:func:`arbi_serve.engine.build_phases_load.assert_config_only_eligibility`.

A no-op when the lever is off, and never a verdict that the config IS
servable: the model-dependent refusals still fire at the install seam.

The object owning the drafter's lm_head seam for this boot.

The bundled MTP head when the checkpoint builds one, otherwise the model
itself. DFlash and external drafters replace that head entirely
(``mtp_attach.maybe_append_mtp_spec`` returns ``False``), so under them it
is never built and the seam has to live somewhere that is.

The seam is deliberately NOT the drafter object. The install phase must run
before ``_profile_and_size_kv`` so the sliced head's VRAM is charged to the
KV budget, and the DFlash drafter is not constructed until after it
(``build_phases_kv`` -> ``build_mtp_driver``). Installing on the model and
having the drafter read it at draft time is what lets one phase serve both
orders.

Width ``N`` of an installed vocabulary-prefix head, or 0 if none.

0 for a head-quant ``drafter-*`` override too: that override replaces the
projection without narrowing it, and every width consumer must keep
reading the full vocabulary through it.

Per-rank width of the global window ``[0, n_keep)`` on a sharded head.

Rank ``r`` of a vocab-sharded EXL3 head holds global columns
``[start_r, start_r + width_r)``, contiguous and in rank order, so the
intersection of the global window with rank ``r``'s shard is always a
LEADING window of that shard: ``clamp(n_keep - start_r, 0, width_r)``.
That is why the per-rank half of this lever is the same operation the
replicated window already performs — an output bound — and why nothing
has to move between ranks.

The shard geometry is DERIVED here (a pure function of the padded
vocabulary, the rank count and the 128-wide Hadamard block) rather than
exchanged, so every rank computes the same table for every rank. The
derivation is then checked against the head THIS rank actually holds: a
table that does not reproduce the local shard is a table whose other
entries are guesses, and trimming a gathered row by a guess is silent
corruption rather than a wrong answer.

Build an EXL3 linear over vocab ids ``[0, n_keep)`` from ``lm_head``.

``lm_head`` must be a bound EXL3 linear (the trellis lane). The window
head shares its ``trellis`` / ``suh`` / ``svh`` tensors by reference and
allocates nothing: output columns live on the trellis's middle dimension
in output order — the same fact
:meth:`~arbi_serve.weight_quant.exl3.linear.EXL3ColumnParallelLinear
.exl3_load` slices a TP output shard on — so ids ``[0, n_keep)`` are the
LEADING ``n_keep // 16`` blocks and ``svh``'s leading ``n_keep`` entries.
``exl3_gemm`` reads exactly those from the full tensors, given ``n_keep``
as its ``size_n_out`` bound, and keeps B's k-row pitch at the stored
width. The lever therefore costs no VRAM: the verify head's copy IS the
drafter's.

Sharing binds the same TENSOR OBJECTS, not just the same storage, so a
``.data`` rebind on the source head carries to the window. What does not
carry is the source's compaction hook, which rebuilds only ITS kernel
descriptor — so this must stay downstream of the flat loader's compaction
(``loader.flat_loader._binders``), where every rebind has already fired.

A VOCAB-SHARDED source head takes the same window per rank, over the
columns that rank holds (:func:`prefix_window_geometry`), and composes
them through the column-parallel all-gather the head already fires — see
:class:`~arbi_serve.weight_quant.exl3.linear
.EXL3OutputPrefixVocabParallelLinear`. The drafter's vocabulary is then
``[0, n_keep)`` at every rank count, which is what makes a TP1-vs-TP2
comparison a comparison of the same drafter.

Build the drafter head for ``reach`` over ``lm_head``'s own tensors.

An identity reach is exactly :func:`build_prefix_draft_head` and takes it,
so the shipped lever keeps its zero-resident window and its whole test
surface. Every further window is bound over a MATERIALIZED slice of the
same trellis and ``svh``, at a 128-aligned output offset — the cut
:meth:`~arbi_serve.weight_quant.exl3.linear.EXL3ColumnParallelLinear
.exl3_load` performs for a TP shard, here at an offset the drafter chose.

The union is refused under TP: the per-rank intersection in
:func:`prefix_window_geometry` is derived from the window being a PREFIX,
so a rank's share of it is always leading and nothing moves between ranks.
A trailing window lands on one rank and the derivation says nothing about
it — the honest state is a named refusal, not a table that is right for
one window and a guess for the other.

Install the drafter's vocabulary-window head on ``model``.

Every refusal names the mechanism it could not find, in the shape of
:func:`~arbi_serve.weight_quant.head_quant.apply_head_quant`'s: an
explicitly requested lever that cannot mean what it says is a boot
failure, never a silent no-op.

Raises :class:`DraftVocabPrefixIneligible` — a ``RuntimeError``, so the
boot fails. There is no setting that asks this module to pick a width:
an explicit ``N`` is an assertion, and a fraction has already been
resolved against this checkpoint's own head before the call. The
fractions' own "cannot be measured" cases disable earlier, in
:mod:`arbi_serve.engine.build_phases_load`, where no width has been
named yet for a failure to be about.

``(row_start, token_id_start, width)`` per window, in row order.

The one form a range-wise consumer needs: a dense-``q`` scatter copies
``q[..., row_start:row_start+width]`` to ``dest[..., id_start:id_start+width]``
per entry, and never has to know how many windows there are.

``(width,)`` int64 row-index → token-id table, cached per device.

Capture-safe: built once at install, read by an advanced index inside
the captured region, so it costs a static gather and no allocation.

Convert a tensor of ROW INDICES to TOKEN IDS.

The identity reach returns its argument, so the shipped prefix path
pays nothing and still goes through this seam — which is the point:
one code path, exercised by every existing test.

Engine: continuous-batching loop, per-request state, per-step batch struct.

``Engine`` is intentionally NOT re-exported here: importing it eagerly
would close a circular import chain
(scheduler -> engine.request -> engine.__init__ -> engine.engine ->
scheduler). Import it as ``from arbi_serve.engine.engine import Engine``.

Re-exports are wired through PEP 562 ``__getattr__`` so importing this
package does NOT eagerly pull ``arbi_serve.engine.batch`` (which
imports torch). Light-weight callers like ``dump_openapi`` (used by
``scripts/generate-client.sh`` in a torchless CI lane) reach
``arbi_serve.engine.request`` via fully-qualified import without
paying the torch tax.

Free functions, control-plane dataclasses, path counters, and the
(de)serialization helpers for :mod:`arbi_serve.engine.distributed_driver`.

Every name defined here is re-imported into ``distributed_driver.py``,
so the public import path ``arbi_serve.engine.distributed_driver`` is
unchanged. This module holds NO driver state and imports nothing from
the driver / mixin modules (so there is no import cycle) — the mixins
import FROM it.

Slate indices that must have the MTP drafter seeded on THIS tick.

A row qualifies iff it opted into MTP (``mtp_k > 0``), emits a token this
step, AND is still a PREFILL row — i.e. this tick is the one finishing its
prompt. That is the only tick on which an SPMD row can be seeded: from the
next tick it is a decode row that the verify pass owns, and the verify
derivation reads ``mtp_next_drafts`` as a precondition (the drafter chain
inside the verify pass refreshes them from there on).

Pure function of ``(slate, mirror, emit_idx)`` — all of which are
bit-identical on every rank — so every rank computes the same list and
takes the same forward branch, keeping the collectives in lockstep.

This is a gate rather than an unconditional flag because a non-empty
result makes the caller ask ``forward`` for the per-token hidden, and
``return_hidden_state=True`` de-captures any graph without a
``hidden_out`` buffer — which is every plain-decode (S=1) graph, and
the seed helper has no decode-flavor replay rung at all (it replays
prefill rungs only). A warm decode row in this list therefore forces
a full eager whole-model forward on every rank every tick.

The ``is_prefill`` term is what makes that unreachable, and it must not be
inferred from the routing: whether a warm MTP row reaches this call site is
a property of the verify-dispatch predicate in the SPMD step body, and a
row that arrives here warm has, by construction, nothing to seed. Locked
down by ``tests/test_mtp_spmd_seed.py``.

Resolve the DFlash x SPMD step-loop routing.

Returns ``(effective_spmd_enabled, mode)`` where ``mode`` is one of:

  * ``"none"`` — no DFlash drafter; SPMD is left as ``spmd_enabled``;
  * ``"driver"`` — DFlash active on the REPLICATED build (default). SPMD is
    forced off. The rank-0-only DFlash draft model cannot be re-derived
    symmetrically across ranks (its replicated ``project_context`` GEMMs
    are not bit-identical across GPUs); DFlash runs on the rank-0-driver
    path, where rank 0 drafts and broadcasts the two vocab-parallel
    collectives to workers via the worker bridge.
  * ``"driver_sharded"`` — DFlash active on the TENSOR-SHARDED build
    (``ARBI_DFLASH_TP_SHARD`` at TP>1). Still the rank-0-driver path (SPMD
    stays OFF so the worker bridge is installed) — this is NOT the removed
    rank-symmetric SPMD step loop. Every rank builds its OWN drafter shard,
    and rank 0 drives the DFlash-specific shard dispatch: it broadcasts the
    whole draft step + the shard-KV observe so every worker runs its shard
    forward and the per-layer row-parallel all_reduces (plus the two
    vocab-parallel legs) rendezvous 1-to-1.

``effective_spmd_enabled`` is ``False`` for BOTH DFlash modes — the worker
bridge carries the DFlash collective dispatch either way.

Return whether boot or any declared resident uses a DFlash drafter.

The distributed step-loop routing is fixed before resident preparation,
so every declared pool member must participate in this boot-time decision.
Both convenience fields and general override keys are accepted by the
pool-member schema.

``True`` iff any emitting row carries a real top_k/top_p/min_p filter.

Consulted only when the filtered fast path is rolled back via
``ARBI_SAMPLER_FILTERED_GUMBEL=0``: filtered rows must then ride the
full-noise Gumbel path (the fused top-k/top-p kernel consumes a
materialized ``(B, V)`` noise tensor) to stay rank-symmetric. Greedy
rows sample nothing (argmax) so they never force the noise path.
Conservative on ``top_k`` (any positive cap counts) — a false
positive only costs the slower path, never correctness.

One per-step broadcast payload from rank 0 → ranks 1..N-1.

Carries the materialized batch as RAW CPU buffers — one numpy
ndarray per scheduling tensor (and a CPU torch tensor for the
vision embeds). Workers rebuild the on-device :class:`ScheduledBatch`
by moving each buffer to the local device, skipping the page-table
walk entirely (rank 0 owns slot allocation; the slot indices are
simply shipped).

Why ndarrays, not Python lists. ``Tensor.tolist()`` builds a nested
Python int object per element — for ``block_table`` (B × pages) that
is millions of boxed ints per decode step, and each ``.tolist()`` is
a separate device→host blocking sync. A raw ndarray pickles as ONE
out-of-band buffer (protocol-5 ``buffer_callback`` in the shm queue),
and rank 0 stages every tensor with a SINGLE batched non-blocking
device→host copy + one sync. The ndarray dtype carries the field's
exact dtype, so the worker rebuild is bit-identical.

Out-of-band control-plane message broadcast on the TP group.

Two scopes:

  - **Outer-loop ops** broadcast on every tick:

    - ``"slate"``: rank 0 has work this tick. Workers enter the
      worker-bridge dispatch loop and consume per-slate ops
      (batch / drafter / seed) until rank 0 broadcasts
      ``"tick_done"`` via the bridge.
    - ``"swap"``: payload is the new ``"kind:name"`` spec string
      for an attention-backend swap. Long-lived.
    - ``"config_override"``: payload is the sparse overrides dict
      from ``POST /v1/admin/config_override``; every rank applies
      the same delta locally so the variant route (noop / live
      overlay / switch) is computed identically in lockstep.
    - ``"profile"``: payload is ``"start"`` / ``"stop"``; every rank
      flips its own live torch profiler so each dumps a per-rank
      Kineto trace (the ``/start_profile`` + ``/stop_profile`` path).
      Unlike swap / config_override it is NOT drain-gated (profiling
      wants live traffic), so in legacy driver mode it can land
      MID-SLATE between the bridge's per-slate tuples; the worker
      dispatch defers it to the tick boundary
      (:meth:`arbi_serve.distributed.worker_bridge.WorkerRankBridge.run_worker_dispatch`).
    - ``"drift_seam"``: payload is a ``{"action": ...}`` dict for
      TP>1 in-process TKV calibration. Rank 0 brackets each
      drift-measure's forwards with an ``install`` (before) and
      ``clear`` (after) on THIS FIFO ring, plus a one-time
      ``load_bundle`` at drift-stage start, so every worker arms
      its own kv-head-shard seam for exactly the candidate's
      forwards. Outer-loop, applied between ticks (see
      :func:`arbi_serve.calibration.drift_seam.apply_drift_seam`).
    - ``"shutdown"``: payload is ``None``; workers exit.
    - ``"noop"``: rank 0 has no work this tick; workers wait.

  - **Per-slate ops** broadcast through the worker bridge
    (``("batch" | "drafter" | "seed" | "tick_done", payload)``),
    not :class:`ControlOp`. See
    :mod:`arbi_serve.distributed.worker_bridge`.

Finish every request in a failed step's slate with ``finish_reason='error'``.

The shared terminal-fail loop for the driver's four step-failure
handlers: the rank-0 step ``except``, the SPMD cross-rank
fail-agreement, and the SPMD rank-0 commit + verify-stream ``except``
blocks. For each request it frees the KV via :meth:`Scheduler.finished`,
marks it FINISHED, signals the streaming waiter, drops it from the live
request map, and emits finish metrics — so a connected client is
released instead of hanging on a step that errored.

Rank-0 only: only rank 0 owns the real :class:`Request` objects (LoRA
is refused at TP>1, so there is no per-request LoRA ref to release
here). The terminal record rides the typed output pipeline — one
``FinishOut`` per request, flushed immediately since the SPMD loop
has no later flush site for this error path.

Move a set of device tensors to CPU with a SINGLE sync.

Issues a non-blocking ``.to("cpu")`` for every tensor (each queues
an async copy on the current stream), then synchronizes ONCE so all
copies land before any host read. CPU-resident tensors pass through
untouched (no sync needed); the sync only fires when at least one
source tensor is on CUDA.

Convert a rank-0 :class:`ScheduledBatch` into a raw-buffer
:class:`StepPlan` workers can reconstruct bit-identically.

Stages every scheduling tensor to CPU with ONE batched
device→host copy (:func:`_batched_d2h`), then hands each off as a
numpy view. No per-element Python int materialization:
``block_table`` ships as one ndarray buffer, not a B × pages nest
of boxed ints.

Module-level because :class:`RankZeroBridge` (in
:mod:`arbi_serve.distributed.worker_bridge`) needs to call this
without holding a driver instance.

Fast path (persistent cudagraph hot path): when ``batch`` carries a
:class:`HostBatchMirror`, the seven core scheduling tensors are
already host-resident — ``_build_batch`` wrote them into the pinned
``h_*`` ring this step (and patched ``h_input_ids`` to match the
post-``gpu_overwrites`` device buffer). Read them straight from the
mirror: NO device→host copy of tensors the host already holds, and
NO per-step ``torch.cuda.synchronize()`` drain on the rank-0
critical path. Only the optional carriers (mm / mrope /
state_indices) — which have no ``h_*`` twin — still ride a batched
D2H, and that D2H syncs ONLY when one of them is actually present
and on CUDA (the K=0 text-only step has none → zero sync).

Slow path (fresh-alloc / CPU-stub / test): ``host_mirror is None``,
so all tensors go through the batched D2H exactly as before.

Rebuild a device tensor from a raw ndarray buffer.

``torch.from_numpy`` is a zero-copy view of the ndarray. The ndarray
already carries the wire dtype (set by :func:`_serialize_batch_to_plan`),
so ``dtype`` here is a no-op assertion of the target — kept explicit so
the worker tensor is bit-identical to the rank-0 source.

H2D via PINNED staging + ``non_blocking=True``, NOT a direct pageable
``.to(device)`` — the latter is a synchronous pageable H2D (the CUDA
runtime serializes it as ``cudaMemcpyAsync`` + ``cudaStreamSynchronize``),
which stalls the host behind everything already enqueued on the stream.
The SPMD TP decode path. ``pin_memory()`` copies the host bytes into a
block owned by the CUDA caching host allocator, which defers reuse
until the recorded copy event completes, so the async copy is safe
even though the pinned tensor is a temporary; the source ndarray is
free to be reused immediately (the pinned copy already happened, on
the host). Values are byte-identical (bit-neutral) — only the copy
is async.

The worker's ``state_indices_long``, matching rank 0's field presence.

``state_indices_long`` is the int64 alias of the recurrent slab-row
mapping. Rank 0 refills it in lockstep with the int32 mapping out of a
persistent buffer so the recurrent block consumes a stable ``data_ptr``;
a worker that left it ``None`` fell back to a per-call ``.to(int64)``
AND — because the compiled recurrent layer guards on this field's TYPE —
ran a different compiled specialization of the same logical layer than
rank 0, on every step. A boot warm is rank-symmetric by construction, so
it can never cover a specialization only one rank reaches.

Materialized from the worker's OWN persistent buffer: the value is
recoverable from ``state_indices``, so shipping it would cost wire bytes
and still hand the layer a fresh per-step tensor.

Falls back to ``None`` when the buffers cannot hold the batch — the same
state rank 0 publishes when its own persistent path refuses.

Arm the scheduler's admission-time activation gate from the boot profile.

The engine-side half of :mod:`arbi_serve.scheduler.activation_budget`: it reads
the boot activation profile and the serving floor, calibrates the step model,
hands it to the scheduler, and — the part that matters most on a ceiling-bound
card — states in the boot log whether the WIDEST slate admission can build fits
the bytes the floor reserved for one step.

That statement is the pair the KV budget has never had. The floor reserves for a
step; admission decides which steps exist; nothing has connected the two. When
the widest admissible slate costs more than the reserve, the excess is not a
theoretical worst case — it is a slate the scheduler will build, whose forward
then OOMs after the layout freeze. Naming the excess, and which term carries it,
is what lets the floor follow the real peak instead of a guess.

THE BUDGET IS RE-READ, not just armed. The boot reading is taken once and
everything the process maps afterwards comes out of it — ``scratch.forward_arena``
re-growing past what the post-capture grow accounted for, a cubin the driver
loads on a kernel no boot phase dispatched — so the gate went on enforcing a
number for a card that no longer existed and the forward met the difference as
an OOM. :func:`rearm_step_budget` takes the SAME measurement again on the two
edges where the card contradicts the gate (a serving OOM, and the first step
that succeeds after one) and re-arms every scheduler with it. Only
``budget_bytes`` moves: the slopes describe the model, not the card.
:func:`minimum_step_bytes` is the other half of that pair — the cost of the
smallest slate the gate can build, so a budget below it is the statement that
narrowing has run out. See :mod:`arbi_serve.engine.memory_pressure`.

A DEFICIT IS NOT AN EMPTY ROW, and the boot refuses on it. Free VRAM below the
reserve the floor holds for consumers OTHER than a step means those bytes are
already gone before admission has priced one slate. That reading used to be
clamped at zero and then restored up to the grow's own step reserve, which
handed the gate a budget better than the card and printed the deficit as
``0 B``; the first heavy request met the difference as a step OOM and every
request behind it was shed. :func:`_step_free_bytes` now carries the sign
through and bounds the restore by the share it repays — the STEP's reserve, the
only one whose bytes the post-freeze phase can have left as this process's
mapped segments — so a raided non-step share stays raided.
:func:`measure_step_budget_bytes` publishes what survives, and
:func:`arm_activation_admission` refuses on a deficit: no chunk narrowing or
deferred row returns bytes to a reserve the step never held. It refuses equally
on a budget that does not outlast :func:`minimum_step_bytes`, the point where
narrowing has nothing left to narrow to; a budget merely short of the WIDEST
slate still narrows, which is the gate working.
``activation_step_reserve_overcommitted`` and
``activation_step_budget_below_minimum`` are the counters that say which fired.
Serving-time re-arms never refuse: there the same reading is what the
memory-pressure ladder narrows and sheds against.

A READING AND AN ABSENCE ARE DIFFERENT ANSWERS, and this module keeps them
apart at every seam it owns. :func:`measure_step_budget_bytes` returns ``None``
when there is no reading to take and an integer — ``0`` included — when there
is; ``StepActivationModel.budget_bytes`` carries the same distinction into the
gate; and :func:`rearm_step_budget` returns ``None`` only for the one fact its
caller can do nothing about, that no scheduler carries an armed model. A zero
reading is the state where a step may spend nothing, which is the state where
admission most needs to narrow and the front door most needs to shut — so it
cannot share a value with the state where the gate is inert.

THE EDGES ARE NOT THE ONLY SIGNALS. :func:`rearm_step_budget` above re-reads on
the two edges where the card has ALREADY contradicted the gate. Three more
events say the card is about to, and each of them already announces itself:

* ``scratch.forward_arena`` re-grows past what the post-capture grow accounted
  for. It is exempt from the Phase-2 live-size cap on the stated grounds that
  the SERVING FLOOR must bound it instead, and
  :func:`~arbi_serve.engine.arena_watch._report_overage` exists for when that
  floor turns out to be short: it computes the exact overage and records it for
  the NEXT boot.
* driver residency grows past what ``transient.serving_step.driver_growth``
  holds — a cubin the driver materializes on a kernel no boot phase dispatched,
  a re-grown kernel-stack pool, a graph instantiated post-boot. Those bytes are
  in no allocator counter at all, which makes this the overage most likely to
  move the budget.
* a serving-time JIT / dynamo compile fires ``jit_compile_serving``, leaving
  both a resident cubin and whatever its autotune sweep allocated.

All three used to end at a log line and a cross-boot record, and the gate served
the rest of the process believing a number the card had stopped honouring.
:func:`request_rearm` / :func:`service_pending_rearm` close that: each signal
latches a request and the next served step re-runs
:func:`measure_step_budget_bytes` against it — the same measurement both paths
share. That turns a permanent silent shortfall into narrower chunks and deferred
rows BEFORE a request pays for it, which is what the gate is for.

ONE READING, TWO CALLERS. Both re-arm paths take the budget from
:func:`measure_step_budget_bytes` and re-arm through the same ``replace``,
because it is one quantity and a second expression for it would mean the
comparison that decides whether to narrow was between two different numbers.
They differ only in what their own call site can do: :func:`rearm_step_budget`
is reached from the OOM edge at world size 1 and may take the all-rank
collective; these three signals are rank-LOCAL — one rank's arena grows, one
rank's kernel compiles — and an all-reduce entered by the rank that saw the
signal and not by its peers is a hang, not an agreement, so that path passes
``collective=False`` and clamps to the boot budget, which keeps every re-armed
value at or below a number every rank agreed to.

Nothing here is periodic, and nothing decays. A timer would put a driver call on
the step path to ask a question whose answer only changes when one of these
events happens — and each of them already says so. The two standing conditions
raise a request only when their overage GROWS
(:func:`~arbi_serve.engine.arena_watch._request_rearm_on_rising`), so a
persistent overage under a metrics scraper cannot turn into a cadence.

Calibrate the step-activation model for this engine, or ``None``.

``None`` whenever the inputs are not there to calibrate honestly — no
activation profile, fewer than two profiled token widths — and the gate
then stays inert rather than enforcing a number nothing measured.

Bind ``per_query_token_bytes`` below by the ``lm_head`` epilogue.

The query slope is fitted from :data:`~arbi_serve.runtime.
activation_profile.SHAPE_VERIFY_FORWARD`, and that probe is
``model.forward`` — whose head runs on the PER-SEQUENCE last-token gather,
one row per sequence, exactly as a plain decode step's does. The served
verify step does not: it takes the forward's hidden state back and runs
``logits_from_hidden`` over every drafted position
(:mod:`arbi_serve.spec_decode.mtp_verify_offload`). So the fitted slope
cannot contain one byte of the head's per-position cost, and a slope that
does not contain it lets admission seat a slate whose epilogue has nowhere
to allocate — which is an allocator failure standing in for a scheduling
decision.

Binds only DOWNWARD, the same rule the serving floor's verify tail follows:
a fitted slope ABOVE the bound already covers the head and is left alone.
Below it, the bound stands in, and the gate then narrows the chunk or defers
the row (:meth:`~arbi_serve.scheduler.scheduler.Scheduler._activation_admits`)
instead of the step OOMing.

A head that cannot be read leaves the model exactly as calibrated.

Bytes the widest slate THIS engine's admission limits allow.

THE ONLY PLACE THE WIDEST SLATE IS EVALUATED. Two consumers ask for it and
they have to get the same answer or the reserve and the gate are about
different steps: the serving floor asks
(:func:`~arbi_serve.engine.inprocess_capture.forward_arena_step_bound_bytes`)
so its activation rows can cover the widest step, and
:func:`arm_activation_admission` asks so the boot line can say whether they
do. They previously evaluated the same model at two different corners — one
at ``max_batched_tokens``, the other at
:func:`~arbi_serve.runtime.activation_profile.reachable_step_tokens` — so
the floor could be sized for one step and the gate report on another.

Every corner is a limit the scheduler enforces STRUCTURALLY, before any
byte accounting, and the model is monotone non-decreasing in all of them,
so no admissible slate evaluates above this. ``reachable_step_tokens`` is
the token corner rather than the raw budget because a step is built from at
most ``max_batch`` rows and every prefill row is capped at
``chunk_prefill``: a budget above their product names a step the scheduler
cannot assemble.

Extra decode-row query tokens the widest slate admission can build.

``max_batch * K``: every row opted into speculation at this boot's resolved
draft depth, charged the ``K`` drafts its verify step carries beside the
committed tail the row's own per-row term already prices. Zero when the
engine does not speculate, which leaves the widest-slate number exactly
what it was.

Resolved through :func:`row_query_tokens`, which reads
:func:`~arbi_serve.runtime.activation_profile.verify_forward_geometry` — the
same predicate that decides whether the shape this term is calibrated from
gets measured at all. One read of the depth for the widest slate and the
smallest one alike: a second reading of ``n_draft`` here could charge for a
depth the profile never ran, or disagree with the floor the OOM ladder
compares against.

Refuse a boot whose non-step reserve is not free at the moment it arms.

``serving_nonforward_reserve_bytes`` is the part of the serving floor that
must stay DRIVER-free while a step runs and that a step may not spend: the
additive spec-decode reserve (the stochastic verify tail, the DFlash draft
transient, and the verify slate's ``lm_head`` epilogue — a fresh
``max_batch x (K + 1)`` by vocab block on the default allocator every verify
step, reserved because no pool holds it), the media and TTS holes, the
separate vocoder process's physical hole, the ``gpu_memory_utilization``
headroom. A free reading below it says those bytes are already gone — spent
by something that allocated after the last residency bracket the floor was
measured from — and the reserve is breached before admission has priced a
single slate. The first step to reach for one of those reserved
allocations then meets an OOM the ledger said could not happen.

NO ADMISSION LEVER REACHES IT. Narrowing a prefill chunk and deferring rows
return bytes to the STEP's own share; the deficit is in a reserve the step
never held, and the consumers it belongs to (another process, a hard
utilization cap) cannot be served out of this process's allocator at all. So
there is no narrower configuration to fall back to and the honest answer is
the one the KV floor and the post-freeze gate already give: refuse the boot
rather than serve a layout whose first heavy request meets the shortfall as
a step OOM, which sheds every request behind it.

Reads the SIGNED row admission published
(:attr:`~arbi_serve.engine.boot_state.BootState.
serving_step_free_restored_bytes`) rather than re-deriving it, so the number
the refusal names is the number the gate was about to arm against.
``activation_step_reserve_overcommitted`` is the counter that says this
fired.

RUNS BEFORE THE ARMED BOOT'S TERMS LINE, which is why it states the terms
itself: that line reports what a scheduler took, and on this path nothing
does. One format string carries both the log and the exception so the two
cannot come to say different things about the same reading.

Refuse a boot whose budget does not outlast admission's last lever.

:func:`minimum_step_bytes` is the cost of the smallest slate the gate can
build — one row, no prefill tokens, no padded gather, and on a speculating
boot the ``K`` drafts that row's verify step carries. Admission narrows
towards it and stops there: the chunk cap floors an empty slate at one token
and ``_activation_admits`` seats a solo row that does not fit rather than
starving it. So a budget at or below that number is the statement that every
narrowing lever is already at its stop before the first request arrives, and
the slate the scheduler is obliged to build costs more than a step may
spend.

NARROWING FIRST, REFUSAL ONLY WHERE NARROWING RUNS OUT. A budget under the
WIDEST slate is not this: the gate answers it by narrowing the chunk and
deferring rows, which is the mechanism working, and the line above reports
it. This refuses only the case with nothing left to narrow to, where serving
means handing every request a step OOM that sheds the queue behind it.

``activation_step_budget_below_minimum`` is the counter that says this
fired. A minimum of zero — a model whose smallest slate costs nothing —
cannot be missed by any budget and is left alone.

Arm every scheduler, then log the fit against the reserve.

Best-effort about the INPUTS: an engine without a profile or a resolvable
floor keeps the page-only admission it has always had, and says so at INFO.
The gate is a tightening, never a prerequisite for serving. A calibrated
model with no scheduler to receive it is a different thing — the inputs were
all there and the gate is unreachable anyway — and that is an ERROR.

THE LOG LINE REPORTS WHAT WAS ARMED, which is why the arming runs first.
A line describing a calibrated model says nothing about whether any
scheduler received it, and an engine whose gate is unreachable reads
identically to one whose gate is enforcing — the boot's only statement
about the gate has to be a statement about the schedulers holding it, or a
silent miss is indistinguishable from a healthy arm until an OOM asks.

Bytes the SMALLEST slate admission can build costs.

One row, no prefill tokens, no padded gather — the shape the scheduler
falls back to when every narrowing lever has been pulled
(:meth:`~arbi_serve.scheduler.scheduler.Scheduler._activation_chunk_cap`
floors an empty slate at one token, and ``_activation_admits`` admits a solo
row that does not fit rather than starving it). It is therefore the floor of
what the gate can do, and a budget below it is the statement that no slate
exists which fits the card. Read from the model, so it moves with the same
calibration the gate enforces rather than standing a second number beside
it.

``query_tokens`` IS THAT ROW'S DRAFTS, and on a speculating boot they are
not optional. The chunk cap narrows PREFILL tokens and is called on prefill
rows only; a decode row goes through
:meth:`~arbi_serve.scheduler.scheduler.Scheduler._activation_admits` at the
full ``K`` its :meth:`~arbi_serve.scheduler.scheduler.Scheduler.
_query_charge` reports, with no lever between "charge ``K``" and "admit it
anyway because one row is the minimum unit of progress". So on a boot that
speculates the smallest slate the gate can be handed is a ``K + 1``-token
verify row, and pricing it at one token would tell the memory-pressure
ladder there is narrowing left where there is none — the door would stay
open for another OOM, on exactly the card the ladder exists for.

Zero — the default, and every non-speculating boot — leaves this the number
it has always been. The spec valve can still zero ``K`` on the emitted
slate, which only makes the executed step cheaper than the one priced here;
it fires on decode-batch WIDTH, so an engine drained to one row under
pressure is precisely where it leaves ``K`` standing.

Extra query tokens ONE decode row's verify step runs beyond its tail.

``K``. The per-row half of :func:`max_query_tokens`, off the same predicate,
so the widest slate and the smallest one cannot disagree about whether this
boot speculates or how deep.

Re-measure what a step may spend and re-arm every scheduler with it.

Returns ``(budget_bytes, minimum_step_bytes)``, or ``None`` on the ONE fact
the caller cannot act on: no scheduler on this engine carries an armed
activation model. There is then nothing calibrated to narrow and no measured
minimum to compare a budget against, and saying so is better than inventing
either.

EVERY OTHER OUTCOME RETURNS A PAIR, including the one an exhausted card
produces. A budget of zero is the reading that a step may spend nothing —
the state in which the ladder's remaining rungs (shut the front door, let
the in-flight rows drain, escalate if an OOM lands before one succeeds) are
the only reaction left. Folding it into the ``None`` return would disarm the
ladder in exactly the condition the ladder exists for, and would report it
as "no armed model" to an engine whose boot armed one.

A card that cannot be READ (no CUDA, no boot state) is the third case, and
it re-arms nothing: the gate keeps the budget it holds, and that budget —
the number admission is enforcing, not an invented one — is what the pair
reports. Re-arming with a fabricated value would replace a measurement with
a guess at the moment the measurement is missing.

THE SAME MEASUREMENT THE BOOT TOOK, taken again. ``budget_bytes`` is the
only field that moves: the slopes and the constant come from the boot
activation profile and describe the model, not the card, so re-calibrating
them here would re-fit a shape nothing re-profiled. What changed is how much
the card has left, which is exactly what :func:`measure_step_budget_bytes`
reads.

WHY IT IS RE-READ AT ALL. The boot reading is taken once, at the moment
serving begins. Everything the process maps afterwards comes out of it —
``scratch.forward_arena`` re-growing past what the post-capture grow
accounted for, a cubin the driver loads on a kernel no boot phase
dispatched — and none of that moved the number admission was enforcing. So
the gate went on admitting slates priced against a card that no longer
existed, and the forward discovered the difference as an OOM. Re-reading on
the edges where the card contradicts the gate is what keeps the two the same
quantity.

OFF THE STEP PATH. ``mem_get_info`` is a driver call this repo keeps off
the hot path, and the reading is collapsed to the all-rank minimum, which is
a collective. Callers reach this only from the OOM edge and the first
success after one, and only at world size 1 (see
:mod:`arbi_serve.engine.memory_pressure` for why the multi-rank ladder stops
before here).

Ask the next served step to re-measure the admission gate's budget.

THE SIGNAL AND THE MEASUREMENT ARE DELIBERATELY SPLIT. Every caller here
fires from somewhere the measurement cannot honestly be taken: the arena and
driver overage reports also run from the metrics scrape thread, a JIT
compile fires from INSIDE the forward that triggered it, and an OOM fires
from the run loop's exception path. A ``mem_get_info`` taken mid-forward
reads the card with that step's whole transient live set on it and would cut
the budget to a number no quiescent step has to live under; one taken off
the engine thread races the allocator. Latching the request and servicing it
at the post-step seam takes every reading at the same point in the step
cycle the boot reading was taken at — after a forward, with the transients
freed — which is the only way two readings are comparable at all.

Free of engine state on purpose: :mod:`arbi_serve.jit_detector` has no
engine handle, and giving it one to raise this signal would be a worse
coupling than a module-level latch. The reasons ACCUMULATE (a set, bounded
by the number of distinct signal sites) so one service call can name every
signal that stood behind it.

Never raises: a signal that breaks the path it observes is worse than the
shortfall it reports.

Whether a signal has asked for a re-measurement (a set truth test).

Read on the post-step seam, so it is the one thing here that a served step
pays for. Deliberately not a function of the engine: the whole point of the
latch is that its writers have no engine.

Record on the boot state the budget the schedulers actually carry.

``measure_step_budget_bytes`` writes its RAW reading here, which is correct
for the boot seam that arms straight off it and wrong for every path that
clamps, declines or fails after taking it. One helper so the three exits
that have to undo it cannot drift into two of them doing it.

Re-measure the step budget and re-arm every scheduler, or ``None``.

``None`` whenever nothing changed — no request outstanding, no reading to
take, the gate was never armed, or the card still holds what the gate
believes. That last case is a RESULT, not a failure: the signals below say
"something permanent may have happened", and only the measurement says
whether it cost the step anything. The pair of counters is what tells the
two apart from outside — ``activation_gate_rearm_requested`` says the
mechanism ran, ``activation_gate_rearmed`` says it found something — so a
zero on the second can never be read as "did not run".

THE SIBLING OF :func:`rearm_step_budget`, and the two now differ in exactly
two things: which callers reach them, and the clamp that follows from it.
They take the SAME reading, from :func:`measure_step_budget_bytes`, because
it is the same quantity — a second expression for it would mean this
function's decision to narrow was a comparison between two different
numbers, which is the defect class this module exists to keep out. That one
is the OOM edge's, runs at world size 1, and may therefore take the all-rank
collective; this one is reached from three rank-local signals during
ordinary serving, so it passes ``collective=False`` and clamps instead.
Both end at the same place: every scheduler re-armed through the same
``replace`` with a model whose only moved field is ``budget_bytes``.

CLAMPED TO THE BOOT BUDGET, never to the last re-armed one. The boot reading
is the all-rank minimum; this one is rank-local, because the signals that
raise it are rank-local. Clamping to the boot value keeps every re-armed
budget at or below a number every rank agreed to, so this can only ever be
TIGHTER than the gate the boot armed, which is the direction the whole
mechanism exists to move in. It does not see a peer's own new shortfall —
but neither did the gate before, so that is the state this leaves unchanged
rather than one it introduces.

Clamping to the BOOT value rather than the last re-armed one is also what
keeps a transient reading from ratcheting: a budget cut by one bad moment
recovers on the next signal that finds the card whole again, instead of
holding the narrow chunk for the life of the process.

A MEASURED ZERO RE-ARMS, and the clamp cannot turn the gate off. The card
reporting that a step may spend nothing is the state where narrowing is the
only reaction left, so it takes the same path every other reading takes:
the model is re-armed at zero, the chunk cap goes to zero tokens, and the
scheduler's empty-slate floor still advances a row one token per step. Only
an UNREADABLE card (``None``) leaves the gate on the budget it holds. This
is the same distinction :func:`measure_step_budget_bytes` and
``StepActivationModel.budget_bytes`` draw, at the one seam that would
otherwise re-fold them: ``min(ceiling, 0)`` is ``0``, and a ``0`` treated as
"nothing to do here" would disarm the reaction on the fullest cards.

Bytes one step's forward may allocate, or ``None`` when unreadable.

THE ONLY PLACE THIS QUANTITY IS COMPUTED. Three callers ask it — the boot
seam, :func:`rearm_step_budget` on the OOM edge, and
:func:`service_pending_rearm` on the signals that fire before one — and they
ask the same question about the same card. A second implementation of it
would be two numbers that only look like one, which is how the sentinel this
module now keeps apart got in: the boot arms a budget from one expression
and a re-measurement compares against it with another, so the comparison
that decides whether to narrow is between different quantities. One
expression, one credit source, one set of reserves.

MEASURED, on this boot, at the moment it is asked for — the boot seam asks
at the moment serving begins, and :func:`rearm_step_budget` asks again on
the edges where a serving OOM contradicts that first reading. The driver's
own free reading, minus the part of the serving floor that belongs to
consumers other than the step (:attr:`~arbi_serve.engine.boot_state.BootState.
serving_nonforward_reserve_bytes`), plus the physical ``scratch.
forward_arena`` already holds and a step can therefore spend without
touching free VRAM (:attr:`~arbi_serve.engine.boot_state.BootState.
serving_arena_resident_credit_bytes`). Nothing here is predicted, cached or
carried from another boot, because nothing else can be: what a step may
spend is whatever the card still holds once every pool is mapped, and only
the card can say what that is.

THE ARENA TERM IS NOT A CUSHION, it is the other half of a reading the
driver cannot give. ``mem_get_info`` reports what is unmapped; a private
cuMem pool's reserve is mapped and so invisible to it, while the model this
budget gates is calibrated on per-shape peaks that include the pool's
allocations. Free VRAM alone therefore prices a step's whole live set
against only the part of the card that can pay for some of it, and the gate
answers by narrowing the prefill chunk — the arena's residency, charged to
the request that cannot use it. The floor and this budget now split the pool
exactly once between them: the floor stops holding it, and this credits it.

Read HERE rather than from the floor's own activation rows, which is where a
gate drifts from the physical truth. The floor is computed at the
post-capture grow and the grow leaves exactly that many bytes free — but the
boot's own serve-kernel warmup and readiness forward then draw from them,
and so does every driver allocation landing after the last residency
bracket. A budget cut from the floor promises a step bytes the card no
longer has; a budget cut from the reading taken after all of that has
happened promises exactly what is there.

The reading is collapsed to the ALL-RANK MINIMUM. Free VRAM diverges across
ranks structurally (rank-0-only pools, per-rank driver residency), the
scheduler that spends this budget runs on rank 0 alone, and the slate it
builds is executed by EVERY rank — so the binding card is the tightest one,
not the scheduler's own. One int64 all-reduce, at boot, on a path every rank
reaches.

``collective=False`` is the caller stating that its own call site is NOT one
every rank reaches, which is a fact about the CALLER and never about the
quantity — the arithmetic below is identical either way. The signals
:func:`service_pending_rearm` services are rank-local (one rank's arena
grows, one rank's kernel compiles), and an all-reduce entered by the rank
that saw the signal and not by its peers is a hang, not an agreement. That
caller clamps its rank-local reading to the boot budget instead, which is
where the all-rank agreement already lives, so its result is still bounded
by a number every rank consented to.

``None`` when there is no reading to take — no boot state, or a driver that
will not answer — which leaves the gate unarmed rather than enforcing a
number nothing measured.

ZERO IS A READING. Free VRAM at or below the reserve that belongs to
consumers other than the step is the card stating that a step may spend
nothing, and it is the state an OOM leaves behind almost by definition: an
allocation failed because the card was full. Returning it as the same value
that means "unreadable" would make every consumer treat the fullest cards as
the ones with nothing to say.

FREE VRAM a step may spend, holding the KV sizing to its own promise.

The grow reserved ``serving_step_free_reserve_bytes`` of free VRAM for one
step and left exactly that much. Then the boot's OWN post-freeze phase — the
serve-kernel warmup, the JIT replay, the readiness forward — ran, and those
are the classes a served step reaches: they took the reserve out of free
VRAM and put it back as caching-allocator segments they no longer hold live.
That phase is measurable at both shipped shapes: ``unpooled.
torch_default_pool`` grows by roughly the reserve between the freeze and the
arm, and ``free - the non-step reserve`` lands at the zero mark — so without
the restore the gate would run on the arena's residency alone and report a
card short of a slate it can serve. The boot's own ledger reports the two
readings for the boot in hand; the terms line
(:func:`arm_activation_admission`) states them side by side, signed.

Those bytes did not leave. They are mapped, they are not live, and the next
allocation of a size the warmup already realized is served from them without
touching the driver — which is the same fact the forward arena's credit
rests on, for the same reason, and this is the rest of the pool table
stating it.

So the row is restored to the promise and NEVER PAST IT. The cap is the
reserve the floor itself computed, so the gate can never admit against more
free VRAM than the KV sizing held back — the reserve and the gate stay one
quantity, and a card that genuinely never had the bytes is unaffected
because nothing is mapped to restore them from. A surplus above the reserve
(a grow that left more free than its floor asked for) still counts in full.

Reads the allocator at the SAME instant as the free reading it corrects, and
nets out the arena's own held-free bytes because
``serving_arena_resident_credit_bytes`` already spends those — each byte is
credited once, by exactly one of the two.

THE ROW IS SIGNED, AND THE RESTORE REPAYS ONE SHARE. ``free_row`` below zero
is the card stating that free VRAM no longer covers even the reserve that
belongs to consumers OTHER than the step: the boot spent past its own step
reserve and into theirs. Clamping it at zero first and restoring on top of
the clamp handed that state the full promise back — a card 60 MiB in the
hole and a card that had merely spent its row into mapped segments returned
the same number — so the sign is carried through instead.

What the restore may return is bounded by what it is repaying: every share
of the floor whose consumer allocates THROUGH THIS ALLOCATOR. That is the
step's own reserve and the in-process part of the non-step reserve — the
verify tail and its lm_head epilogue, the draft transient, the codec
scratches, the logprobs tile, the media and TTS holes — because a segment
the boot mapped and no longer holds live serves any of them without a
driver call, which is the same fact the arena's credit rests on.

Excluded is the share owed outside it: a separate vocoder PROCESS's
physical hole, which that process must find as real free VRAM, and the
driver's own growth. A segment in this allocator repays neither.

The reading is capped at that sum before it is added, which changes nothing
for a row at or above zero (the sum was already capped there) and stops an
externally-owed share from being repaid out of recoverable bytes. What
survives goes out signed to
:attr:`~arbi_serve.engine.boot_state.BootState.
serving_step_free_restored_bytes`, where the boot seam refuses on it.

Caching-allocator physical held but not live, OUTSIDE the forward arena.

``memory_reserved - memory_allocated`` is every segment the allocator holds
across every pool, private ones included; the arena's own share comes out
because admission already spends it as
``serving_arena_resident_credit_bytes``. 0 when either reading raises, which
leaves the free row exactly as the driver reported it.

Per-StateKind active-state builder + teardown.

:func:`build_active` constructs the production
:class:`MultiStatePool`, page-table, scheduler, metadata builders, and
per-layer attn ops; called both at boot (by
:func:`arbi_serve.engine.build.build`) and during per-StateKind
hot-swap. :func:`teardown_kind` drops a kind ahead of a rebuild.

:mod:`arbi_serve.engine.build` re-exports both for legacy import
paths.

Worst-case ``num_seqs`` a capture-time scratch pool must pre-size for.

Two boot-time consumers request rows against this pool:

  1. the scheduler's decode slate (tops out at ``resolved_max_batch``),
     plus the boot mixed-batch warmup's one extra row
     (:data:`_CAPTURE_SYNTHETIC_EXTRA_ROWS`); and
  2. :func:`arbi_serve.runtime.activation_profile.profile_engine`'s pure-
     prefill probe, which spreads ``max_batched_tokens`` q-tokens across
     ``prefill_probe_num_seqs(...)`` rows whenever one row cannot hold
     them — a row carries at most
     :func:`~arbi_serve.runtime.activation_profile.prefill_row_cap`
     (``min(max_context, chunk_prefill)``), so e.g. 8192 tokens at
     ``chunk_prefill=2048`` is 4 rows. At a low ``--max-batch`` that probe
     row count EXCEEDS ``resolved_max_batch``, so sizing the pool to
     ``max_batch + 1`` alone made the probe's ``TQBufferPool.ensure``
     overflow the pre-sized cap and abort the boot.

Pre-sizing to ``max(resolved_max_batch, probe_seqs) + extra`` covers BOTH
from ONE call: any boot-time ``ensure(B)`` with ``B`` from either consumer
can never exceed the cap, so the pool never needs a post-capture realloc
(which would invalidate captured-graph pointers) and the batch dimension
can't silently clamp below the resolved value. Both inputs are config
quantities (no magic literals); the probe row count is derived from the
same :func:`prefill_probe_num_seqs` the probe itself uses, so the two can
never diverge.

Worst-case flat ``num_tokens`` a capture-time scratch pool must hold.

The token-dim sibling of :func:`capture_pool_num_seqs`, and it exists
for the same reason: a boot-time consumer that requests more than the
pre-size gets a refusal, not a realloc — a realloc would invalidate
captured-graph pointers, so :meth:`TQBufferPool.ensure` fails the
capture instead.

Two consumers:

  1. the scheduler, bounded at ``max_batched_tokens``; and
  2. the MIXED capture sweep, whose graph runs a PADDED
     ``(max_batch - 1) + bucket_n`` flat tokens. That width is
     reachable only by the capture — a live mixed step still fits the
     scheduler's budget — so sizing from ``max_batched_tokens`` alone
     is exactly ``max_batch - 1`` tokens short whenever the top
     mixed bucket equals the budget, which is the shipped
     configuration. Measured on the served 27B at max_batch 4 /
     chunk 2048: ``ensure(num_tokens=2051)`` against a 2048 cap, and
     every mixed graph failed to capture.

``mixed_capture_flat_tokens`` returns 0 unless the sweep is armed, so
a boot without ``ARBI_MIXED_CAPTURE`` pre-sizes exactly as before.

Build the multi-tier prefix-cache store from runtime flags, or None.

Returns a :class:`PrefixTierStore` when ``ARBI_PREFIX_TIER`` is set,
else ``None`` (legacy drop-on-evict). The warm (RAM) tier is always
present; the cold (disk) tier is enabled only when
``ARBI_PREFIX_TIER_DISK_DIR`` is set — it encrypts blobs with an
ephemeral per-process key (documented confidentiality tradeoff). A
break-even threshold left at 0 in the env falls back to
``DEFAULT_{WARM,COLD}_BREAK_EVEN_TOKENS``, which are UNVALIDATED
placeholders, not a measured crossover — see
:mod:`arbi_serve.cache.prefix_tier_store` for what would replace them.

Clear per-call side-channel stashes so old pool slabs can free.

GDN / Mamba blocks live on ``eng.model`` (engine-lifetime) and stash
the per-call ``GdnLayerView`` / meta on ``self`` for the custom-op
real-impl. A stale stash pins the OLD pool's recurrent + conv slabs
AND the MTP snapshot buffers through a teardown→rebuild, doubling
the recurrent footprint. Safe to clear between forwards: every op
call re-publishes first.

The model runner's cached flush hooks are the same kind of stash — bound
methods OF the pool, held until the next step re-resolves them — so they
are released here too, by the same argument.

The engine's scheduler: one per attention-DP set, or just one.

At ``attn_dp_size == 1`` this is the plain
:class:`~arbi_serve.scheduler.scheduler.Scheduler` over the engine's
own page table — the exact pre-existing construction.

At ``attn_dp_size > 1`` each attention-DP set admits and owns its own
requests and their KV, so admission has to be decided per set. Rank 0
is the only rank with a scheduler under SPMD, so it runs one member
per set: member ``attn_dp_rank`` over the REAL page table (rank 0's
own KV) and every other member over a page-count mirror of that peer's
pool (:class:`~arbi_serve.cache.attn_dp_shadow_pool.AttnDpShadowPagedPool`).
A worker rank builds the plain scheduler and never schedules with it —
its rows arrive as the broadcast per-set delta.

The peers' page tables are FLAT regardless of what rank 0 runs: a
radix tree per set needs prefix-affinity routing (a request must land
on the set whose tree already holds its prefix) to be worth anything,
and the SPMD boot gate refuses radix at ``attn_dp_size > 1`` for that
reason.

Admission control / backpressure.

High-watermark gates on:
  - KV-pool occupancy (default 90% — refuse at/above)
  - Queue depth (refuse when more requests are waiting than the engine
    can plausibly serve in a couple of steps)

...and, ahead of both, the states in which the engine cannot serve AT ALL, so
the door does not admit into one: a poisoned CUDA context, an exhausted card,
live memory back-pressure (:mod:`arbi_serve.engine.memory_pressure`), a failed
admin swap, a released/sleeping engine, a drain, and the self-healing
infrastructure latch. Those are not thresholds and carry no ``limit`` — they
are the server saying what it already knows.

Orthogonal to per-token rate limits in :mod:`arbi_serve.server.auth`:
rate limiting protects the user's fair share, admission control
protects the engine from runaway queue growth. The gate fires HTTP 429
with ``Retry-After``.

The single admission bound: max requests in flight (running + queued).

``override > 0`` (an explicit ``--queue-depth-max`` / embed-fan-out
sizing) wins verbatim; otherwise derive ``INFLIGHT_PER_BATCH × max_batch``
(floored at :data:`INFLIGHT_FLOOR`). Both the HTTP semaphore and the
engine queue gate resolve their capacity through THIS function, so the
server has exactly one admission number with one derivation.

**This bound is deliberately WIDER than the number of sequences that can
RUN.** ``max_batch`` — and, on a hybrid GDN / Mamba / short-conv model, the
equally-sized recurrent slab-row pool — bounds the RUNNING set. This number
bounds the running set PLUS the queue behind it. Do not clamp it down to
either: that was the regression this docstring replaces (``a27a1ec00``,
2026-09-05). It sized the HTTP semaphore at ``max_batch``, so with one long
request resident a burst of 8 on the shipped ``--max-batch 8`` config lost
its last request to a spurious 429 at the door — never queued, never shown
as queued, just refused.

Row exhaustion is a SCHEDULING constraint, not an admission one.
``Scheduler.add`` does NOT raise on it (that premise died in ``e2ae43b1b``,
three months before the clamp asserted it): it DEFERS — rolls the request
back to a holding pen that owns no GPU state and re-attempts it between
forwards as running requests free rows. ``check_admission`` below refuses
interactive only at ``queue_depth_max``, which counts that pen, so a
SUSTAINED overload still backs off cleanly while a burst queues and drains.
Continuous batching IS concurrency > max_batch; a bound that forbids it is
not a bound, it is an outage.

The ENGINE-side queue bound derived from the connection bound.

For chat / completions one HTTP request is one engine request, so the
engine queue IS the connection bound. For embed / rerank one HTTP
request FANS OUT to one engine request per input, so the engine queue
must be deep enough to absorb that fan-out — a per-task allowance, NOT
a second knob and NOT a change to the connection bound.

Boot (``build_phases_kv``) and the live ``/v1/admin/config`` override
both resolve here, so the running bound can never follow a different
derivation than the one the engine booted with.

Outcome of :func:`check_admission`.

A refusal carries the GOVERNING THRESHOLD (``limit``) and the value
that crossed it (``observed``) so the HTTP layer can write a receipt
naming what refused and what would have to change to stop it — see
:func:`arbi_serve.server.shed.record_shed`. Both stay ``None`` for
refusals that are not threshold-driven (a drain, a latched fault).

Whether the self-healing infra latch is set AND has not decayed.

Reads the same two fields and applies the same
:data:`~arbi_serve.engine.infra_health.INFRA_UNHEALTHY_WINDOW_S` decay
``engine_unhealthy_reason`` does, so the front door and ``/health/ready``
can never disagree about whether the engine can run a forward. One
predicted-not-taken attribute read on the common path; the clock is read
only once a latch is actually set.

Whether anything the run loop will step is still in the scheduler.

Running OR waiting: both are rows ``schedule()`` builds a slate from, so
either will produce the successful step that clears a latch. An empty
scheduler produces no steps at all, which is why a gate that waits for one
must not be armed against it.

Decide whether to accept a fresh request right now.

Order of checks matters, and it is one order: every state in which the
engine CANNOT SERVE comes first, cheapest and stickiest first — a poisoned
context, a card that ran out of memory to narrow into, live memory
back-pressure, a failed admin swap, a released (sleeping) engine, a drain,
a live infrastructure latch — and only then the threshold gates that decide
how much work a serving engine should take: queue depth (cheap), then KV
occupancy.

THE FIRST GROUP USED TO BE THREE ENTRIES LONG (fatal fault, swap fault,
terminating), and everything missing from it was a way to accept a request
into an engine that was going to fail or ignore it: an OOM'd card kept
taking work, and so did a released one, whose requests then sat in the
scheduler with no error and no progress. Each new gate reads one attribute
on the served path.

Priority-aware backpressure (opportunistic batch).
=================================================
The ``priority`` of the arriving request shapes the queue-depth
gate so a large low-priority ``"batch"`` backlog cannot 429 an
interactive request:

  * INTERACTIVE requests count only OTHER interactive requests
    against ``queue_depth_max``. A queue full of batch work does
    not block interactive admission — interactive is admitted as if
    the batch backlog were not there.
  * BATCH requests are gated more tightly: refused as soon as ANY
    interactive request is still waiting (so batch never deepens
    the queue ahead of interactive demand), and counted against the
    FULL queue depth otherwise. This is the admission-side mirror
    of the scheduler's opportunistic admission: batch is accepted
    only with genuine spare room.

The default ``priority="interactive"`` keeps every existing call
site (embed / rerank / un-plumbed paths) on the original full-depth
gate, so behaviour is unchanged for all-interactive traffic.

Return the O(1) interactive counter ``sched.<attr>`` when concrete.

The Scheduler maintains ``_n_waiting_interactive`` /
``_n_running_interactive`` as plain ints. Read them in O(1) on the
hot path. When the attribute is absent or not a concrete ``int``
(e.g. a ``MagicMock`` scheduler in a unit test that fabricates the
``waiting`` / ``running`` lists directly), fall back to the exact
list scan so behaviour is identical to the pre-counter gate.

True iff a resident batch DECODE job can be parked for its row.

The suspended-job offload can only snapshot a DECODING job (stable
KV); a PREFILLING batch job's KV is incomplete and not offloadable.
On hybrid models the recurrent-state offload must also be enabled
(``_offload_recurrent``) — otherwise parking a hybrid job is refused
and the row can't be reclaimed this way, so admission must yield.
Requires the scheduler to have a suspended-store wired.

The per-step caches that hold a block in ``scratch.forward_arena``.

``scratch.forward_arena`` is per-step forward scratch, and the only way to hand
its physical back is to destroy the pool
(:func:`~arbi_serve.engine.inprocess_capture.release_idle_forward_arena` states
why: ``empty_cache`` skips a private ``MemPool``, and torch 2.12 has no
per-pool variant). A destroy needs the pool EMPTY — and emptiness is decided by
the whole pool, not by the block: **one** retained block forfeits every wholly-
free segment in it, whatever the block's size.

So a consumer that MEMOISES an allocation it made inside the arena blocks the
release even when the memo is a few hundred bytes and even when the data in it
is dead. That is not a leak and not a mis-placement: the allocation is per-step
forward scratch and the arena is where per-step forward scratch belongs. What
makes it a problem is only that the memo outlives the step, so the pool is never
empty at the one moment the engine could give it back.

THE DISTINCTION THIS MODULE DRAWS is between a memo and a tenant. A memo exists
to save the NEXT step an allocation; dropping it costs that allocation and
nothing else, and it is sound to drop precisely when the next step cannot hit it
anyway. A tenant is a buffer whose CONTENTS a later step (or a captured graph)
reads back; dropping one is unsound, and
:func:`~arbi_serve.engine.inprocess_capture._refuse_reclaim_live_arena` still
refuses on it, by name. Every entry below is a memo, and each entry says what
makes the next hit impossible at the seam that drops it — a memo whose key
could still match is a tenant in disguise and does not belong here.

WHERE THIS RUNS, and why it costs nothing. Only inside
``release_idle_forward_arena``: at boot before the freeze, and at the engine's
own idle quiesce point (no running row, no waiting row, both deferred pipelines
drained, no duplex deadline pending) once the pool has overrun what the
post-capture grow accounted for. Not on the step path, not on a timer. The
first step after that release re-maps what it needs out of a pool that was just
destroyed and recreated — so the allocations these drops give up are ones that
step was going to make regardless.

ADDING A CACHE. A new process- or engine-level memo of an arena allocation
REQUIRES an entry here, for the same reason
:data:`~arbi_serve.engine.member_scratch_retire.MEMBER_SCOPED_DEVICE_CACHES`
requires one: the failure is silent (the release simply stops firing and the
pool's whole idle tail stays resident) and it is invisible in any steady-state
gauge. ``tests/test_arena_step_caches.py`` asserts every entry still resolves,
so a rename fails the suite instead of quietly forfeiting the pool.

One memo of an arena allocation, and the seam that gives it up.

``drop`` returns the number of entries it released (0 when the memo was
already empty, which is the common case at boot). It must never raise: the
release is a repair, and a repair that can take the engine loop down is
worse than the bytes it recovers.

``why`` records what the memo caches and — the half that makes the drop
sound — why the next step cannot hit the entry being dropped.

Give up every memo of an arena allocation. Returns entries dropped.

Best-effort per entry: a drop that raises is logged and the rest still run,
because a memo this seam cannot reach costs the release, not the engine.

Owner attribution for tensors still live in a pool that must be empty.

A pool snapshot names bytes and shapes; a boot that has to *fix* a live
allocation needs the owner. This resolves each tensor to the first named root
that reaches it — a module global or an instance attribute — and renders a
dotted path such as ``pkg.mod._SCRATCH[(torch.bfloat16, 0)]`` or
``GatedSiLUMLP._fused_gate_up``.

Two passes, because neither alone is sufficient:

* top-down from ``sys.modules``, which names module globals and the
  containers hanging off them without consulting the collector at all;
* bottom-up via ``gc.get_referrers``, which names instance attributes and
  caches that live only inside a closure.

Diagnostic-only and deliberately bounded: it runs on a refusal path, never on
a serving path, and every step degrades to ``"<unattributed>"`` rather than
raising.

True (having logged) if ``what`` must not run on the calling thread.

The heap walks in this module and in :mod:`~arbi_serve.engine.memory_budget.pool_residency`
are a whole-heap ``gc.get_objects`` pass plus a referrer BFS — measured at
1.63s on a warm 27B boot — and they hold the GIL for the whole of it.

The GIL is PROCESS-wide, so moving the walk to a worker thread does not make
it safe: it stalls the engine loop from wherever it runs. A thread pool is
NOT an escape hatch, and a future reader who "fixes" a slow boot by pushing
this off-loop reintroduces exactly the stall this refuses. The only safe
placements are (a) boot, before anything is being served, or (b) never.

Keyed on the thread rather than on who calls it, so the property is a fact
about the function instead of a fact about today's call sites — which is the
difference between a guarantee and a convention.

``(key, value)`` pairs of an exact container, or ``None``.

Closure cells count: a cache held only by a closure is otherwise
unreachable, and that is a common shape for a lazily built scratch.

Paths from module namespaces down to any tensor in ``targets``.

Deliberately independent of the collector: ``gc.get_objects`` and
``gc.get_referrers`` do not see the permanent generation, so a process that
has called ``gc.freeze()`` hides its whole import graph — including exactly
the module-global caches this attribution most needs to name.

``(the namespace it owns, that namespace's name)``, or ``None``.

The inverse of the question the walk actually has. Climbing from a tensor
reaches its holder's NAMESPACE — an instance ``__dict__``, a function's
``__closure__`` — and a namespace does not carry its own name; the instance
or the function one hop above it does. This reads that name off a single
candidate referrer.

It replaces a ``_terminal_labels`` map built by walking every gc-tracked
object to label every namespace in the process, of which the walk then used
a handful. Measured on a warm 27B boot, that map was 336-373 ms — the
largest single term in the attribution once the storage-rep count fell.
Asking per referrer costs one ``__dict__`` read against the few objects a
climb actually reaches.

That ``__dict__`` read is also load-bearing, not just a lookup. CPython 3.12
keeps instance attributes in an inline values array and MATERIALISES the
``__dict__`` only when something asks for it — so before the read, the
referrer of an attribute value is the INSTANCE and no dict exists to be one.
The map this replaces got that materialisation as a side effect of touching
every object on the heap, which is why removing it silently changed the
referrer graph rather than just costing less. Reading it here restores the
same effect for the objects the climb reaches.

Name an object that holds a tensor but that the walk cannot descend.

Covers the C-level holders — an ``lru_cache`` wrapper, a frame, a bound
method, a ``torch`` autograd/graph node — where "who holds this" is the
whole answer and there is no key path below it.

Climb from ``t`` through its referrers, describing each link.

The last resort when no link in the chain is a named root: the chain of
container shapes usually identifies the cache on sight, and it always says
where to put a breakpoint. Queried one tensor at a time so each holder is
attributed to the right tensor, which a batched ``gc.get_referrers``
cannot tell you — affordable only because a pool that must be empty holds
few tensors by definition.

Of ``unresolved``, the ids some Python object refers to.

One ``gc.get_referrers`` heap scan for the whole batch, then a per-referrer
``gc.get_referents`` read — a pointer read, not a scan — to say WHICH
targets each referrer holds. ``get_referrers`` cannot be asked that
question per object without one heap scan per object, which is what forces
callers to cap how many tensors they probe; this pays a single scan and
lets the expensive per-tensor chain walk run only on tensors that can
actually yield a chain.

A target absent from the result has no Python referrer at all: it is
reachable only from C++.

Paths from each unresolved tensor up to the first instance attribute or
enclosing function.

``gc.get_referrers`` does the heap scan in C, which is the only affordable
way to do this with a model loaded.

Map ``id(tensor)`` to a dotted owner path for every tensor given.

``exclude`` carries the ids of the CALLER's own containers — the list it
gathered the tensors into, the per-storage records it folded them through.
Those refer to every tensor here, so without this they are found first and
the answer is the diagnostic naming its own bookkeeping.

A tensor nothing Python-side holds resolves to :data:`NO_GC_REFERRER`; one
that is held but that no pass could name within its hop bound resolves to
:data:`UNATTRIBUTED`. Both are verdicts a caller can act on — the first by
describing the allocation physically, the second by widening the walk.

Report which pass an expensive attribution spent its time in.

Every pass here is a whole-heap operation whose cost is set by the process
being described, not by the question asked — so "the walk took N seconds" is
not actionable without the split, and the split is the only thing that says
whether a targeted fix exists. Off unless ``ARBI_BOOT_PROFILE``.

Render one line per live storage group: size, dtype, shapes, owner.

``groups`` is the ``{storage_ptr: (nbytes, [tensors])}`` mapping a pool
snapshot walk produces. Sorted widest first so a truncated log still
carries the tenants worth moving.

The serving forward-arena watch: one owner for the pool's serving reading.

``scratch.forward_arena`` is a private cuMem pool. ``empty_cache`` skips it and
the Phase-2 freeze exempts it from the live-size cap, so whatever it maps while
serving is permanent physical taken out of the free VRAM the serving floor is
holding for the transient rows.

WHAT COVERS IT IS THE BOOT'S OWN REALIZATION. Every eager serving forward runs
inside this pool, and so do the boot's -- the activation profile probes at the
widest admissible shapes, the verify-width probe, the drafter pre-sweep. Boot
therefore MEASURES the pool's working set by realizing it, and keeps it: the
physical is resident before the KV grow reads free VRAM, so KV was never
offered it and a serving step re-uses it (:func:`~arbi_serve.engine.
inprocess_capture.forward_arena_regrow_plan`). Nothing here predicts the pool's
serving size, and nothing here carries a reading to another process -- a
reserve persisted across boots was tried and under-reserved a production boot
from a reading a process took after serving one short request.

SO WHAT THIS MODULE DOES IS CHECK THE REALIZATION HELD. A pool that grows past
what the grow accounted for is taking bytes the floor is holding for the other
transient rows, after the KV layout is frozen, and the only honest place to
learn that is while it happens.

WHY THIS IS AN ENGINE MODULE. The observation has to be taken because a step
was SERVED. It used to be taken only inside the metrics export's
reserved-bytes gauge callback, which is an observer and not an owner:

* with no OTEL provider installed, ``get_meter`` resolves the no-op provider
  and an observable gauge's callback is never invoked at all, so nothing was
  ever read;
* with a Prometheus exporter installed, the callback fires on COLLECT, so the
  reading happened when -- and only when -- somebody scraped ``/metrics``.

A deployment with no exporter or no scraper therefore served its whole life
with the pool unwatched: a memory check whose correctness depended on whether
anyone was looking at a dashboard.

:func:`note_step_served` is the causal edge -- the engine's post-step seam
calls it because a step ran. :func:`observe_forward_arena_reserved` is the
single implementation of "a serving reading of this pool exists"; the metrics
export feeds the SAME function with the reading it already has, so a scrape
still contributes and there is one implementation with two triggers rather than
two copies of the check.

AND THE POOL CAN BE GIVEN BACK WHILE THE PROCESS SERVES. :func:`note_engine_idle`
runs at the engine's own quiesce point -- an empty slate with both deferred
pipelines drained, which is the one moment nothing is live in this pool -- and
releases an arena that has overrun what the post-capture grow accounted for.
The KV layout is frozen by then, so what it repairs is not KV but this
process's own free VRAM: without it, a pool that overruns once is short by the
overage for the rest of the process's life, and the shortfall lands on the
transient rows (the verify tail, the DFlash context-assemble) rather than
anywhere it can be seen.

THE SAME EDGE CARRIES A SECOND OBSERVATION, on the other side of the allocator
boundary. :func:`observe_driver_serving_growth` records the driver-resident
physical that appears after the boot brackets close -- a cubin the driver
materialises the first time a served step dispatches a kernel no boot phase
launched (module loading is lazy unless an operator overrides it), the
kernel-stack pool it re-grows for a deeper frame, a graph instantiated
post-boot. That quantity had the identical defect and worse: the card already
NAMED it (``driver.modules_serving``), but the only thing that ever re-measured
it was the admin memory route, so on a process nobody inspected the row read 0
-- indistinguishable from the quantity being zero, and meaning the opposite --
and no term of the serving floor reserved a byte of it. It is now reserved by
``transient.serving_step.driver_growth``, which is sized from what this module
records.

WHY THE DRIVER READING RIDES THE ARENA'S GATE rather than carrying one of its
own. The arena gates on the cuMem ledger because its reading is a
``memory._snapshot`` walk. The driver reading has no such ledger to gate on --
CUDA exposes no counter for module residency, which is the whole reason the
term is measured rather than derived -- and its own cheap half is
``cudaMemGetInfo``, a driver call this repo keeps off the step path. So it
takes the arena's gate, and the two are coupled by more than convenience: the
eager serving forward allocates its activations INSIDE the arena, so a step
that presents a shape the pool has not served maps new physical under that tag
-- and a shape the process has not run before is exactly when a kernel it has
not dispatched before gets dispatched, which is when a cubin loads. The audit
that found the quantity measured both moving over the same suite (the arena
+408/+552 MiB while the driver added +131).

WHAT THAT GATE DOES NOT PROMISE, said plainly: a cubin loaded on a step that
maps no new arena physical is not seen when it happens. It is picked up by the
next reading this module takes, by a metrics scrape where one exists, or -- if
this process takes neither -- by the next boot, because the record is a
cross-boot high-water and the quantity saturates. That is convergence over
boots, not within one, and it is weaker than the arena's guarantee. Closing it
needs one measurement this repo does not have: what ``cudaMemGetInfo`` costs on
the step path. If it is under the step's noise floor the gate becomes the
driver's own cheap read and the coupling above stops being load-bearing.

A served step finished: read the pool if it grew, and check it.

Called from the engine loop's post-step success seam
(:func:`arbi_serve.engine.run_step.loop._run_forever_inner`), which is the
one place that knows a step ran rather than that somebody asked.

GATED ON THE POOL HAVING MAPPED NEW PHYSICAL, because the reading itself is
not free: the pool's reserved bytes come from a walk of
``torch.cuda.memory._snapshot`` (see
:meth:`~arbi_serve.runtime.named_pool.NamedMemPool.snapshot`), which
:meth:`~arbi_serve.runtime.named_pool.NamedMemPool._check_size_target_on_exit`
already refuses to do per ``use()`` on the decode path for measured
reasons. The gate is the cuMem allocator's per-tag mapped-bytes ledger,
which is a dict read kept in step with every map and unmap -- and it is the
same event: this pool's reserve can only move when the pool maps or unmaps
physical under its tag.

The gate cannot silently mean "did not run". Its hint only advances after a
reading is actually taken, so a step that skips the read (the allocator
lease is closed) retries on the next one, and a deployment with no cuMem
ledger at all -- ``--no-cumem-pools``, the unsupported backend -- says so in
the log instead of quietly watching nothing.

The share of the floor's draw that a SERVED STEP is responsible for.

Pure arithmetic, separate from the reading so the deduction is testable
without a card. Clamped at zero: the three readings come from three
instruments and the deductions can momentarily exceed the draw, which is
"no served step has taken anything yet", not a negative draw.

Check the boot realization covered what real traffic takes.

The floor's terms price what a step allocates through a pool we own. The
CUDA caching allocator also takes fresh segments OUTSIDE every named pool
when a served step presents a block size no boot phase realized, and those
segments come out of the same free VRAM the floor holds. Boot MEASURES that
by RUNNING the widest step this configuration can issue, through the
scheduler and the shipped batch materializer
(:func:`~arbi_serve.engine.serving_realization.realize_serving_step`), and
the floor holds what that step was measured to take.

THIS IS THE NULL CONTROL FOR THAT READING, and nothing more. Traffic taking
more than the boot step took is traffic reaching a shape the realization
did not present — the floor is then covering it out of whatever it has
spare, and the only honest place to learn that is while it happens. It is a
check, not a reserve: nothing here is carried to another process, because a
reading one process took under one workload is a forecast about the next
boot rather than a fact about it, and the point of running the step at boot
is that a boot no longer has to inherit one.

The quantity is how far FREE VRAM fell below the floor the KV grow left,
because that is what the floor is: a free-VRAM reserve. The caching
allocator's own counters cannot see the driver's share or a cuMem map, and
all three come out of the same free.

BOTH SIDES OF THE TEST MUST BE THE SAME QUANTITY, and the whole free-VRAM
drop is not. Three different things draw on the floor and only one of them
is what the realization measured:

* the BOOT's own post-freeze phase — the serve-kernel warmup and the
  readiness forward — whose draw is measured and gated by
  :func:`~arbi_serve.engine.post_freeze_budget.assert_post_freeze_floor_intact`
  and is permanent residency from before the first request;
* the DRIVER's post-freeze growth, which has its own floor row
  (``transient.serving_step.driver_growth``) and its own warning in
  :func:`observe_driver_serving_growth`;
* the served steps, which is the only draw the realization's reading is a
  reading of.

So the first two are deducted here, each from the measurement its own
surface already took, and the remainder is what this compares. Charging
them here instead would make the test fire on every boot that warms a
kernel — a null control that fires when nothing is wrong reports nothing —
and would double-report bytes another warning has already named.

Nothing here samples on a schedule and nothing depends on an observer being
attached: the reading is taken by the step loop itself. It is taken AFTER
the step, where the step's transients are already freed — so it reads the
steady state, not an in-step trough.

POST-GROW ONLY, the same condition its driver-growth sibling uses: the
floor publishes its rows when it sizes, and their presence is what says a
reading taken now describes the card the grow left behind.

The engine has nothing in flight: give back an arena that has overrun.

Called from the engine loop's idle branch
(:func:`arbi_serve.engine.run_step.loop._run_forever_inner`), after the
async-output and deferred-verify drains and only when the scheduler holds
no running and no waiting row. Returns the bytes handed back to the driver
(0 when nothing was).

WHY THERE IS ANYTHING TO GIVE BACK. The arena is a private cuMem pool: a
freed block stays mapped, and the allocator serves a request from a free
block only if it FITS, so the widest allocation of each new size class
takes a fresh segment and every segment a narrower class left behind stays
mapped and idle. The pool's reserve is therefore a sum over the classes a
workload has presented, not a maximum over them, and most of it can be
wholly-free segments while the live set is a fraction of the pool.

WHEN IT RUNS, and why that condition and not a timer: only when the pool
holds MORE than the post-capture grow accounted for (its reserve plus what
it already held). Below that line the pool's physical is already booked —
handing it back would swap resident bytes for free bytes and change
nothing. Above it, the process is short by the overage for the rest of its
life: the KV layout is frozen, so those bytes come out of the free VRAM
the serving floor is holding for one in-flight step, which is the reserve
the verify tail and the DFlash context-assemble fail out of first. That is
a MEASURED condition on this process's own card, so the re-tightening
fires when there is something to repair and never otherwise.

WHAT IT COSTS THE STEP LOOP: nothing. It runs only with no step in flight
and no row admitted, and the ordinary idle poll pays one dict read on the
cuMem ledger — the same read the served-step gate uses — which is what
keeps a fast-polling idle loop from walking ``memory._snapshot``. The walk
and the release happen only after that read says the pool has moved past
the accounted mark, and the attempt is not repeated until it moves again.

WHAT IT DOES NOT DO: hand anything to KV. The layout froze at boot, so the
returned physical is free VRAM for the transient rows this process is
already short of — the repair is in-process, and nothing about it reaches
another boot.

One serving reading of ``scratch.forward_arena``: check it against the grow.

The single entry point for both triggers -- the served step above and the
metrics export's reserved-bytes gauge, which passes the reading it already
produced for the gauge. Both share the engine's warn latch, so the two
cannot drift apart or double-report.

The reading is CHECKED and never carried to another process. What the pool
holds is a property of the workload this process served, and a reading
taken under one workload says nothing about a boot that will serve another
-- the version of this that wrote the number into the budget cache
under-reserved a production boot from a process that had served a single
short request. This boot's cover for the pool is the physical the boot
itself realized (:func:`~arbi_serve.engine.inprocess_capture.
forward_arena_regrow_plan`); the job here is to say LOUDLY when serving
goes past it.

One serving reading of the driver's own residency: record it, then check it.

The single implementation of "a serving reading of the driver-side growth
exists", with the same two triggers its arena sibling has: the served-step
edge above, and the metrics export
(:meth:`~arbi_serve.server.metrics._collectors.MetricsCollectorsMixin.
_check_driver_serving_growth`), which owns nothing and calls this.

WHAT IS MEASURED. ``DriverResidencyMeter`` opens a bracket at
:meth:`~arbi_serve.engine.memory_budget.driver_residency.
DriverResidencyMeter.seal_boot` that stays open for the engine's life. The
growth across it is ``non_torch_bytes`` now minus at the boot close: our
own share of the device, less the torch allocator's reservation, less the
regions cuMem maps directly. That is the superset of the three rows the
live card splits it into -- ``driver.modules_serving`` plus the
kernel-stack and graph-exec growth carved out beside it -- and the superset
is the right quantity for a FLOOR, because all three are free VRAM the grow
handed to KV that the driver then took.

BOTH SAMPLES MUST SPLIT THE CARD THE SAME WAY. ``non_torch_bytes`` is built
on this process's share of the device, and which of the three readings
resolved that share (NVML naming us, the pre-context mark, or the whole
card as a last resort) can differ between two samples. A difference taken
across two different splits is not growth, it is the two instruments
disagreeing, so it is discarded rather than recorded as a reserve the next
boot would hold.

Best-effort throughout: an accounting read must never break a step.

Signed non-torch growth over the serving bracket, or ``None``.

``None`` on every reading that cannot be trusted rather than a 0 that
cannot be told apart from one: the allocator lease is closed, the probe
failed, the bracket has no boot close, or the two samples split the card
differently.

Under :func:`~arbi_serve.runtime.named_pool.allocator_query_lease` for the
same reason the arena's read is: the probe reads
``torch.cuda.memory_reserved``, which takes the caching allocator's mutex,
and the lease is also what serialises this sampler against the admin memory
route's -- the two share the meter's single serving slot.

Put a rising driver-growth reading on record for this configuration.

Cheap on the common path: an int compare against the largest reading this
process has taken.

MONOTONIC in the cache too, unlike the modules BASELINE stored under the
same key: that one is a gate a leak could walk upward a boot at a time, so
it is written once; this is a budget bound on a quantity that only grows
within a process -- a resident cubin stays resident -- so the largest
reading is the one a later boot must hold.

Fire and (once) narrate driver growth above what the floor holds.

Read from the published term rather than re-resolved, for the reason its
arena sibling states: re-resolving would consult the record
:func:`_record_driver_growth` just wrote and check the quantity against
itself, so the overage could never fire.

Ask the gate to re-measure only when the overage GREW.

Both reports above fire on every READING taken while their condition holds,
and one of the two triggers for those readings is the metrics export -- so
a standing overage under a Prometheus scraper reports on a cadence. A
re-arm request raised from every report would inherit that cadence and
become the periodic re-measurement this mechanism deliberately does not
have.

A RISE is the new information; persistence is not. The budget is a function
of what the card holds, so a re-measurement against an overage the gate has
already been re-armed against can only produce the answer it already gave.
One int compare on the engine.

Bytes the cuMem ledger currently maps under the arena's tag, or ``None``.

``None`` means there is no ledger to ask -- the process never built the
cuMem allocator, which is what ``--no-cumem-pools`` produces. Read through
:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.peek` so this
never brings the native allocator into existence.

A CHANGE, in either direction, is the trigger -- not a rise. The reading it
schedules is absolute, so an unmap that is later re-mapped is still
checked.

The pool's reserved bytes, or ``None`` when the allocator is busy.

Reads through :func:`~arbi_serve.engine.inprocess_capture.
arena_resident_reserved_bytes`, which is the serving floor's own reader for
this quantity -- the recorded number and the number the floor subtracts
residency from have to come from one place, or the pair stops closing.

Under :func:`~arbi_serve.runtime.named_pool.allocator_query_lease`: the
walk takes the caching allocator's mutex without releasing the GIL, and the
engine's pipelined build thread allocates into cuMem pools holding that
mutex and wanting the GIL. That is the inversion the lease exists for, and
this caller is on the engine thread but not the only thread allocating.

Fire and (once) narrate a pool above what the grow accounted for.

NOT gated on the warning latch: the WARNING latches (the overage is a
standing condition, not an event) but the COUNT must not, or a reader
cannot tell one reading's crossing from a whole process spent over the
line -- which is the difference between a blip and the degradation this
watch exists to name.

Say ONCE that this backend leaves the serving arena unmeasured.

Reached only when the engine has an arena pool but the process has no cuMem
allocator, i.e. ``--no-cumem-pools``. There is then no per-map ledger to
gate the reading on, and taking the reading unconditionally would put a
``memory._snapshot`` walk on every decode step. Saying nothing would leave
a floor that silently seeds from the boot mark forever, which is the defect
this module exists to remove -- so the boot is told which of the two it is.

Bind per-head learned attention sinks onto the per-layer attn ops.

A sink is an extra logit per query head that enters only the softmax
denominator. The model owns it as
:attr:`arbi_serve.models.attn.AttentionBlock.sinks` (allocated whenever
the layer's :attr:`LayerSpec.attention_sinks` is set, and filled by the
weight loader); the kernels take it as a per-call argument. This module
is the seam between the two: it runs at boot, after ``eng.attn_ops`` is
built and after the weights are loaded, and pushes each block's tensor
onto its op through the duck-typed ``set_attention_sink``.

A backend that cannot carry a sink refuses in its ``make_attn_op``
(:class:`arbi_serve.backends.mla_backend.MlaBackend`), so an op that
reaches here without the hook on a sink-declaring layer is a wiring
bug, not a silent degradation — :func:`bind_attention_sinks` raises.

Per-step batch metadata — typed per-state-kind dataclasses.

A :class:`ScheduledBatch` carries the routing fields every step needs
(``is_prefill``, ``input_ids``, ``positions``, scheduling tensors) plus
optional ``*Meta`` slots for each state kind. Per-layer ``AttnOp``
instances read exactly one typed object — their kind's metadata view —
never the whole batch.

The engine keeps the raw scheduling tensors on :class:`ScheduledBatch`
as the bookkeeping copy used to *build* the per-kind metadata; per-kind
builders return the typed dataclasses ``AttnOp``\s consume.

Per-step metadata for ``StateKind.PAGED_KV`` layers.

Carries the FA / TKV per-layer kernel inputs:
``cu_seqlens_q``, ``seq_lens``, ``block_table``, ``slot_mapping``,
``max_seq_len``, ``max_query_len``, ``causal``, ``sliding_window``.
The TKV ``TQRunState`` and the Turbo prefill scratch attach as
additional fields populated by :class:`TkvMetadataBuilder`.

Per-step metadata for ``StateKind.MLA_SHARED`` layers.

Same paged shape as PAGED_KV; slot byte layout differs (see
:class:`MLAStatePool`).

``total_kv`` is ``sum(seq_lens)`` — the exact row count of the
gathered bf16 latent workspace the multi-token path materialises,
and by construction the value ``cumsum(seq_lens)[-1]`` holds. It is
a plain Python ``int`` carried from the host-side scheduling data
(:attr:`ScheduledBatch.total_kv_tokens`): reading it off the device
is a host sync, which serialises the step and is illegal inside a
CUDA-graph capture region.

``is_prefill`` mirrors :attr:`AttnPagedKVMeta.is_prefill` — the
step class the scheduler assigned, as opposed to the launch shape
the op routes on (``max_query_len`` / ``mtp_block_m``).
:class:`~arbi_serve.models.mla_block.MLAAttentionBlock` reads it to
gate the calibration seam's fresh-diagonal split.

Per-step metadata for ``StateKind.MAMBA`` layers.

Per-request state lives in :class:`RecurrentStatePool` as a flat
slab tensor; here we carry the per-row slab-row index + past-
token-count plus the prefill boundary tensor that Mamba-2's
varlen path consumes.

Convention:
  * ``state_indices`` is per-SEQ — shape ``(B,)`` int32 mapping
    batch row i → slab row in :class:`RecurrentStatePool`. The
    block reads ``slab[state_indices[i]]`` for each row.
  * ``has_initial_state`` is per-SEQ — shape ``(B,)`` bool.
    True iff the request already pushed tokens through this
    layer's recurrent state in a prior step (i.e. seed from the
    live slab row; otherwise the row is freshly zeroed).
  * For multi-token-per-request prefill (n_tokens > B), Mamba-2
    blocks bucket tokens via ``cu_seqlens_q`` (shape ``(B+1,)``).
    Mamba-1's existing path predates that and synthesizes single-
    request bucketing from ``state_indices`` directly.
  * ``verify_pass`` — multi-token-per-row batch with ``mtp_meta``
    set and ``is_prefill=False``. :class:`Mamba2Block` routes it
    through the PER-TOKEN decode kernels (``causal_conv1d_update``
    + ``selective_state_update``) instead of the varlen chunk scan,
    which writes only the post-all-K+1 state and leaves nothing to
    roll back to.

Where this step's GDN prefill row splits its recurrent fold.

A savepoint can only be restored at a position on the KV page grid, and
a co-batched prefill chunk almost never ENDS on one — so the state that
position names exists only inside the step, swept through by one kernel.
This plan is how it is made to exist: the forward issues the fold as two
calls around ``offset`` and hands the intermediate state to
``snapshot`` while it is live.

``row`` indexes the batch, ``offset`` is tokens into that row's segment
of this step, and ``boundary`` is the absolute prompt-token count the
resulting state will be labelled with — ``offset`` tokens past the row's
chunk start. All three are decided in ONE place
(:meth:`~arbi_serve.scheduler.scheduler.Scheduler._arm_fold_split` via
:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_fold_split`) and
carried here; the forward re-derives none of them, because a boundary
derived twice is a coverage label that can disagree with the state it
names.

``offset`` is a multiple of the recurrent kernel's fold grid by
construction — ``savepoint_fold_split`` refuses rather than rounds — and
that is what makes the two-segment fold bit-identical to the one it
replaces (``tests/test_savepoint_fold_split_gpu.py``).

Per-step metadata for ``StateKind.GDN`` layers (Qwen3-Next,
Qwen3.5 / Qwen3.6).

Carries the per-row state-slab mapping + the prefill-only
``cu_seqlens_q`` boundary tensor that FLA's varlen prefill
kernel (``chunk_gated_delta_rule``) consumes. For per-request-
one-token decode (``n_tokens == B``), ``cu_seqlens_q`` may be
``None`` — the block routes to FLA's batched
``fused_recurrent_gated_delta_rule`` and reads ``state_indices``
directly. Prefill (multi-token, ``is_prefill=True``) routes to
``chunk_gated_delta_rule``. The MTP verify pass (multi-token per
row at T = K+1, ``is_prefill=False``) routes to
``fused_recurrent_gated_delta_rule`` with ``cu_seqlens`` so the
bonus-slot output shares numerics with the K=1 decode kernel —
the bundled Qwen3.5 MTP head was trained against fused-recurrent
output and is sensitive to the bf16 reduction-order ε divergence
of the chunk kernel. ``verify_pass`` selects the kernel.

``state_indices`` is per-SEQ — shape ``(B,)`` int32 mapping batch
row i → slab row in :class:`RecurrentStatePool`. The block reads
``slab[state_indices[i]]`` for each row.

``has_initial_state`` is per-SEQ — shape ``(B,)`` bool. True iff
the request already advanced this layer's recurrent state in a
prior step (i.e. seed from the live slab row; otherwise the row
is freshly zeroed and effectively a no-op).

Per-step metadata for ``StateKind.SHORT_CONV`` layers (LFM2 family).

LFM2's ShortConv mixer needs, per request, (a) the row index into
the per-layer conv-state slab and (b) whether this step continues
a previous prefill chunk (``has_initial_state``) so the kernel
knows whether to seed from the saved buffer or zero. ``cu_seqlens_q``
is shared with PAGED_KV scheduling — but layers reading this meta
only see the ShortConv-relevant subset.

Fields:
    state_indices: ``(B,)`` int64 — for each request in the batch,
        its slab row index into the conv-state slab.
    has_initial_state: ``(B,)`` bool — True iff this request has
        already pushed tokens through the conv in a previous step
        (i.e. seed conv-state from the slab; otherwise zero-seed).
    cu_seqlens_q: ``(B+1,)`` int — per-request token boundaries in
        the flat ``(N_tokens, ...)`` input. Same tensor as
        :class:`AttnPagedKVMeta.cu_seqlens_q` for the attention
        layers in the same step.
    seq_lens: ``(B,)`` int — per-request total token count
        (prefill length or running decode position).
    is_prefill: True if this batch is a prefill step (kernel
        dispatches `causal_conv1d_fn`); False for decode (kernel
        dispatches `causal_conv1d_update`).

Per-step metadata for ``StateKind.DSV4_SPARSE`` layers (DeepSeek-V4).

The DSv4 attention advances three per-request stores — a sliding-window
ring, a compressed stream and the indexer's stream — plus a pooling
accumulator, and it does so for the whole flat batch at once. To stay
free of device→host reads it needs each request's slab row and its
token count on the HOST, the same reason
:attr:`ScheduledBatch.total_kv_tokens` and :attr:`max_query_len` are
plain ints: deriving them from ``cu_seqlens_q`` would sync the step and
is illegal inside a captured graph.

Fields:
    state_indices: ``(B,)`` int64 — each request's row in the per-layer
        stores, the same role it plays for the recurrent kinds.
    query_lens: per-request token count in THIS step, on the host. A
        uniform decode is ``(1,) * B``; a prefill chunk carries the
        chunk's own lengths.
    verify_pass: multi-token-per-row batch with ``mtp_meta`` set and
        ``is_prefill=False``. The DSv4 attention then keeps what
        :meth:`~arbi_serve.cache.dsv4_pool.DSv4StatePool.rollback_batch`
        needs to undo the rejected tail: a copy of the ring rows the
        block overwrites, and a frame of the pooling accumulators.
    transient: the step's tokens are never committed — a DSpark
        drafter's own block. Nothing is published to the window ring,
        because nothing will roll it back.

Host-resident twin of a batch's core scheduling tensors.

Populated by :meth:`ModelRunner._build_batch` on the persistent
(cudagraph) hot path, where every core tensor was just written into
the pinned-host ``PiecewiseBuffers.h_*`` ring before the H2D copy
into the device twins. Carrying the host views here lets the TP=2
rank-0 ``_serialize_batch_to_plan`` build the worker ``StepPlan``
straight from host memory — no device→host copy of tensors the host
already holds, and no per-step ``torch.cuda.synchronize()`` drain on
the rank-0 critical path.

Every field is a CPU tensor sliced to this step's active extent
(``[:N]`` token-flat, ``[:B]`` / ``[:B+1]`` per-row). ``input_ids``
is the post-``gpu_overwrites`` host value: ``_build_batch`` writes
each decode row's authoritative ``last_token_id`` back into
``h_input_ids`` so this mirror is bit-identical to the device
``input_ids`` the forward consumes. ``cu_seqlens_k`` has no ``h_*``
twin (it is a device cumsum), so it is recomputed on the host from
``seq_lens`` — an exact integer cumsum, bit-identical to the device
one.

Per-step MTP verify metadata carried on :class:`ScheduledBatch`.

Per-position draft mask + per-request K bookkeeping. Built by
the engine in :func:`mtp_step` from the slate's per-row K and
the draft slots returned by
:meth:`MultiStatePool.allocate_draft_slots`. Consumed by the
verify-pass logits gather (:func:`gather_verify_logits`) and by
the post-step draft-slot free path.

Fields:
    per_req_k: ``(B,)`` int32 — per-row K (0 for non-MTP rows).
    draft_token_mask: ``(N_tokens,)`` bool — True at positions
        that are draft tokens (one of the K appended after each
        row's ``last_committed`` token); False at the leading
        committed-token slot.
    per_req_offsets: ``(B+1,)`` int32 — flat per-token boundaries
        (``cu_seqlens_q``-shaped, but specific to the verify-row
        layout where each row contributes ``K+1`` flat tokens).

The richer working copy (with draft_slots / request_ids) lives in
:class:`arbi_serve.spec_decode.mtp_meta.MtpStepPlan`; this batch-
side dataclass holds only the fields per-block forwards may need
to consult mid-step.

One step's worth of work, ready for the model.

Field shapes:
  - ``input_ids``    : int32 ``(N_tokens,)``        flat over all seqs
  - ``positions``    : int32 ``(N_tokens,)``        absolute pos per token
  - ``cu_seqlens_q`` : int32 ``(B+1,)``             token-boundary csum
  - ``seq_lens``     : int32 ``(B,)``               full seq len per req
  - ``block_table``  : int32 ``(B, max_pages)``
  - ``slot_mapping`` : int64 ``(N_tokens,)``        flat slot index per token
  - ``max_seq_len``  : int                          longest full seq
  - ``max_query_len``: int                          longest per-step query

The ``*_meta`` fields are populated by per-backend metadata
builders only for state kinds present in the model. AttnOps read
only their own kind's meta — never the rest of the batch.

Route a per-kind built ``*Meta`` onto the correct slot.

Single source of truth for the StateKind → ``*_meta`` field
dispatch. Both the regular step path and the MTP verify path
call this after invoking each kind's :class:`MetadataBuilder`,
so adding a new state kind only touches this method.

The boot clock's origin and the breakdown of the region before ``build()``.

``build()`` is entered after the interpreter has started, torch has imported and
the engine has been constructed — several seconds on a cold page cache. A clock
started there reports a boot shorter than the one an operator waited through,
and the phase table sums to that shorter number, so the accounting looks exact
while excluding the time. The origin resolves to process start so the total is
what ``docker run`` to serving actually costs.

That region is several distinct costs — interpreter start, the torch import,
the rest of the import graph, CLI + config resolution, app construction, CUDA
primary-context creation and the engine ctor — and one number for all of them
cannot be optimised. This module stamps each boundary and hands ``build()`` the
telescoping parts:

  interpreter start → torch import → python imports → cli + config resolve →
  app + server start → cuda context → engine ctor

The stamps are torch-free and cost a ``perf_counter`` read, so
``arbi_serve/__init__`` can take them before anything heavy is on the path.

A ``perf_counter`` value marking process start.

Falls back to ``fallback_perf`` (typically the engine ctor stamp) and then
to now, so a platform without ``/proc`` still gets a monotonic origin — just
a later one, which under-reports rather than inventing time.

``sys.meta_path`` finder that times the first import of each top-level module.

It resolves the spec through the finders behind it and wraps the loader's
``exec_module``, so the recorded seconds are the module's own execution plus
everything it imports. Nested top-level imports (torch pulling numpy) are
recorded at a non-zero depth and dropped, leaving a disjoint set that sums
to the import graph's cost.

Only dotted-free names are timed: a submodule's cost already belongs to the
root that pulled it, and wrapping every spec would cost more than it
measures. A name is wrapped on every lookup until it has actually executed,
because a bare ``importlib.util.find_spec`` probe resolves a spec that is
then discarded — the import that follows gets a fresh loader.

Create this process's CUDA primary context and record what it cost.

The context is created lazily by whichever CUDA call happens to run first,
which buries it inside an unrelated phase. Forcing it here — a device
synchronize, which needs a context but allocates nothing through the
caching allocator, so the boot VRAM baseline still reads the bare context —
gives the boot table a number for it. No-op without CUDA.

Telescoping ``(phase, seconds)`` for everything before ``build()``.

The parts sum exactly to ``now - origin`` so the boot phase table still
accounts for the whole region. The CLI stamps are optional: a programmatic
embedder that constructs an ``Engine`` directly gets the import graph as one
part instead of three. Returns the single fused phase when a required stamp
is missing or the arithmetic would go negative — an under-split total is
honest, an invented one is not.

Boot-time configuration SUBSTITUTIONS — counted once, reported once.

Several boot searches MEASURE a value on the card and, when the card cannot
run the probe, install a closed-form (or eager) stand-in instead. Keeping the
server up that way is legitimate. Doing it QUIETLY is not: the server then
serves a configuration nobody asked for, and a per-probe warning mid-boot is
not a way anyone learns that.

Two properties of the boot make the silence structural rather than accidental,
which is why the fact needs a home of its own rather than a better log line:

  * the ready banner is identical for a boot that installed what it was asked
    for and one that did not — it reports readiness, and readiness is true
    either way;
  * ``_capture_oom_failures`` — the ledger the post-capture VRAM gate refuses
    on — deliberately EXCLUDES a prefill rung at or above
    :data:`~arbi_serve.engine.capture_admin.layer._HIGH_RUNG_FLOOR`, so a
    small card is not refused a boot. The top rung is ``chunk_prefill``
    itself, so the rung the operator named is precisely the one that gate
    cannot see.

This module is the ONE place a substitution is recorded, and it publishes the
fact three ways so no single one of them has to be trusted alone:

``boot_autotune_oom_fallback``
    A :mod:`arbi_serve.flag_truth` counter, one fire per probe that ran out
    of memory AND lost its measurement to a stand-in. Interrogable long after
    boot at ``GET /v1/admin/flag_truth``, and the thing a test can assert a
    magnitude on. Probes that OOM without costing a measurement (a
    self-bounding ladder that still chooses among the widths that ran) do NOT
    fire it: the counter counts fallbacks, not exceptions.

``boot_state.boot_substitutions``
    The structured ledger — what was requested, what was installed, why. The
    rows that cost the capture sweep a bucket also carry
    ``truncates_capture_pool``, and :func:`capture_pool_truncations` is what
    the graph-pool persist reads to refuse writing a budget measured over a
    ladder that did not run into the cross-boot cache.

:func:`log_boot_substitutions`
    ONE structured record at the boot banner, grouped so a sweep that
    substituted fourteen rungs is one line and not fourteen.

Recording decides nothing. A site that records a substitution still installs
it, and a boot with substitutions still serves; whether a particular stand-in
is safe enough to serve belongs to the site that makes it. What the site needs
from here, besides the ledger, is the vocabulary to refuse in:
:class:`RequestedConfigurationNotInstalled` is the refusal a site raises when
the value it could not install is one the OPERATOR named rather than one the
engine picked.

True when ``exc`` is, or was raised from, a CUDA out-of-memory error.

The boot probes convert an OOM into a typed refusal
(:class:`~arbi_serve.runtime.activation_profile.ActivationProbeFailed`,
raised ``from`` the original), so the handler that later swallows that
refusal sees a ``RuntimeError`` and would book a memory fallback as an
unrelated failure. Following the ``__cause__`` chain is what keeps the
counter's name true.

A value the OPERATOR named could not be installed, so the boot refuses.

The discriminator is provenance, not severity. When the engine chose a
value, the engine may re-choose it: a stand-in is recorded here and the
server serves. When the operator named the value, a stand-in serves a
configuration nobody asked for under a ready banner that cannot say so —
the substitution defeats the request instead of adapting it.

Carrying a distinct type (rather than a bare ``RuntimeError``) keeps the
two apart for a caller that wants to tell "this card cannot do what was
asked" from "this boot hit a bug". The message is the whole payload and
must name three things: the value requested, the value that would have
been installed in its place, and the knobs that move the outcome. A
refusal that says only that something failed is a worse version of the log
line it replaces.

One installed value that is not the one the config asked for.

``requested`` and ``installed`` are rendered for a person, because the
reader's question is "what am I actually running", not "what type is
this". ``oom_probes`` is 0 for a substitution no OOM caused — a
pre-emptive headroom gate, a probe that failed for another reason — so
the counter and the ledger can never disagree about what memory did.

``truncates_capture_pool`` is the site's own declaration that this
substitution REMOVED graphs from the capture pool a later
``total_cudagraph_bytes`` reading measures. It is set by the sites that
drop a capture bucket and by nobody else: a stand-in installed elsewhere
in the boot (an unmeasured activation profile, a closed-form DFlash tile)
changes what the server runs without changing what the capture sweep
paid, and must not be read as a truncated pool.

Record that ``what`` came out as ``installed`` instead of ``requested``.

Fires :data:`_OOM_FALLBACK` once per OOM'd probe that was lost to this
substitution, and appends the record to ``boot_state.boot_substitutions``.

Tolerates an engine double with no ``boot_state``: the counter is the
part a test and an operator interrogate, and losing the ledger must never
turn a working boot into a crash.

The substitutions that cost this boot's capture sweep a bucket.

The signal a persist decision needs, and the reason it cannot use
``_capture_oom_failures``: that ledger answers "is this card too small to
boot", and it answers it by EXCLUDING every prefill rung at or above
:data:`~arbi_serve.engine.capture_admin.layer._HIGH_RUNG_FLOOR` — the
rungs whose loss shrinks the measured pool the most. A boot that dropped
only those reads as zero failures, so a budget measured over the rungs
that survived is indistinguishable from one measured over the full
ladder, and the cache keeps no record of which rungs ran.

"Refuse the boot" and "refuse to persist" are different questions. This
is the second one's signal: whether the number about to be written
describes the configuration the cache key claims it describes.

Emit :func:`boot_substitution_report` at WARNING, once, if non-empty.

WARNING and not INFO: the boot banner two lines above says the server is
ready, and a degraded boot that renders at the same level as a clean one
is a boot an operator reads as clean.

Progress line for boot phases that would otherwise run silently.

Several boot phases do minutes of work and emit nothing until they
finish: the phase-start line, then a gap, then the completion line. On a
cold kernel cache that gap is long, and it is **indistinguishable from a
hang** — no output, no error, no non-zero exit, and (for a compile) no
GPU activity either. An operator, an orchestrator, and a bisect all read
that the same wrong way.

So a running phase says so. While a phase is open a daemon thread emits
one line every :data:`HEARTBEAT_INTERVAL_S`, naming the phase and how
long it has been running. The line advancing means the process is alive;
the line *stopping* is the actual hang signature. That distinction is the
whole point — it converts "no output" from ambiguous into diagnostic.

Daemon thread, so it can never hold the process open, and it only ever
logs — a heartbeat that could fail a boot would be worse than the silence
it replaces.

The same module owns the two things a hung boot needs in order to say
WHERE it hung: :func:`arm_stack_dump_signal` (on-demand, ``kill -USR1``)
and the wedge deadline (unattended). Both go through :mod:`faulthandler`,
which writes from C without taking the GIL — the only mechanism that
still works when the interpreter itself is frozen.

Make :data:`STACK_DUMP_SIGNAL` dump every thread's stack. Idempotent.

A boot that can hang must be able to say where. Until this is armed
SIGUSR1 is unhandled, so the one signal an operator would reach for
KILLS the container and takes the diagnosis with it.

:mod:`faulthandler` is the mechanism, not :mod:`signal`, and the
difference is the whole point: a Python-level handler runs at a
bytecode boundary, so it cannot fire while another thread holds the
GIL inside a blocking native call — which is the exact shape of the
hangs worth dumping. ``faulthandler`` writes the tracebacks from the
signal handler in C, GIL or no GIL.

The process CONTINUES after dumping (``chain=False``): the signal is
a probe, not a kill, so a stuck boot can be sampled repeatedly to
tell a deadlock (identical stacks) from slow progress (moving ones).

``chain=True`` re-arms on top of an already-installed Python handler
for the same signal and runs it after the dump — how
:func:`arbi_serve.runtime.capture.dispatch.coverage.install_signal_dump`
keeps its post-boot coverage dump without giving up the thread dump.

Returns whether the handler is armed. Never raises: a platform
without the signal loses the probe, not the boot.

(Re-)arm the boot's wedge deadline. No-op when disabled (``0``).

The timer is process-global (one ``faulthandler`` timer), so this is
an ARM-OR-RESTART: each call gives the boot a fresh full budget from
now. Boot work is sequential, so restarting it at every phase
boundary is the same thing as a per-phase budget — with the
difference that the BETWEEN-phase stretches stay covered, and those
are not small: more than half the boot's wall clock runs outside any
``_phase_start``/``_phase_done`` bracket — the weights-DMA join, the
cuBLAS workspace warmup, the KV sizing and the post-capture grow all
do. Cancelling on phase close leaves every one of them with no
deadline at all, and a wedge there then sleeps forever.

Hand the deadline back. Call ONCE, when the boot is over.

Serving is not boot: a request that legitimately runs longer than the
boot budget must not meet a timer that exits the process.

The handle :func:`announced_boot_step` yields.

The body sets what only the body can know: whether the bracketed work
actually RAN, and in what words. A cached ``.so`` load and a 210-second
``nvcc`` build pass through the same bracket, and only the step itself
can tell them apart — a duration cutoff would be an invented constant
that hides a slow cached load and shows a fast fresh build.

Left unset, both stay ``None``: unknown, which a timeline must render
rather than fold away.

Bracket a mid-boot step that the build's phase sequence does not.

The build procedure opens a phase around each step it OWNS. Work
triggered lazily from inside a dependency is not one of those: it can
land in any phase or between two of them, and until it is bracketed it
is the one stretch of the boot that cannot name itself — no phase open,
no heartbeat, no line. A JIT cpp/CUDA extension build is exactly that
shape, and it is minutes long: a boot sat 210 s reporting ``phase:
null`` with ``nvcc`` running.

So it gets a window of its own, and the window is NESTED
(:func:`~arbi_serve.boot_progress.nested_phase_opened`) — the enclosing
phase keeps running and is restored on close, because it genuinely is
still running. Widening the enclosing phase to cover the step instead
would report the seconds under a name that does not describe them,
which is the same defect with a better name on it.

Only while the boot is in progress. Past ready this machinery must not
fire: the heartbeat re-arms the boot's wedge deadline
(:func:`arm_stall_deadline`), which the end of the boot deliberately
hands back, and re-arming it under a long serving request would exit
the process.

The announcement is best-effort in one direction only — it never
suppresses the step's own exception, and a failure to announce is
LOGGED rather than swallowed, since a silently lost announcement is
the failure this exists to prevent.

The OPENING line is unconditional and stays that way: the step's length
is unknowable until it ends, so deferring the announcement would put
the longest silent stretch of a cold boot back where it was. What the
step learns about itself goes on the CLOSING line and into the phase
record, via the yielded :class:`BootStep`.

This must survive the exact failure it exists for. The heartbeat above
is a Python thread, so it only runs when the interpreter can schedule
it: a phase blocked in a driver call that holds the GIL freezes the
heartbeat too. That is not hypothetical — it is the observed
signature. A wedged cuMem pool allocation at drafter attach stopped the
heartbeat dead, so a deadline living in that thread could never have
fired, and the process sat silent until it was killed by hand.

The deadline itself is boot-wide, not phase-scoped: see
:func:`arm_stall_deadline`. A phase opening simply restarts it.

Stop the progress line. Idempotent; safe from any thread, twice.

Deliberately does NOT touch the wedge deadline. A closing phase is
not the end of the boot: the stretch between two phases is still
boot, and several of the longest ones — the weights-DMA join, the
cuBLAS workspace warmup, the KV sizing — live entirely there. The
deadline is restarted at the next boundary and handed back once, at
the end of the boot, by :func:`disarm_stall_deadline`.

Wall-clock accounting for boot work that runs off the main thread.

A linear phase timeline bills an overlapped task to whichever phase is
open when the task is JOINED. That is a true statement about the wall
clock and a useless one about the task: the phase's duration cannot say
whether the main thread spent those seconds doing other work or standing
still at ``Future.result()``.

Each task recorded here carries three numbers instead:

``duration``
    submit → the worker returning. What the boot had to wait out.
``blocked``
    the main thread's own wall inside the join.
``overlapped``
    ``duration - blocked`` — the seconds the overlap actually bought.

``overlapped ≈ 0`` means the async task is a serial one wearing a
thread, and the phase that joins it is a wait, not work.

Run the task with the submitter's CUDA device ambient on the worker.

The ambient CUDA device is THREAD-LOCAL: ``torch.cuda.set_device`` on the
boot thread does not reach a pool worker, whose ambient device is cuda:0.
Backgrounded boot work builds its tensors on explicit devices and so reads
as correct, but any device-less CUDA call it makes — a bare
``torch.cuda.synchronize()`` is the common one — targets cuda:0 and
INSTANTIATES that device's primary context. That context is hundreds of MiB
of permanently-resident VRAM on a card this rank does not serve from, and at
TP>1 it lands on a peer rank's GPU, which books it as
``driver.foreign_process`` and sizes its KV pool around it.

Pinning at this seam rather than inside each task is what makes the
guarantee hold: the ambient device is a property of the THREAD, so the place
that hands work to a thread is the only place it can be set once for
everything that runs there. The same pin, for the same reason, is
:func:`~arbi_serve.realtime.thread_device.device_guard` on the
``asyncio.to_thread`` seam.

Publish every recorded task to the metrics surface.

Kept off the phase gauge on purpose: an overlapped task's seconds run
CONCURRENTLY with a phase's, so adding them to the phase map would
make the phases stop summing to the boot.

Typed container for the engine's boot-computed sizing/profile state.

The boot path (:mod:`arbi_serve.engine.profile`, :mod:`arbi_serve.engine.build`,
:mod:`arbi_serve.engine.capture_sizing`) measures the model's activation peak,
sizes the paged-KV pool, and predicts the persistent cudagraph-pool VRAM
budget. Those results are produced in one phase and read by later phases
(active-state build, capture sweep, the Phase-2 freeze reconcile, the in-process
capture/grow path) and by the admin/metrics layer.

:class:`EngineBootState` holds that state in one typed place instead of a bag of
ad-hoc ``eng._*`` attributes. It is mutable: boot phases progressively populate
and refine the fields (e.g. the profiled page count is clamped after capture,
the graph-pool budget is replaced with the measured value once the sweep runs),
so the container is not frozen. Every field has a default so a partially-booted
or non-paged-KV engine reads a well-defined value rather than ``AttributeError``.

The ONE block-table width, monotone across the boot.

Three things must agree on it: the persistent ``PiecewiseBuffers`` table,
the TKV page-metadata scratch pre-sized against it, and the captured graphs
that read both. They are sized at different points of the boot, and the
pool ceiling they derive from is de-rated in between by the activation-arena
profile — so a later consumer re-deriving the width gets a NARROWER table
than the scratch was built for. The realized pool then maps past that table
and ``max_context=auto`` widens the served window onto pages no captured
graph can address: the request is admitted, pays its whole prefill, and
raises in ``CapturedGraph.replay``.

So the width only ever grows. The first caller records it; later callers
take the recorded one when it is wider than what they would derive.

The one KV pool ceiling, tolerant of duck-typed ``boot_state`` stubs.

See :meth:`EngineBootState.kv_pool_ceiling_pages`. Callers and tests pass
bare namespaces, so fall back to the raw fields instead of requiring the
accessor to exist.

The ONE ceiling: the most pages this KV pool can ever hold.

There were two, and they disagreed. ``kv_serving_ceiling_pages`` is
FORECAST-derived (it subtracts a profile-time serving floor whose
caching-overhang term is not measurable until the Phase-2 freeze);
``kv_va_ceiling_pages`` is what the growable slab actually RESERVED
address space for. Different consumers picked different ones: the grow
capped at the forecast while ``active.py`` reserved VA at the generous
value, so the pool could never use address space it already held —
a pessimistic guess silently costing KV that was physically free.

VA. So: one
accessor, the max of the two, used by the grow, the block table and the
spec-decode verify buffers alike.

The forecast keeps the job it is actually good for — the pre-capture
viability gate — and stops bounding a quantity measured directly after
capture.

Wait for a departing process's VRAM to drain before the KV pool is sized.

``docker restart`` and ``docker compose up -d`` recreate a container without
waiting for the previous process to release the card. The new boot then sizes
its KV pool against whatever VRAM is free at that instant, which is a fraction
of the card, and the pool is locked before the predecessor finishes exiting —
so the server comes up serving a fraction of its context and never recovers.

A co-resident sibling is legitimate and must not be refused, so "wait for the
card to be empty" is the wrong gate. The distinguishing signal is direction: a
departing process's footprint FALLS, a sibling's is STABLE. This waits for
foreign residency to stop falling, then lets the boot size against a settled
card. It never refuses; it reports what it waited for and what remains.

Telling the two footprints apart is the whole problem, and NVML's per-process
walk is not enough on its own: it reports pids in whatever namespace the driver
picks, so inside a container it can report the HOST's, no entry ever matches
``os.getpid()``, and the walk attributes this process nothing at all. This
module therefore also takes a mark BEFORE the process creates its CUDA context
(:func:`mark_pre_context_residency`) — everything resident at that instant is
someone else's by construction — and derives the split from it when the walk
cannot. That mark is the one source of the ours-vs-theirs split for the whole
boot; the driver-residency brackets read it from here.

Bytes THIS process holds on ``device_index`` per NVML, or None.

Two ways to find ourselves in NVML's per-process walk:

  * our pid appears in it — the direct match, and the only one that works
    with several tenants on the card;
  * the walk lists ONE process and accounts for the whole device — then
    that process is this one, because every caller of this function already
    holds a context on the device and must therefore be in the list,
    whatever pid the driver reports it under.

The second is not a nicety. NVML reports pids in whatever namespace the
driver picks, and inside a container that can be the HOST's, so the direct
match never fires there and a card with nothing else on it would otherwise
be unsplittable.

None when neither holds, which is not zero: without our own share the device
total cannot be split at all.

Bring NVML up, without dragging the metrics stack in with it.

The pre-context mark runs at the process entry point, ahead of every heavy
import, so it initializes NVML directly rather than through
:func:`~arbi_serve.server.nvml_metrics.init_nvml` — that module pulls
OpenTelemetry and the distributed state with it, and a mark that has to run
FIRST cannot afford to decide when those are imported. ``nvmlInit`` is
reference-counted, so the metrics stack's own init later is unaffected.

Device residency per NVML, on ``cudaMemGetInfo``'s basis. None if absent.

``nvmlDeviceGetMemoryInfo``'s original struct folds the driver's own
RESERVED pool — hundreds of MiB — into ``used``; ``cudaMemGetInfo`` does
not. Version 2 of the struct splits ``reserved`` out and its ``used``
matches ``cudaMemGetInfo`` exactly, which is what every consumer here needs:
both readings are subtracted from each other, and mixing the two bases
would put the whole driver reserve on whichever side of the split came from
NVML. Returns None rather than the v1 number when v2 is unavailable — a
reading on the wrong basis is worse than no reading.

Whether this process has already created a CUDA context.

Checks ``sys.modules`` before touching torch, so calling this from the CLI
entry point does not pull the torch import forward: a process that has not
imported torch cannot hold a context created through it.

Mark every device NVML can see. Returns how many marks were taken.

Called from the CLI entry point, which is the earliest moment in the
process that is guaranteed to precede the CUDA context — the engine
constructor is not: by the time it runs, something on the import path has
usually made a CUDA call already, and a mark taken then contains our own
context and splits the card wrong.

The device is not known that early, so every device is marked and the
engine reads back the one it lands on. Indices are NVML's, matching the
convention the rest of this module and the per-process walk already use.

Record device-wide residency BEFORE we create a CUDA context. Bytes/None.

Called once per device, from the engine constructor, immediately ahead of
the primary-context creation. At that instant this process holds nothing on
the card, so the reading is FOREIGN by construction — a split of the card
that needs no per-process counter and cannot be defeated by a pid the
driver reports in a namespace that is not ours.

Read through NVML rather than ``cudaMemGetInfo``, which would create the
very context the mark has to precede. Idempotent: the first mark for a
device wins, so a re-entered build cannot overwrite a true pre-context
reading with one taken after our own allocations landed. Refuses to record
anything once this process has a CUDA context, for the same reason.

One device's residency, split into ours and every other process's.

Bound to a device AND to a span — a single settle wait. Our own share is
derived once and held for that span, because the wait's entire signal is
the direction of the other tenant's footprint: our own is the constant (no
weights are loaded yet) and theirs is what moves, so re-deriving ours from
each poll's device total would track the card instead of the tenant and no
fall would ever be visible.

builds more than once — a hot model swap rebuilds the engine, and every
stable-VA member and runtime resident is another ``build()`` — and each
build runs this gate again with a footprint of its own. A share memoized at
the first boot and reused would book this process's own weights as a
departing process's bytes, and the gate would wait out its own residency
while a live server's model switch stalls on it.

The three readings are injected so the split can be exercised where there
is no NVML and no card at all, which is where it has to be checked before
it reaches one. Each defaults to this module's reading for the bound
device.

Block until foreign VRAM stops falling. Returns what was observed.

Reports ``settled`` (the wait ended because residency held steady),
``waited_s``, and the first/last foreign readings. ``settled`` is False on
timeout or when the residency cannot be read — the boot continues either
way, loudly, because refusing would strand a card whose sibling never
departs.

The ours-vs-theirs split is built HERE and discarded with the wait, so a
later build in the same process derives its own share against the footprint
it actually has rather than the one the first boot had.

Our own residency per the pre-context mark. Bytes, or None.

``device_used_now − foreign_at_context``: everything the card has
gained since we created our context is ours. Valid only while the OTHER
tenants' footprint is what it was at the mark, which is the situation
this span runs in — the mark and the first poll are a context creation
apart.

``device_used`` lets the caller pass the device reading it already
took, so the derivation and the reading it is subtracted from are the
same instant. Held after the first success, for this span only.

Bytes held on this device by processes OTHER than this one.

Device-wide residency less this process's own share. NVML enumerates
only the processes visible in the CALLER's pid namespace, so a sibling
CONTAINER's entry is absent from that walk entirely and no per-process
sum can see it; subtracting our own share from the device total can.

Our own share comes from NVML's per-pid walk when it can attribute us,
and from :meth:`own_bytes` when it cannot — which is not an edge case:
NVML reports pids in the namespace the DRIVER chooses, and where that
is the host's, no entry ever matches our ``os.getpid()`` and the walk
attributes us nothing at all.

None when neither split is available, which is not the same as zero: a
caller must not read "no other process" from "cannot tell".

Procedural engine boot.

Walks the boot dependency order as a straight-line procedure with no
manifest/artifact/content-addressed-cache indirection.

The boot body (:func:`_build_impl`) and the sync/async entrypoints
(:func:`build` / :func:`abuild`) live here, together with the cudagraph
capture-region helpers (:func:`_no_gc_during_capture`,
:func:`_stable_va_namespace`). The per-phase helpers live in sibling
modules and are re-exported below so every existing
``arbi_serve.engine.build`` import site keeps working:

  - :mod:`arbi_serve.engine.build_phases_load`   — pre-flight, backend
    registry, model load, head-quant.
  - :mod:`arbi_serve.engine.build_phases_kv`     — KV profile/size, LoRA
    + MTP, activation arena, deferred KV resize.
  - :mod:`arbi_serve.engine.build_config_resolve` — auto max-batch/context
    resolution, calibration/generation defaults, the boot-config banner.
  - :mod:`arbi_serve.engine.build_verify`        — synthetic-forward
    warmups, the readiness gate, phase measurement.
  - :mod:`arbi_serve.engine.build_memory_sizing` — capture-sweep / KV-floor
    / post-capture VRAM headroom assertions.
  - :mod:`arbi_serve.engine.build_graph_pool`    — cold-boot cudagraph
    pool sizing + measured-budget persistence.
  - :mod:`arbi_serve.engine.profile`      — activation profile + KV sizing.
  - :mod:`arbi_serve.engine.active`       — pool / page-table / scheduler
    / metadata-builder / attn-op build.
  - :mod:`arbi_serve.engine.tkv_install`  — TKV codec install + calibration.
  - :mod:`arbi_serve.engine.mtp_attach`   — MTP driver build + attach.
  - :mod:`arbi_serve.engine.build_helpers` — cold/warm flat-dump model
    load + post-load compaction.

Sync entry: :func:`build`. Async entry: :func:`abuild` (await-able
from inside an existing event loop, e.g. the FastAPI lifespan;
trivial async wrapper since boot is fully synchronous).

Log a ``[boot-profile]`` span for one step inside a boot phase.

A phase line reports the phase; a phase that hosts several independent
steps (the capture sweep's post-sweep KV grow, budget persist, reclaim and
freeze) needs the steps named too, or a regression in one of them is only
visible as a slower phase.

Tag-namespace context for the boot model when stable-VA residency is on.

When ``cfg.stable_va_residency`` is set, every cuMem allocation this boot
makes is stamped ``"<model-key>/<pool>"`` so the residency controller can
later park / wake exactly this model's physical pages (keeping every other
model's VA reservations alive — the disjoint-arena guarantee). A no-op
context otherwise, so the default single-model path is byte-identical.

Warm ``torch._dynamo`` / ``torch._inductor`` on a boot thread.

These pure-Python imports take a noticeable fraction of a second.
Kicking them off here lets them run concurrently with the GPU/IO-bound
weights-DMA / KV-profile / pool-sizing phases instead of serially in
front of Phase 4b's ``precapture_compile_warmup()``. Pure prefetch: any
failure is swallowed and the warmup re-imports synchronously (Python's
per-module import lock makes the concurrent import race-free — a
main-thread import blocks until this one finishes rather than
double-importing).

Recorded through :class:`~arbi_serve.engine.boot_overlap.BootOverlap`
and joined by :func:`_join_compile_dep_prefetch` at the head of the
compile phase. The join is behaviour-neutral — the first
``import torch._dynamo`` inside the warmup already blocks on this
thread's import lock — and it is what makes the prefetch's
``overlapped`` seconds a measured number instead of an assumption.

Do not add ``arbi_serve.models._gdn_kernels`` (fla) prefetch here:
prefetching an ``arbi_serve.models`` submodule on a daemon races the
main thread's ``load_model`` import of the same package and can raise
``ImportError: cannot import name '_gdn_kernels'`` (partial-init of a
package being imported concurrently). ``torch._dynamo``/``_inductor``
are safe because the main thread doesn't concurrently import them
mid-package.

Settle the overlapped weight-fill worker and retire its pool.

Idempotent; a worker that raised is settled the same as one that
returned — the caller is already unwinding a build error.

Say loudly, once per build, that the per-step timing profiler is on.

``--timing-debug`` / ``ARBI_TIMING_DEBUG`` instruments every request on
the decode path and holds its records until the request finishes, so a
server left with it on serves measurably slower and holds memory in
proportion to output length x concurrency. Nothing an operator normally
wants — the swimlane, the request timeline and the latency histograms
read the always-on lifecycle stamps — so an operator who turned it on
for a debug session must be able to find it in the log rather than in an
unexplained regression. WARNING level on purpose: a debug-level line is
exactly the line nobody reads.

Synchronous procedural boot.

Runs the full cold-boot sequence — pre-flight, model load,
tokenizer + xgrammar, profile, active-state build, LoRA + MTP
attach, activation arena, cudagraph capture sweep, post-capture
VRAM gate. Sets ``eng._ready = True`` on success.

Wait for a departing process's VRAM to stop falling. Seconds waited.

Returns 0.0 when the card is clear, when the gate is disabled
(``ARBI_BOOT_VRAM_DRAIN_TIMEOUT_S=0``) or when the residency cannot be read
— the wait is a guard on the KV sizing, never a boot blocker. The verdict is
stashed on ``eng.boot_state.foreign_vram_wait`` for the boot ledger.

Turn the optimistic output advance off for a host-row model.

Qwen4-Exp's Per-Layer Embedding prefetches its n-gram rows from host
memory using each token's local context, so the host has to KNOW every
fed id before the forward that consumes it. The async-output advance
leaves the sampled id device-only and the next row's context
unnameable — the architecture's own requirement, decided once here
rather than surfacing as a failed step per request.

Refuse a KV page size the loaded checkpoint's attention cannot use.

A model whose attention addresses fixed-size blocks through the page
table — Qwen4-Exp sparse attention, whose index block IS one page —
declares ``required_cache_block_size``. Checked at boot, where the remedy
is one flag, rather than at the first forward.

The per-layer counters the accept-commit warm must move on THIS config.

Returns ``(witnesses, unregistered)``: the counters whose fire count the
warm is expected to raise, and the names that name a path this config
routes through but whose registration site has not been imported — which
is a check that could not run, never a pass.

Derived from the same knobs the launch sites read, so a boot that legally
routes elsewhere (a pool with no GDN layers, either fused arm switched
off) expects nothing from that arm rather than reporting a false miss.

Whether any recurrent view already holds a captured replay graph.

The graph is per KIND view (``pool._views[StateKind.GDN]``), not on the
pool facade that exposes ``rollback_partial_accept``; a read of the
facade is always ``None`` and would silently turn this witness off.

Pre-compile the fused GDN accept-commit, and prove it covered it.

One fused accept-commit launch compiles + caches the pointer-table plan
against the REAL slab/snapshot pair set. Row 1 (the first allocatable row
— row 0 is the guarded zero-sentinel and the pool refuses it) is free at
boot, and the readiness gate re-zeroes every recurrent slab after this
phase, so the warm write cannot leak state into serving.

Both rollback modes warm through ``rollback_partial_accept``: in
``replay`` the eager pass JIT-warms the kernels AND the pool captures its
side cudagraph HERE, so the one-time capture sync never lands on a live
serving tick; in ``recompute`` it JIT-warms the host-side replay kernels.
(The post-warmup slab re-zero only rewrites VALUES — the captured graph's
``data_ptr`` fingerprint is untouched.)

The witness counters are read either side of the call because "the warm
ran" and "the warm covered the kernels the first live commit launches"
are different claims, and only the second one is worth a green: a warm
that routes past a launch site leaves that kernel to compile inside the
first request, where ``jit_compile_serving`` reports it — after the
latency has been paid.

Async wrapper around :func:`build`.

Boot is fully synchronous (CUDA-bound) — the async surface exists
so callers already inside an event loop (FastAPI lifespan, the
distributed driver's startup coroutine) can ``await eng.abuild()``
without ``asyncio.run`` complaining about a nested loop.

Pause Python GC for the duration of the cudagraph capture sweep.

A ``torch.cuda.MemPool`` that became unreferenced on a PRIOR teardown (the
classic case: a hot-swap reload's previous-model ``graph_pool``, parked in
``_LEAKED_POOLS`` then unpinned when its captured graphs were cleared) is
collectable by the time the NEXT model's build runs its capture sweep. If
Python GC fires mid-capture and destructs that pool, ``~MemPool`` calls
``synchronize_and_free_events``, which aborts the WHOLE process with
``captures_underway.empty() INTERNAL ASSERT FAILED`` (a cuda graph is
recording on the stream) — the swap-N≥2 crash. Disabling GC for the capture
window makes that impossible by construction; on exit we restore the prior
GC state and run one collection at a SAFE point (no capture underway) so any
deferred ``~MemPool`` runs cleanly. Cheap + boot/reload-only.

Close the open boot window and bill it to ``name``.

The windows telescope — each starts where the previous one closed —
so the recorded phases sum to the boot. ``charge`` moves seconds off
``name`` (work it hosted but does not own); ``charge_to`` names the
phase they move ONTO. Pass both together: a deduction without a
destination leaves the sum, which is the thing the sum exists to
catch.

Close a VRAM bracket WITHOUT a wall-clock phase line.

The capture sweep is one wall-clock phase but several distinct driver
consumers (one cubin set per capture family, plus the instantiate
that named nothing, which is the granularity problem this ledger exists
to fix — so the sweep is bracketed per family here.

Run one capture family inside BOTH its instantiate-meter scope and
its own driver-residency bracket, and file the bracket on the meter.

The two spans have to be the same span. A family that meters zero
instantiates is either dormant (its sweep is gated off, or its
duck-typed binding is absent on this model) or captured through a path
the meter cannot see, and the ONLY thing that separates those is
whether the driver grew while the family ran. Deriving that from two
constructs written side by side is how they drift; deriving it from one
is how they cannot.

Adds no sample of its own — it reads the bracket ``_mark_vram(label)``
already closes, so the per-phase ledger line is unchanged.

Announce a phase before its (silent, possibly long) work begins.

The matching ``_phase_done`` reports the MEASURED elapsed once the
phase finishes; this start line just marks that the phase is running
so a live tail isn't an unexplained gap. No time estimate is printed
— the completion line carries the real number.

A heartbeat runs for the duration so a phase that does minutes of
silent work is distinguishable from one that is stuck.

Log a compact banner of the resolved boot config + non-default flags.

Two blocks: the engine config a reader most wants to see (model, backend,
memory, batch/context, MTP, parallelism) and every ``ARBI_*``/``TKV_*``
runtime flag whose value differs from its default (so the banner shows
exactly what was overridden, not 60 lines of defaults). All values are
read defensively — the caller wraps this in ``suppress(Exception)`` so a
missing attribute can never abort boot.

Enforce the calibration policy for an active TKV ``PAGED_KV`` backend.

When the active PAGED_KV backend is TKV and no calibration resolves
for it: raise ``ValueError`` unless ``cfg.cache.allow_uncalibrated_tkv``
is set, in which case log a warning and return (the codec runs at
default centroids — lower quality). No-op when the active PAGED_KV
backend is not TKV or a calibration is present.

``ValueError``, never ``SystemExit``: ``build()`` also runs under
``areload_model`` on a LIVE server, and ``SystemExit`` is a
``BaseException`` that the admin dispatcher's ``except Exception``
cannot see — the operator got a bare ``500 Internal Server Error``
naming no cause. ``ValueError`` is in ``/v1/admin/model``'s error map.

Load the model's sampling defaults + log which will apply.

Stashes a :class:`~arbi_serve.generation_defaults.GenerationDefaults`
on ``eng.generation_defaults`` for the request path. Under
``generation_config_mode == "none"`` the file is NOT read and an EMPTY
(apply-nothing) result is stashed — the request path then keeps the
server's hardcoded defaults. LOUD either way: the boot log states
exactly which sampling defaults will backfill an omitting request.

Parse a TKV calibration JSON, failing LOUD and ACTIONABLE at boot.

A missing/unreadable/malformed calibration never silently falls back to
codec defaults (a silent tkv→uncalibrated path is a known footgun). The
two common operational failures get an actionable message instead of a
raw traceback:

- **missing file** (e.g. the calibration volume/NFS mount did not come
  back after a host reboot) — report the path, the likely cause, and list
  what *is* present in the parent directory so the operator can see the
  real available calibrations (or that the mount is absent entirely).
- **truncated/invalid JSON** (a ``kill -9`` mid-write) — say so explicitly
  and point at regenerate/restore.

``model_path`` is the checkpoint this engine is booting. Given one, the
bundle's ``fingerprint`` stamp is compared against it — the only check
that asks whether the bundle was fitted for THIS engine rather than
merely being a well-formed bundle. This is the single door every load
goes through (boot AND the live ``/v1/admin/calibration`` reload), so
the identity guard cannot be present on one path and absent on the
other.

Once the file exists and parses, BUNDLE-level validation is tkv's:
:func:`tkv.runtime.calibration.load_calibration_file` runs the
schema-version gate, the finalized / byte-budget-coverage checks, the
peer-scale guard and the per-layer-quantizer guard. The two pre-checks
above stay here because that loader warns and returns ``None`` on a
missing or unparseable file (vLLM bootstraps one); on arbi a missing
calibration is a boot failure.

Resolve ``cfg.cache.max_context == "auto"`` to the model ceiling.

The model dims are available (weights loaded) but profiling has not run,
so the concrete value is the model's trained context window
(``max_position_embeddings`` — the RoPE table size). This is the static
cap the engine sizes capture/admission against; runtime memory only
narrows the *advertised* per-request window downward
(``derive_effective_max_context``), never the engine's ceiling.

Frozen-dataclass note: ``cfg.cache`` is frozen, so we write the resolved
int via ``object.__setattr__`` (mirroring how the build mutates other
derived engine state in place rather than rebuilding the whole cfg).

Gate the configured ``max_batch`` against the KV pool.

``max_batch`` is a KNOWN quantity — the config default or the operator's
``--max-batch N`` — so this is a GATE, never a resolution: the number is
already fixed before the engine profiles anything, and the capture-time
pools pre-size to it. If the profiled pool cannot hold it at the
memory-realistic per-seq window, serving fewer rows (a quiet clamp, or a
downstream ``TQBufferPool.ensure`` cap) would honour the config on paper
while degrading concurrency in the dark.

Two outcomes when it does not fit. When the context is the engine's to flex
(``--max-context auto``, or a recipe default) the context NARROWS to the
largest window at which the full batch fits and the boot continues, loudly.
When the operator pinned the context, the boot refuses, naming the
configured batch, the count that fits, and the limiting budget term.

TWO PASSES, and only the second may refuse on an estimate. ``realized=False``
(pre-capture, from ``_profile_and_size_kv``) reads
``boot_state.profiled_num_pages`` — a FORECAST cut from the predicted
non-KV pools, one of which is the graph-pool reserve. When that reserve is
the closed-form COLD ESTIMATE rather than a measured reading
(``boot_state.graph_pool_budget_is_estimate`` — every budget-cache miss,
which includes every capture-affecting live config override), the forecast
is a lower bound the post-capture grow routinely multiplies, so a refusal
here would reject a config the engine goes on to serve. The gate then WARNS
and defers. ``realized=True`` (post-capture, after ``grow_kv_after_capture``
+ ``finalize_served_max_context``) reads the pool's own mapped page count
and refuses on that — truth, not a forecast. The realized pass does not
narrow the context: ``finalize_served_max_context`` /
``assert_kv_pages_floor`` already narrowed against the same realized pool.

Skipped for non-paged-KV models (no profiled page count) and for the
``--num-pages`` override path: the operator who hand-sizes both knobs owns
the trade-off (the override already fail-loud-validates against the budget
in ``profile_and_size_kv_pool``).

Boot-time guard: ``max_context`` vs the model's RoPE table size.

The RoPE cos/sin tables are built to
``model.dims.max_position_embeddings``
(:class:`arbi_serve.models.layers.RotaryEmbedding`); indexing
``cos[positions]`` past that row count is an out-of-bounds gather
(an ``IndexError`` on CPU or a device-side assert that poisons the
CUDA context on GPU). ``cfg.cache.max_context`` (recipe / CLI
default 32768) is set INDEPENDENTLY of the model — admission only
rejects prompts ``>= max_context`` (``_build_request``), so a
``max_context`` larger than ``max_position_embeddings`` lets a long
request walk the RoPE table OOB.

Hard-REJECT at boot (``ValueError``) rather than clamp: the config
dataclasses are frozen (clamping ``max_context`` in place would need
a full ``cfg`` rebuild) and a context window the model cannot serve
is an operator mis-config best surfaced loud at boot, before any
request, with an actionable message (lower ``--max-context``). Set
``ARBI_ALLOW_OVERSIZE_CONTEXT=1`` to downgrade to a warning and boot
anyway (the server then crashes only if a request actually reaches
``max_position_embeddings`` — useful when the operator knows their
prompts stay short).

(``chunk_prefill`` vs the prefill cudagraph ladder is handled upstream
in :func:`precapture_prefill_graphs`, which auto-extends the ladder to
cover ``chunk_prefill`` via :func:`_prefill_buckets_covering_chunk`, so
there is no eager-prefill cliff to guard here.)

Set ``dims.rope_cache_seq_len`` to the largest position any forward reaches.

Stores ``min(max(max_context, max_batched_tokens), max_position_embeddings)``
on the model's (frozen) :class:`ModelDims` so the per-arch RoPE-cache
ctors and the ``rope_cache_pool`` budget predictor build the cos/sin
tables to the same row count. The cache must cover every position a
forward gathers: an admitted request reaches ``max_context``, and the
boot activation profiler's synthetic single-sequence prefill reaches
``max_batched_tokens``. Capped at the model's trained window so it is
never larger than the rotary table actually spans.

Under ``--max-context auto`` this runs while ``max_context`` is still the
MODEL CEILING (:func:`_resolve_auto_max_context` resolved it there), and
that is the correct row count to build — not the narrower window the boot
eventually serves. The served window is not knowable here (it comes out of
the KV budget, which is sized by an activation profile these very tables
must already exist for), it varies boot to boot with realized VRAM, and it
moves in BOTH directions during the build: the pre-capture gate narrows it
and the post-capture grow widens it again against the realized pool
(``build_memory_sizing``). Sizing the tables to a narrowed window would
make a later widen gather past their last row — a silent wrong-position
bug, not a crash — and by then the tables' storage is baked into captured
graphs and cannot be re-grown. The ceiling is the only row count correct
for every window the engine can end up serving; the rows the realized
window never reaches are the price of that guarantee.

Refuse a boot whose grouped MoE dispatch would overrun ``MGEMM_MAX_SLOTS``.

``exl3_mgemm``'s expert-range filter carries a fixed slot list, so a
shard-filtered call is capped at ``MGEMM_MAX_SLOTS`` slots and the kernel
``TORCH_CHECK``s above it. Only the grouped (slot-major) leg filters, and
only under expert sharding; ``moe_expert_major_min_rows`` is what decides
how wide a step reaches it. The default (4) keeps that to three rows, but
the flag is operator-settable and ``0`` sends every prefill step there,
which trips the check partway through capture as a bare kernel error.

Runs AFTER the weights future joins: ``exl3_bound`` is what separates a
model whose routed experts reach ``exl3_mgemm`` from one that takes the
per-expert walk, and it only answers once the packs have bound.

DFlash drafter build + attach.

Mirrors :mod:`arbi_serve.engine.build_external_drafter`: loads the
DFlash draft checkpoint, runs boot asserts against the main model,
arms the target model's hidden tap, and wires the verify-pass
scaffolding (verify buffers + recurrent snapshot buffers + lifecycle
hooks). Routed from :func:`arbi_serve.engine.mtp_attach.build_mtp_driver`
when ``cfg.mtp.dflash_draft_path`` is set.

Rank>0 placeholder drafter — holds sizing scalars, no model weights.

Under TP>1 the DFlash draft model + per-request draft-KV state live
only on rank 0; worker ranks never call :meth:`DFlashDrafter.draft`
(they only join the ``embed_tokens`` / ``lm_head`` collectives via
``DFlashDraftOp``). But the verify scaffolding shared by every rank
(``build_for_engine`` → ``VerifyBuffers``, snapshot-buffer sizing)
reads ``mtp_driver.max_k`` off the installed drafter, and
``eng.mtp_driver`` (a forwarding alias for ``eng.drafter``) must
stay non-``None`` and symmetric across ranks. This stub supplies
exactly those scalars without loading the draft checkpoint on the
worker.

Whether the TARGET's ``config.json`` declares a RoPE scaling schedule.

``None`` when it cannot be read — the draft head then keeps whatever its own
checkpoint declares, since an unreadable target proves nothing.

Resolve :func:`_target_declares_rope_scaling` from a server config.

``ServerConfig.model`` is a ``ModelConfig``, not a path — reading it as one
yields ``None`` (unknown) and leaves every draft head's declaration intact,
so the reconciliation never fires and nothing reports that it did not.

Why this boot keeps the draft forward eager, or ``None`` when it compiles.

The rule behind ``ARBI_DFLASH_COMPILE``'s default-ON: the drafter rides the
verifier's compile path, so it compiles exactly when that path is armed
(``compile_on`` wires the Inductor caches) and the drafter is on CUDA (the
trampoline gates on the piecewise-cudagraph context, which has no CPU
leg). Each string is the name ``dflash_compile_armed`` refuses under and
the boot line prints, so a boot that resolves OFF says so by name rather
than serving an eager drafter silently.

Build + attach the DFlash block-diffusion drafter.

Boot asserts (fail loud, before any state is accepted), read from the
draft ``config.json`` so every rank validates without loading the
model:
  1. main model exposes the DFlash tap (``_dflash_tap_ids`` seam);
  2. draft ``num_target_layers`` / ``target_layer_ids`` match the
     main model's decoder depth;
  3. draft ``hidden_size`` matches the main model's;
  4. a checkpoint-declared draft vocab matches the main model's. A
     checkpoint that owns neither embedding nor LM head may omit it and
     borrows the target vocabulary by construction.

TP>1: the draft model + per-request draft-KV state live on **rank 0
only**. Worker ranks install a :class:`_WorkerDFlashStub` (sizing
scalars, no weights) and join rank 0's ``embed_tokens`` / ``lm_head``
collectives in lockstep via ``DFlashDraftOp`` (see
:meth:`DFlashDrafter.draft`). The hidden tap, the verify scaffolding
(verify buffers + GDN snapshot buffers), and the greedy config
mutation run on **every** rank so the per-rank target forwards stay
in step — a rank-0-only tap would diverge rank 0's captured target
graph from the workers' and desync the vocab-parallel collectives
during the capture sweep.

No-op accept-counter commit on a worker rank.

Rank 0's real
:class:`~arbi_serve.spec_decode.dflash_driver.DFlashDrafter` owns
the lifetime aggregate + the per-request counters; the worker holds
no draft state, so it absorbs the call rather than AttributeError-ing
on an unconditional commit and wedging rank 0 on the next collective.

External draft-model drafter — boot wiring.

Builds an :class:`ExternalModelDrafter` end-to-end:

  1. Resolve drafter dtype (override via ``cfg.mtp.draft_model_dtype``
     or inherit the main model's).
  2. Hash both tokenizers (drafter dir + main dir) so the boot assert
     can verify they match.
  3. Load the drafter model via ``load_model`` — same factory as the
     main model — through the warm flat-dump cache
     (:mod:`arbi_serve.engine.drafter_flat_cache`), so a boot with a
     published dump never re-reads the drafter checkpoint.
  4. Build a per-drafter :class:`MultiStatePool` sized for
     ``draft_model_max_seqs × (cache.max_context + max_k)``. Bytes
     land in the engine's existing ``kv_pool`` named bucket so
     /metrics doesn't spawn a parallel column.
  5. Build a :class:`FlatPageTable` against the drafter pool +
     ``pool.bind_page_table(...)``.
  6. Build per-layer :class:`AttnOp` via the same factory the main
     engine uses (``backend.make_attn_op(...)`` per layer).
  7. Build per-:class:`StateKind` :class:`MetadataBuilder` for the
     drafter (PAGED_KV only — boot assert refuses other kinds).
  8. Construct :class:`ExternalModelDrafter` (which fires its own 8
     boot asserts).
  9. Hook engine's ``_on_admitted`` / ``_on_finished`` so the drafter's
     page table tracks the engine's active set.
 10. Run ``warmup()`` once so first-call JIT cost amortizes off the
     first request's latency.

A captured ``cfg.cuda_graphs`` (see step "captured-chain pool" below + the
precapture sweep in :mod:`arbi_serve.engine.cudagraph_admin`).

Current limitations: no sleep/wake integration for the captured chain
pool, no multi-group teardown.

Sha256 over the raw ``tokenizer.json`` bytes in ``model_dir``.

Sharing :class:`tokenizers.Tokenizer` instances across drafter +
main is risky (added_tokens metadata is mutable state); the hash
is the cheapest fail-fast that catches "different vocab merge
order" / "different special tokens" without actually touching the
tokenizer object.

Returns the hex digest. Raises :class:`FileNotFoundError` when the
file is absent — the boot path expects standard HF layout.

Resolve drafter dtype from cfg + return the explicit-override flag.

Returns ``(drafter_dtype, override_dtype_or_none)``:
  - ``drafter_dtype`` — the dtype to load drafter weights into.
  - ``override_dtype_or_none`` — the same dtype iff the operator
    explicitly set ``cfg.mtp.draft_model_dtype``, else ``None``
    (which signals to the boot assert that drafter dtype must
    match main dtype).

Build one :class:`AttnOp` per drafter layer.

Mirror of the per-layer construction in
:func:`arbi_serve.engine.build.build_active`, but bound to the
drafter's own backend + layer_specs. PAGED_KV only — the boot
asserts refuse MLA / Mamba / GDN / ShortConv drafter archs.

Build + attach the external draft-model drafter.

Caller (`build_mtp_driver`) has already verified
``cfg.mtp.draft_model_path is not None``. Sets ``eng.drafter``,
hooks admit/finish callbacks, runs warmup, and returns.

Failures bubble up as :class:`RuntimeError` with actionable boot
messages — no silent degradation.

Wire ``eng._on_admitted`` / ``eng._on_finished`` to call the drafter.

The engine's existing free-function hooks (``on_admitted`` /
``on_finished`` in :mod:`arbi_serve.engine.run_step`) bump
/metrics counters; we wrap them so admit/finish ALSO drives the
drafter's per-request state.

Hook is attached as a bound method override on the engine
instance — the engine class itself is unchanged. Idempotent: a
second attach is a no-op (we tag the engine instance).

Canonical spelling of the active draft-tree geometry, "" for the chain.

Resolved through :func:`~arbi_serve.spec_decode.tree_spec.active_tree_spec`
— the same resolver every sized-per-node buffer reads — so the cache key
and the buffers it prices can never disagree about which geometry is live.

The NVIDIA driver release, ``""`` when NVML cannot answer.

Cubin load sizes and the per-graph ``cudaGraphInstantiate`` cost are
driver-dependent, so a driver upgrade must fork every measurement kept
under the budget-cache identity. NVML needs no CUDA context and answers
the same on every rank.

``{param: value}`` for every capture-affecting param in the registry.

Enumerated from :data:`~arbi_serve.config_overrides.PARAMS` — the one
place a param is classified as capture-affecting — never a list kept
here. RuntimeFlags params are read through
:func:`~arbi_serve.runtime_flags.runtime_flags`, so the active member's
overlay is what keys the cache; config params are read off ``cfg``. These
decide which kernels and cubins load and which graphs are captured, which
moves the graph pool, the instantiate cost and the driver-module residency
the identity guards. A value ``cfg`` cannot resolve keys itself as
``"unresolved"`` rather than aliasing a real value.

Build the ``cache_key_inputs`` dict for ``predict_graph_pool_bytes``.

The cache key components are: model_id, tp_size, kv_backend,
shape_buckets, the device identity (``gpu_arch`` — compute capability —
and ``gpu_name``, because two cards of one arch differ in SM count and so
in local-memory and module residency), ``driver_version``, the
kernel-library identity (``torch_version`` — which carries the CUDA build
tag — and ``turbo_attn_version``), ``capture_affecting_flags`` (see
:func:`_capture_affecting_values`), and ``profile_cache_version`` (an
explicit bump knob for changes to the capture/budget logic itself). The
key is keyed on the inputs that actually move the measured graph-pool
size, NOT the app/image git SHA, so the cached budget survives unrelated
code edits. Every measurement persisted under the budget-cache identity
(graph pool, instantiate cost, driver-module baseline, serving overhang)
reads this one dict off ``boot_state``, so an input added here forks all
of them at once.
Each input must be a JSON-serializable value so the SHA reproduces
across boots at the same configuration.

Per-family ``cudaGraphInstantiate`` bytes/exec + where they came from.

Returns ``({family: bytes_per_graph}, source)`` where ``source`` is
``"measured"`` (this configuration has been booted before and the rate came
out of the budget cache) or ``"seed"``. The dict always contains the
``"*"`` key — the rate for a family with no measurement of its own.

A family the measurement covers is priced by its OWN measured rate. This is
what settles the drafter question empirically instead of by assumption: a
drafter-chain graph bakes ``K`` steps of a ONE-layer head rather than a
backbone pass, so it *might* be cheaper per exec — but the node-count model
that would have predicted "much cheaper" is refuted (see
:data:`_COLD_SEED_PER_GRAPH_INSTANTIATE_BYTES`), and on this driver the cost
looks flat per exec. Rather than choose between two guesses, the meter
records each family separately and the cache remembers which one was right
at this configuration.

``(B, S)`` decode+verify shapes the capture sweep will instantiate.

``_decode_capture_shapes`` only auto-extends with the MTP verify
``(B, K+1)`` cross-product once ``build_mtp_driver`` has attached the
driver + verify buffers. Two of this module's callers run BEFORE that
(:func:`arbi_serve.engine.build_memory_sizing.assert_kv_context_window_fits`
via ``predicted_realized_kv_pages``, and the cold-boot graph-pool estimate),
so reading the raw helper there silently drops EVERY verify graph from the
count. Union in the config-derived verify set — the same escape hatch
:func:`cold_boot_graph_pool_upper_bound_bytes` already uses — so the count
is the same before and after driver attachment.

``{family: graphs the Phase-5 sweep will instantiate}``.

The boot sweep (``build.build_active``) instantiates EIGHT families. A
budget that open-codes ``len(_decode_capture_shapes(eng)) * bucket_mult``
counts the decode/verify family and nothing else, leaving seven free:

  * decode/verify  — ``(B, S)`` shapes × kv-page × LoRA buckets
  * drafter chain  — ``(B, K)`` buckets, MTP with a bundled/external head
  * dflash         — one draft forward per batch width, capture armed
  * prefill        — one graph per bucket, ``prefill_capture='full'``
  * mixed          — one graph per bucket, ``ARBI_MIXED_CAPTURE``
  * piecewise      — per-layer graphs, ``prefill_capture='piecewise'``
  * shared_kv      — ``(B, K)`` buckets, a shared-KV assistant head
  * media          — one graph per media-encoder bucket

The family KEYS are the same ones ``build.build_active`` scopes the
instantiate meter with, so prediction and measurement cannot drift apart in
vocabulary.

Best-effort per family: a family whose helpers raise contributes 0 rather
than wedging a boot, and every caller logs the breakdown so a zero is
visible rather than silent.

Driver ``cudaGraphInstantiate`` reserve for the WHOLE boot capture sweep.

Returns ``(total_bytes, {family: graph_count})``.

ONE source of truth for a term four call sites must agree on — the
pre-capture feasibility gate (:func:`assert_capture_sweep_fits`), the
deferred-KV holdout (``build_phases_kv``), the realized-KV forecast
(``predicted_realized_kv_pages``) and the boot KV ledger. Counts come from
:func:`predicted_instantiate_graph_counts`; the per-exec rate comes from
:func:`resolve_instantiate_bytes_per_graph` — MEASURED on a prior boot at
this configuration, or the labelled cold seed until there has been one.

None of the four consumers is a second PHYSICAL reservation of these bytes:
the gate only refuses, the deferred-KV holdout reserves free VRAM the sweep
then spends, the forecast drives the ``max_context`` auto-narrow, and the
ledger reports. The pool that serves is sized by ``grow_kv_after_capture``
against ``mem_get_info`` AFTER the sweep, by which point the driver has
already taken these bytes out of free. Getting the number right therefore
buys HONESTY (and a correct auto-narrow), not KV — and getting it wrong in
either direction costs: too low and the residual has to absorb the miss, too
high and a config is narrowed for memory that was never spent.

``cudaGraphInstantiate`` bytes this process has NOT yet allocated.

Returns ``(unpaid, metered, note)`` — what free VRAM still has to cover,
what the instantiate meter has already watched the driver take, and a
self-describing clause for the caller's log line and ledger row.

A ``cudaGraphExec`` lands on the default device heap, outside every cuMem
pool and outside the torch caching allocator, so nothing but free VRAM pays
for it and nothing but ``cudaMemGetInfo`` sees it. That makes WHEN it is
paid the whole question: a term already on the card must not be held out of
free a second time, and a term still to come must be.

Both halves are measured. :func:`predicted_instantiate_reserve` is the
sweep's total at this configuration (per-family counts x the per-exec rate
METERED on a prior boot here, or the labelled cold seed until there has been
one); ``instantiate_meter.total_bytes`` is what this process has actually
watched the driver take, and it is the same meter that publishes the
``driver.cudagraph_exec`` ledger row. Before the capture sweep the
difference is the whole reserve; after it, the bytes are on the card and the
difference is zero — which is the point, and the reason a caller that runs
AFTER the sweep must not hold anything.

The meter is process-wide and cumulative, so this is only meaningful to a
caller asking about the sweep that has just run or is about to. The
PRE-capture holdout (``build_phases_kv``) asks
:func:`predicted_instantiate_reserve` directly and must keep doing so.

One-line, self-describing breakdown of the instantiate reserve.

MiB [measured]`` — count, rate, and PROVENANCE, so a
reader of any boot log can tell a measurement from a seed without reading
this module. The acceptance bar is that every held-out term is named and
measured; a term that prints only its total cannot meet it.

Conservative PERSISTENT block-table + buffer floor of ONE captured graph.

This is the SHAPE-INDEPENDENT part of a captured graph's private pool: the
baked KV/page block-table (``max_blk`` int32 entries × ``max_batch`` rows)
plus a small bounded per-graph driver/buffer overhead. It does NOT include
the captured forward-activation working set — that term scales with the
graph's ``(B, S)`` token count and is added per-shape by
:func:`_captured_graph_working_set_bytes`. One helper, one constant, so the
KV sizer and the capture-fit gate cannot drift apart on the floor term.

Per-token elements live at the gated-MLP peak of ONE captured layer.

``GatedSiLUMLP.forward`` on the fused hot path is ``down(act(gate) * up)``
over one ``[gate|up]`` matmul
(:meth:`~arbi_serve.models.layers.GatedSiLUMLP.forward`). At the widest
instant THREE intermediate-width tensors are live: the fused ``gate_up``
output (``2 × intermediate``, still referenced by the ``gate`` and ``up``
chunk views), the activation ``act(gate)``, and the product it multiplies
into. That is ``4 × intermediate`` per token, TP-sharded because gate/up are
column-parallel and ``down_proj`` is row-parallel. The two-matmul and merged
forms allocate the same three tensors, so the count does not fork on the
gate/up strategy.

Per-token elements live at the peak of ONE captured attention layer.

``Hq = num_heads x head_dim``; ``Hkv = num_kv_heads x head_dim``; ``H`` is
``hidden_size``. Head widths are TP-sharded (column/row-parallel
projections); ``H`` is not. Each phase is an explicit sum of named widths
and the largest wins — the block does not hold them at the same time.

  * PROJECT — the previous layer's ``mlp_out`` (the ``pending_add`` this
    layer's fused input_layernorm consumes), the normed ``h``, and the fused
    QKV output. ``q_proj`` emits ``2 × Hq`` when the layer carries the
    packed output gate
    (:meth:`~arbi_serve.models.attn.AttentionBlock._project_gated`), ``Hq``
    otherwise; K and V add ``Hkv`` each, or one ``Hkv`` under ``tie_v_to_k``.
  * ROPE — the projection output is still live (Q/K/V are views into it)
    alongside the ``gate_flat`` reshape (a COPY: the gate is a strided chunk
    of the packed Q output), the Q/K-norm outputs, and RoPE's fp32
    temporaries (:meth:`~arbi_serve.models.layers.RoPECache.apply`).
  * OUTPUT — the attention output, the flattened gate, ``sigmoid(gate)`` and
    the product, then ``o_proj`` and the post-attention normed hidden
    (:meth:`~arbi_serve.models.attn.AttentionBlock._post_attn`).

The attention kernel's own scratch is NOT here: under
``split_attn`` the op runs eager between the captured pieces so its
allocations never reach the capture pool, and with the flag off it is the
backend's workspace, which is not a per-token width this module can derive.

Per-token elements live at the peak of ONE captured GDN layer.

Widths come from :class:`~arbi_serve.models.layer_spec.GDNConfig`, which
carries the checkpoint's own head counts and head dims. ``KD = key_dim``,
``VD = value_dim``, ``C = conv_dim = 2 × KD + VD``, ``Kc = conv_kernel``.

  * PROJECT — ``mlp_out`` + normed ``h`` + the fused QKVZ projection
    (``2 × KD`` for Q/K, ``2 × VD`` for V and the Z gate) and the ``ba``
    projection (two scalars per value head)
    (:meth:`~arbi_serve.models.gdn_block.GDNBlock.forward`).
  * MIXER, ``seq == 1`` — the packed decode kernel writes into a
    pre-allocated output and allocates nothing itself, so the widest instant
    is the normed hidden + the QKVZ output + that output + the ``flat_z``
    reshape, which is a COPY (Z is a strided slice of the fused QKVZ).
  * MIXER, ``seq > 1`` — the multi-token verify path rebuilds the causal
    conv by unfolding the window: it materializes the ``(T, B, C, Kc)``
    snapshot AND the product it reduces, so the conv stream alone costs
    ``(4 + 2 × Kc) x C`` per token. This is the widest term anywhere in the
    model at ``Kc = 4``, and it is exactly what a captured VERIFY graph
    retains that a decode graph does not.
  * OUTPUT — the gated-norm output and ``out_proj``.

``split_attn`` runs the mixer body EAGER between the captured pieces, so its
allocations never enter the capture pool and the mixer phases drop out —
which is why that flag keys the graph-pool budget cache.

``(elements, provenance)`` — the per-token activation working set ONE
captured ``(B, S)`` graph retains, per token, at the widest instant.

A captured decode / verify graph runs the WHOLE layer stack, but each
layer's intermediates die at that layer's exit — only the model-owned
``residual_buf`` slab crosses layers, and it lives in ``graph_buffers_pool``,
not in the capture pool. So the retained per-token working set is ONE
layer's peak, taken over the layer kinds THIS model stacks: derived per kind
from that kind's own projection widths, never maxed over every kind the
codebase supports.

``seq`` selects the body a captured graph of that shape actually runs: the
single-token mixer at ``S = 1``, the multi-token one above it. On a GDN
hybrid the two differ by the verify path's unfolded causal conv, which is
the dominant term — so a decode ladder and a verify cross-product do NOT
share one per-token width.

Returns ``"SEEDED"`` with ``hidden x`` the seed multiple when the model
exposes no priceable decoder body (a stub engine, or an arch whose block
nothing here counts), so the caller can report which of the two it got
instead of letting a seed pass for a count.

PERSISTENT private-pool bytes ONE captured ``(B, S)`` graph retains.

A captured cudagraph holds references to every private allocation made
during its capture window, so its forward-activation working set CANNOT be
freed back to the pool after capture — it stays resident for the whole boot
and ACCUMULATES with every other captured graph in the cuMem MemPool.

PRECONDITION — ``(batch, seq)`` is a DECODE (``S=1``), MTP-verify
(``S=K+1``), or drafter-chain (``S=K``) shape. Never a prefill bucket: the
logits term charges the per-token lm_head output at EVERY token, which holds
only where the forward materializes per-token logits — the verify graph
(per-token logits) and the drafter chain (``S=K`` per-step logits over K
steps). On a prefill bucket it over-counts by a factor of ``num_tokens`` — at
and the prefill activation peak is a separate budget line
(``activation_reserve``) that this estimate does not model.

``vocab_override`` / ``logits_elem_bytes`` size the per-token logits term for
a caller whose retained lm_head output is NOT full-vocab in the model dtype.
The default (both ``None``) is the full ``dims.vocab_size`` in the model
dtype — the verify / decode graphs, which materialize the gathered per-token
row. The greedy drafter step instead runs the distributed local-argmax and
retains only its shard-local ``(B, vocab/tp)`` row upcast to fp32, so it
passes ``vocab_override = ceil(vocab/tp)`` + ``logits_elem_bytes = 4``.

The working set is derived from the model dims and the graph's token count
``T = B × S``: a per-token logits buffer, a per-token activation working set
(:func:`_peak_per_token_activation_elems` — one layer's widest concurrently-
live set, from the model's own projection widths, TP-split), the labelled
per-graph activation residual, and the block-table/buffer floor
(:func:`_per_graph_persistent_bytes`). The post-capture hook persists the
EXACT measured pool for warm boots, so this estimate binds boot 1 only.

The cold-boot graph-pool estimate as ``((row, bytes, note, provenance), ...)``.

The sum over ``bytes`` is the estimate itself, so a reader can see WHICH of
its terms were counted off the model and which are stand-ins instead of
taking one number on faith. Provenance values are the shared
:data:`~arbi_serve.runtime.pool_taxonomy.PROVENANCE` vocabulary.

Three rows carry ``SEEDED``: the per-graph activation residual, the fixed
cuMem pool overhead, and whatever the cold headroom multiplier adds on top.
A fourth — the per-token activation working set — is ``DERIVED`` when the
model exposes the widths and ``SEEDED`` when it does not, and says which.

True iff the boot captures the TRUE-STOCHASTIC drafter twin pool.

``ARBI_TRUE_STOCHASTIC_DRAFT != "0"`` (modes ``"1"`` / ``"auto"``) captures a
SECOND ``(B, K)`` drafter ladder into ``drafter_graphs_stoch`` alongside the
greedy one; only an explicit ``"0"`` captures the greedy ladder alone. The
field DEFAULT is ``"auto"``
(:attr:`~arbi_serve.runtime_flags_fields_spec.true_stochastic_draft`), so a
STOCK boot captures BOTH ladders — size the pool for the twin unless the
deployment has opted out.

PERSISTENT graph-pool bytes the captured drafter chains retain, PER RANK.

SUM over every captured ``(B, K)`` drafter bucket of its captured working
set (K per-step head forwards, priced at token count ``B*K``). Summing the
FULL cross-product (every B in the ladder × every served K) bounds the WORST
intermediate bucket, not just ``max_batch``.

Every term is priced at the DRAFTER's row width
(:func:`~arbi_serve.spec_decode.drafter.resolve_draft_vocab_size`), which is
the verify vocabulary unless ``--draft-vocab-prefix`` gave the drafter its
own lm_head over ids ``[0, N)``. A budget that charges these graphs for a
projection the drafter does not perform buys VRAM back from the KV pool for
nothing.

Greedy chain (captured on every boot). The step runs the DISTRIBUTED local-argmax
(:meth:`Qwen3_5MtpHead._distributed_greedy_argmax`): per rank it materializes
only its shard-local ``(B, vocab/tp)`` logits row upcast to fp32 and
all-gathers the ``(max_val, idx)`` scalar pair — it never retains a full
``(B, vocab)`` row. So the per-rank logits term is sized at
``ceil(vocab/tp)`` in fp32, NOT full vocab in the model dtype (which
over-charges ~tp/2× at tp>1 and needlessly refuses depth).

Stochastic twin (captured whenever ``ARBI_TRUE_STOCHASTIC_DRAFT != 0``, which
the ``"auto"`` default satisfies — so on a stock boot this term is LIVE). The
step needs the WHOLE distribution: it retains the ``(B, vocab)`` bf16 logits
AND a persistent ``(K, B, vocab)`` fp32 ``q_chain`` — both un-sharded (the
drafter forms its whole row on every rank). Added here per stochastic graph
on top of the greedy ladder.

The ``(B, K)`` drafter-chain buckets the capture sweep will produce.

``[]`` when MTP is off, on the DFlash drafter (no head chain), or when the
bucket helper is unavailable. Mirrors the guard in
:func:`arbi_serve.engine.capture_admin.drafter_chain.precapture_drafter_chain`
so the pricing walks exactly the buckets the sweep captures.

Split a MEASURED capture-pool total into its two budget rows.

``graph_pool_bytes`` is the sum of BOTH capture pools
(``capture.cudagraphs`` + ``capture.io_buffers``) — that is what
:meth:`~arbi_serve.engine.engine_boot.Engine.total_cudagraph_bytes` reads.
The plan carries a row per pool, so the total cannot go into either row
whole: doing that booked the io pool twice, once inside the total and once
as its own analytic line.

Returns ``(cudagraphs_bytes, io_buffers_bytes, provenance)``:

  * ``MEASURED_CACHED`` — the entry carries the per-tag reading, so each row
    takes its own measured value and the two sum to the total exactly.
  * ``DERIVED`` — the entry has a total but no split (a boot with no cuMem
    per-tag counter). The io row keeps its analytic prediction and the
    cudagraphs row takes the REMAINDER. Subtraction, not duplication: the
    two rows still sum to the measured total, and the term that keeps its
    modelled value is the one the freeze reconcile already diffs against a
    live reading.

``total_bytes <= 0`` (no measurement) returns ``(0, io_analytic, "DERIVED")``
— the caller's cold path owns that case.

Cold-boot upper bound for the PERSISTENT cudagraph private-pool reserve.

Warm boots use the measured ``total_cudagraph_bytes()`` (authoritative,
persisted by :func:`persist_graph_pool_budget`). On the very first boot
there is no measurement, so the KV sizer needs a stand-in.

What this budget line IS: the PERSISTENT cudagraph pool the warm boot
caches — for every captured graph in the (decode S=1 ladder + verify
(B, K+1) + drafter) × (kv-page × LoRA) cross-product, its baked block-table
+ persistent buffers PLUS its captured forward-activation working set (which
a captured cudagraph references and so cannot free back to the pool). It is
NOT the capture-time PREFILL transient: that is the ``activation_reserve``
line (the prefill peak ``raw_peak_bytes``), budgeted separately AND held out
of free for the capture sweep by the deferred KV reserve / B2 prefix-
backing. Counting ``raw_peak_bytes`` here too DOUBLE-COUNTS that transient.

MTP path (this function's body). SUM the per-shape captured footprint via
:func:`_captured_graph_working_set_bytes` — the block-table floor PLUS a
captured activation working set that scales with each graph's ``(B, S)``
token count. Both the verify ``(B, K+1)`` cross-product AND the drafter
``(B, K)`` chains are priced this way (:func:`_drafter_chain_pool_bytes`),
NOT the flat block-table floor. The greedy drafter step retains only its
shard-local ``(B, vocab/tp)`` fp32 argmax row (distributed local-argmax), so
it is priced PER RANK; the stochastic twin (when captured) retains the full
``(B, vocab)`` logits + a global ``(K, B, vocab)`` fp32 ``q_chain``. All GROW
with K. The post-capture persist writes the EXACT measured pool, so the warm
boot needs no estimate at all.

No-MTP path: returns ``raw_peak_bytes + bounded headroom`` unchanged —
that regime is measured separately (the plain prefill peak tracks its
captured pool there).

Log what the instantiate meter saw, and what it did not. Returns the
per-family bytes/graph rates for the persist.

A family ENTERED that metered nothing is the one state the rates line
would otherwise render as success: the rates look complete because the
unmetered family simply is not among them.

Zero instantiates alone does not say WHICH state that is. A sweep gated
off by config, or one whose duck-typed binding is absent on this model,
enters its family and captures nothing — metering zero exactly like a
capture that reached the driver without passing ``instantiate_exec``. The
separator is whether the family's OWN driver bracket grew, filed alongside
the samples by the boot's capture-family seam. Reporting the two as one is
a red that means "did not run", and it buries the family that has the
defect among the several that do not.

Model-load helper consumed by :mod:`arbi_serve.engine.build`.

:func:`load_model_into_engine` builds the model graph + fills weights
+ runs post-load compaction. Two branches:

  * Warm flat-dump path: ``set_skip_weight_load(True)`` → meta-shaped
    graph → background-thread mmap fill via ``load_flat_weights``.
  * Cold path: per-quant-backend bind → streaming compaction into
    ``weights_pool`` with bit-identity gate.

Group every CUDA param/buffer of ``modules`` by backing storage.

Returns ``[(nbytes, [tensor, ...]), ...]`` in first-seen order, one entry
per unique storage so aliased views (weight tying) stay aliased and each
byte is copied exactly once. Zero-numel placeholders are skipped.

Point every tensor in ``tensors`` at ``storage`` starting ``base_bytes``
into it, preserving each view's own offset/shape/stride.

``base_bytes`` must be a multiple of every tensor's element size (the slab
slots are 256-byte aligned, which covers every slab element size). Byte-
identical: the bytes were copied verbatim; only the backing storage changes.

Byte-copy each storage into its OWN ``pool`` allocation, rebinding views.

Peak overhead is ONE storage at a time (each source is reclaimed the moment
its views are rebound), so a card that cannot hold two complete module
residencies still migrates. This is the low-peak default the drafter load
relies on. Returns bytes migrated.

``verify`` compares each destination against its source, every byte, while
both are still live — see :func:`_check_copy_verbatim`.

Raise unless ``dst`` holds exactly ``src``'s bytes.

Every byte, not a sampled signature: a migration that lands a weight
half-copied does not fail, it serves subtly wrong tokens, and the only
moment the evidence exists is while source and destination are both live.
The error names a tensor over the storage so the failure is locatable.

Byte-copy every storage into ONE flat slab in ``pool``, rebinding views.

Each storage lands in a 256-byte-aligned slot of a single contiguous
allocation (the alignment a pointer-direct kernel — Marlin's INT4 GEMM —
reads from, and the same guarantee the per-dtype weight slabs give). The
whole slab is live while the sources are copied in one at a time and freed,
so the caller must have gated on the slab fitting alongside the sources.
Returns bytes migrated.

``verify`` compares each slot against its source, every byte, while both are
still live — see :func:`_check_copy_verbatim`.

Migrate every CUDA parameter/buffer of ``module`` into ``pool``.

Loads a separate spec-decode drafter (DFlash / external draft model) in
the DEFAULT torch allocator first — its dense-construct + staging
transient then returns to the driver via ``empty_cache`` — and only the
FINAL persistent tensors move into the cuMem-backed ``model.drafter``
pool. Loading directly INSIDE the pool strands the transient: a
cuMem-backed :class:`NamedMemPool` never returns freed blocks to the
boot OOMs at MTP attach.

Mirrors the streaming compactor's storage dance:

  1. run the compaction RELEASE hooks (quant linears tear down derived
     kernel descriptors + registry refs so the old GPU storage is owned
     solely by the module tensors);
  2. per unique storage: byte-copy into a pool allocation, then rebind
     every tensor view of that storage onto the pooled bytes
     (peak overhead = ONE storage at a time, not the module);
  3. ``empty_cache`` returns the old default-pool storages to the driver;
  4. run the compaction REBIND hooks (descriptors rebuilt over the pooled
     tensors — same contract as the flat-dump warm reload).

Returns bytes migrated. No-op (0) off-CUDA.

Migrate every CUDA tensor of ``modules`` into ONE flat slab in ``pool``.

The per-tensor :func:`move_module_tensors_into_pool` gives each unique
storage its OWN pool allocation, so every migrated tensor lands in its own
granularity-rounded cuMem segment and the rounding remainder stays pinned
as pad. This variant packs every storage of every module into a single
contiguous slab (one 256-byte-aligned slot each), collapsing those padded
segments into one so the per-segment pad is reclaimed for KV. Byte-identical
to the per-tensor path: each storage is copied verbatim into its slot and
every tensor view rebound onto the same bytes.

Falls back to the per-tensor migration when the slab would not fit alongside
the still-resident sources — the one-slab pack holds the whole destination
live before the sources free, a peak the per-tensor path avoids — so a
tight-VRAM boot never regresses. Runs the compaction release/rebind hooks
around the move exactly as the per-tensor path does. Returns bytes migrated;
no-op (0) off-CUDA.

Keep only the storages the TORCH DEFAULT allocator owns.

Residency is read off the allocator snapshot's live-block address ranges —
the same ranges the ``unpooled.torch_default_pool`` ledger row and its
live-owners table are computed from, so what this selects is exactly what
that row reports. Anything already inside a named pool is left alone.

Returns the input filtered, in order; empty on a snapshot failure, because
migrating nothing is always correct and migrating the wrong set is not.

Migrate ``model``'s DEFAULT-allocator weight storages into ``pool``.

The warm flat-dump path builds the model graph outside ``model.weights`` on
purpose: a cuMem-backed pool never returns a freed block (pytorch#145168),
so the meta-build and codec-swap transients would become a permanent
multi-GiB free list on the weights pool. The DENSE checkpoint parameters
that build walk materializes are collateral — every quantized buffer is an
empty placeholder the loader presizes inside the pool, but a dense one
arrives with real default-allocator storage and the flat dump fills it in
place, so it stays there. The cold paths bind the same tensors into the
weights slab. This is the seam that makes the warm path agree: the build
transients are already back with the driver by now, so only the surviving
weights move.

Byte-for-byte identical, and gated on being so — every migrated byte is
compared against its source while both are live, and each tensor's bit
signature is re-taken over the rebound view, so a wrong slot offset or a
short copy fails the boot rather than serving subtly wrong tokens.

Migration rebinds storages, which changes device pointers, so it is sound
only before anything bakes one. Call it after the weights load and before
the derived weights, the RoPE tables, and cudagraph capture.

``exclude`` names the modules whose weights a LATER boot phase can drop —
the head-quant and embed-quant swaps replace a dense shard precisely to free
it, and a pool that never returns a freed block (pytorch#145168) would turn
that reclaim into stranded bytes. A weight something still intends to free
belongs in the allocator that can free it.

Returns ``(bytes_moved, storages_moved)``. No-op ``(0, 0)`` off-CUDA, with
no pool, and on a cold boot, whose load bound these same tensors into the
weights slab already. What it found is logged either way.

Return ``True`` if a weight-quant backend (AWQ / FP8 / EXL3 /
NVFP4) owns the checkpoint at ``model_dir``.

Drives the cold-path route choice: dense checkpoints take the
allocate-then-fill path (pre-bind to flat slabs → 1× peak), quant
checkpoints keep the load → ``stream_compact_per_layer`` route
because the quant swap's GPU-side repack changes buffer shapes that
can't be pre-sized from the pre-load graph. Detection only inspects
the safetensors tensor names; it does not materialize any weights.

Return the owning quant backend's ``name`` for the checkpoint at
``model_dir`` (``"exl3"`` / ``"awq"`` / ...), or ``None`` for a dense
checkpoint / detection failure.

Drives the cold-path EXL3 DIRECT-slab route (only the EXL3 backend
supports streaming trellis straight into the slab). Inspects only
safetensors tensor names — materializes nothing.

Boot INFO log of packed-weights stats for a load ``strategy``.

``stats`` may omit ``tensors``/``groups``/``bytes`` on some binder paths
(e.g. the allocate-then-fill ``verify()`` shape, which doesn't populate
them on hybrid-Mamba loads) — read defensively so the log can't abort boot.

Find a stable-VA resident record whose weights this build can SHARE.

Same-model members of the residency pool (e.g. a live config-override
variant that differs only in capture-affecting knobs) must not pay a
second full weight copy — weights are immutable at serve time and a
parked member's weights stay VRAM-mapped (a park excludes the weights
pool by design — see ``_park_mechanics``), so the donor's storage is
always readable by the active sharer.

Matched on ``(model.path, dtype)`` and NOT on the rest of the config: two
members of the same checkpoint can differ in STRUCTURE (``mtp.enabled`` off
→ on builds the bundled MTP head), and rejecting such a donor would send a
27B to a cold full build — a second whole weight residency, which is the
cost this function exists to avoid. A donor that covers only part of the
member's graph is ACCEPTED; :func:`cover_donor_share_gap` loads the rest
from the checkpoint and books its bytes.

Returns ``(donor_key, donor_model)`` when ALL of:

  * stable-VA residency is engaged with at least one prepared record,
  * a record's snapshot matches this build's ``(model.path, dtype)``
    exactly,
  * the donor's model is a real materialized ``nn.Module`` (CPU
    bookkeeping tests park string sentinels — never shareable),
  * the donor performed NO in-place weight mutation (the tkv o_proj
    fold writes ``o_proj.weight`` with the member's OWN calibration —
    a fold-mutated donor cannot back a member with a different
    bundle; full build instead).

Dense AND weight-quantized (EXL3 / AWQ / FP8 / NVFP4) checkpoints share.
A quant member builds its graph under ``skip_weight_load`` (deferred
``backend.bind()``), leaving zero-element quant weight placeholders exactly
like a warm flat-dump reload; :class:`DonorWeightBinder` aliases the donor's
populated weight buffers into them (the immutable weights already stay
VRAM-mapped through the donor's park) and the post-load rebind hook rebuilds
each quant linear's derived kernel descriptor over the shared storage. This
is what lets a live config-override flip of a capture-affecting knob on a
27B EXL3 model share its weight residency instead of paying a
second full copy. (A tkv o_proj fold never applies to a quantized o_proj —
it has no dense ``.weight`` — so the fold guard above never trips for EXL3.)

Otherwise ``None`` → the regular full-build path runs (cross-arch
members are unaffected by construction: their ``model.path`` differs).

Fill the part of ``model`` the donor does not cover, from the checkpoint.

A donor is matched on ``(model.path, dtype)``, which two members can share
while their graphs differ in STRUCTURE: ``mtp.enabled`` off → on builds the
checkpoint's bundled MTP head the donor never constructed, and every one of
its params is a name the donor has no tensor for. The share is a memory
optimisation, so the answer is to take the saving on the common part and pay
for the rest — the alternative, a cold full build, is the second whole weight
residency the share exists to avoid.

Two gap kinds, both recorded by :class:`DonorWeightBinder`:

  * ``unshared_params`` — dense params it materialized as empty storage;
    filled here by a ``weight_map``-restricted :func:`load_model_weights`.
  * ``unshared_modules`` — quant modules with zero donor coverage, whose
    buffers only ``backend.bind()`` can size; bound here from the swap's
    recorded deferred-bind paths.

TRANSACTIONAL: the checkpoint is checked for EVERY missing name before a
single byte is written, and a checkpoint that cannot supply them (a drafter
path aimed at a model with no MTP head) raises a
:class:`RuntimeError` naming them — never a bare ``KeyError`` from a binder
and never a half-filled graph. The caller publishes ``eng.model`` only after
this returns.

Returns the bytes the gap cost — VRAM a full share would not have paid, so
the caller records it for the boot log and the KV budget.

Construct THIS member's module graph and bind the donor's weight
storage into it — zero new weight VRAM for every weight the two graphs
share, no safetensors read for those.

The module graph is constructed fresh (meta device) so per-member module
state stays distinct; :class:`DonorWeightBinder` rebinds every meta
param/buffer the donor also holds to the donor's live tensor by qualified
name. ``set_skip_weight_load(True)`` skips the safetensors fill — the
donor's bytes ARE those weights.

A donor matched on ``(model.path, dtype)`` need not cover the member's whole
graph (``mtp.enabled`` off → on adds the bundled MTP head). What it does not
cover is loaded from the checkpoint by :func:`cover_donor_share_gap`, whose
cost is recorded on ``eng.boot_state.donor_share_gap_bytes``.

Stamps ``eng._weights_donor_key`` so the residency controller records
the dependency (park/unmap guard + shutdown ordering) and the o_proj
fold knows the storage is shared (it must not mutate it).

``eng.model`` is published only after the gap is covered, so a checkpoint
that cannot supply the missing params leaves the engine untouched.

Build the model graph and fill weights, exactly as ``build.py`` does.

Sets ``eng.model``, ``eng._f_weights``, ``eng._weight_load_t0`` so
later phases (post-load empty_cache, RoPE pool attach, profile)
have the same engine state regardless of which call path
constructed them.

Three branches:

  - **Weight-share (same-model residency member):** a prepared
    stable-VA record with the same ``(model.path, dtype)`` donates its
    weight storage — fresh module graph, donor-bound tensors, zero new
    weight VRAM (see :func:`find_weight_share_donor`).

  - **Warm (flat-dump cache hit):** model graph constructed on
    meta-shaped tensors (``set_skip_weight_load(True)``), then a
    background thread mmap-fills weights via ``load_flat_weights``.
    ``eng._f_weights`` carries the future the caller must
    ``.result()`` before the first forward.

  - **Cold (no cache):** weights load through the per-quant-backend
    bind path; post-load streaming compaction packs every
    persistent tensor into one flat slab per dtype inside
    ``weights_pool``. Bit-identity gate guards against silent
    corruption.

Reclaim, then DMA the dump into the live storage. Unchanged path.

The DESTINATION storage must be allocated under the weights pool, so
cuMem-backed MemPool never returns a freed segment (pytorch#145168),
so a ring allocated here became a permanent free list on
serving freeze. Route it to ``model.load_scratch``, whose whole
contract is "dropped after boot", and which the post-load
``release_load_scratch`` actually hands back.

Post-fill pool accounting (logging only; never raises).

Splits the warm load's pool growth into "the graph the materialize
walk built" vs "what the dump fill added on top". Cold's pool ends
it, and that delta IS the warm-vs-cold KV gap (81 pages). Knowing
which half of the load owns it is the difference between fixing the
walk and fixing the fill.

Cheap advisory early-out: refuse an egregiously over-budget capture sweep.

The real guarantee is the cuMem ``graph_pool`` map-time cap (see
:func:`arbi_serve.engine.pool_caps.arm_boot_caps`): the captured-graph
working set routes through the pluggable allocator
(``torch.cuda.graph(pool=graph_pool.id)``), so an over-budget capture is
DENIED at map time — a contained in-pool OOM that falls to eager (TP1) or
raises loud + symmetric (TP>1), never a card OOM. This predictor is NOT
load-bearing for containment; it is a fast, conservative backstop that fails
LOUD with the exact over-budget byte count BEFORE the sweep runs, so an
operator who configured a (max_batch, K) so far over the headroom that even
eager fallback would leave no usable coverage gets an actionable message
instead of an all-eager boot.

This gate runs AFTER the MTP driver is attached (so the verify shapes are
real) and BEFORE the sweep. It predicts the captured graph-pool footprint
for the WHOLE sweep — the full ``(B, S)`` decode + verify cross-product ×
the kv-page × LoRA bucket multiplier — and compares it against the free VRAM
minus the runtime floor. If the prediction is over, it raises
:class:`MemoryBudgetError` with the over-budget byte count and the knobs that
fix it.

No-op when CUDA is unavailable, cuda_graphs is off, or the prediction is
unavailable. The prediction reuses the SAME per-shape captured footprint
(:func:`_captured_graph_working_set_bytes`) as
:func:`cold_boot_graph_pool_upper_bound_bytes` — each captured graph retains
its block-table floor PLUS a forward-activation working set that scales with
its ``(B, S)`` token count (a captured cudagraph references its private
allocations and cannot free them back), so the peak is the per-shape SUM of
those (these ACCUMULATE) PLUS a SINGLE capture-time prefill warmup transient
(only one warmup forward is live at a time; it overlays the partially-built
pool). The verify ``(B, K+1)`` cross-product dominates and grows with K.
Because the cuMem cap is the actual containment, a borderline config the
predictor judges to fit but the real capture cannot is still caught by
the cap (a contained denial), so a predictor under-shoot never becomes a
silent card OOM.

Make the decode-graph pool VRAM-budget-bound, not count-bound.

Reads the graph_pool VRAM budget profiled by
:func:`arbi_serve.engine.profile.profile_and_size_kv_pool`
(``eng.boot_state.graph_pool_budgeted_bytes``) and hands it to the decode
:class:`CapturedGraphPool` as ``max_bytes``. From that point the
pool retains every captured decode/verify shape that fits the
budget and only LRU-evicts when adding a graph would exceed it —
so ``cudagraph_max_shapes`` stops being the binding constraint
(VRAM is). Eviction, when it fires, is logged loudly by the pool.

No-op when:
  * cuda_graphs is disabled, OR
  * the budget is zero/absent (cold path with no estimate) —
    the pool stays on its legacy count-bound ``max_shapes`` LRU.

Idempotent: safe to call again on a stable-VA resume (it just
re-stamps the same budget).

Post-capture headroom gate that runs AFTER the cudagraph capture sweeps.

This is the ONE place ``cfg.vram_mode`` has meaning. Memory management
(measure-then-fit, in-process capture, pool caps, phase-2 freeze) ALWAYS
runs regardless of mode; ``vram_mode`` decides only what happens when the
engine came out of capture INCOMPLETE / undersized (capture-OOM fallbacks
to eager, or post-capture free below the runtime floor):

  * ``vram_mode == "bench"`` (strict) → RAISE ``MemoryBudgetError``. A
    half-warm engine (eager fallbacks, degraded coverage) would pollute a
    benchmark, so a benchmark must fail loud rather than report misleading
    numbers.
  * ``vram_mode == "production"`` (lenient, default) → WARN and PROCEED.
    Real serving never refuses to start over capture degradation: it
    serves with whatever coverage it has (degraded buckets run eager /
    compile-and-capture on the fly), which is always correct, just slower.

The deterministic budget (:func:`profile_and_size_kv_pool`) reserves
VRAM for ``weights + peak_activation + extra_safety + KV slabs``.
The captured-graph pools are allocated AFTER ``build_active``,
piecewise bucket and leave the live fallback path with sub-MiB
free VRAM — the next prefill request crashes mid-forward. With the
by-construction pool cap armed (always, when cuMem + cuda_graphs are on)
this cannot card-OOM, but coverage can still degrade — which bench flags
and serving tolerates.

Skips cleanly when CUDA is unavailable.

``(pinned_prefix_pages, config_upper_bound_prefix_pages)``.

The capture-sentinel prefix is only PINNED on ``boot_state`` once the
boot-strategy phase (``_setup_graph_pool_reserve``) decides the KV slab is
growable/prefix-backed — which is AFTER the first viability gate runs. So the
FLOOR reads the pinned value (0 until then ⇒ the gate stays a pure necessary
condition and can never false-refuse), while the QUOTED CAP reads the config
upper bound (:func:`capture_valid_prefix_pages`, ``max_batch + 3``) — a quote
must be safe on the FIRST refusal an operator sees, and over-subtracting a
handful of pages only makes it more conservative.

The REALIZED servable KV page count after the post-capture grow + freeze.

``boot_state.profiled_num_pages`` is the pre-capture FORECAST ceiling. The
post-capture :func:`grow_kv_after_capture` "constrained" degrade branch
(honest scarcity — the card is too full to grow KV to the forecast) WARNs
and serves the small pool it landed WITHOUT rewriting ``profiled_num_pages``
back down — so a floor gate that trusts the forecast is blind to exactly the
KV starvation it exists to catch. Read the live growable pool's mapped page count
(ground truth — the SAME source admission and the metrics collectors already
trust: :mod:`arbi_serve.engine.admission`,
:mod:`arbi_serve.server.metrics._collectors`); fall back to the forecast for
a static (non-growable) pool or a CPU stub that exposes no live pool.

Tokens the captured block table can address, or 0 when unknown.

Read from the boot-state bus rather than from ``eng.piecewise_buffers``:
the buffers are not guaranteed live at this point of the boot, and the bus
carries the one width every paired consumer agreed on
(:func:`~arbi_serve.engine.boot_state.block_table_width`).

The model's ``max_position_embeddings``, or 0 when unknown.

The hard upper bound on any served context: the RoPE cos/sin tables are
built to at most this many rows (``ModelDims.effective_rope_cache_seq_len``
caps there), so a position beyond it has no table entry.

Narrow the served context to the largest window ``pages_available``
proves. Returns the new value, or ``None`` when it cannot narrow.

THE auto contract, in ONE place, because it now has three seams: the
post-capture boot gate, and the live-swap seam where a re-woken member's
re-provisioned pool comes back smaller than the window it was booted at.
Same event, same counter, same operator-visible WARNING — a second copy is
how one of them keeps telling the older story.

Only ever NARROWS. That is what makes it safe after capture: the block
tables, the drafter slabs and the verify buffers were all sized at the
wider value, so a smaller ceiling is bounded by what they already hold.
Widening past them is not this function's to do (the realized-pool reclaim
:func:`finalize_served_max_context` owns that, before capture locks).

Refuses to narrow below one block and returns ``None`` so the caller falls
through to its own loud refusal: a "servable" window that cannot hold a
single block is not a narrower promise, it is a broken one. The caller must
check the return value — treating ``None`` as success is how a check that
could not run reports green.

``margin_bytes=0``: this runs against a MEASURED, already-mapped pool, so
there is no prediction error and no allocator growth left to round for.

True when the served context is the engine's to flex, not the operator's pin.

``--max-context auto`` obviously flexes. A RECIPE-supplied default does
too: an operator who passed only Only a real, operator-typed ``--max-context`` still refuses loud.

Both bits must be read at EVERY narrow site or the same value behaves as a
pin at one gate and a default at another.

Fail loud when the captured graph pool starved the KV pool below a
servable floor — the "4 KV pages" footgun.

The captured cudagraph pools (``graph_pool``: MTP verify/decode graphs AND
the piecewise-prefill bucket sweep) are allocated against residual VRAM
AFTER weights + activation reserve, and the KV pool then GROWS into
whatever free VRAM remains (:func:`grow_kv_after_capture`). Two capture
sources can balloon that pool:

  * the MTP verify ``|B-ladder| × K`` cross-product (a large verify pool
    can leave only a handful of KV pages on a tight card), and
  * the piecewise-prefill bucket sweep (the reason
    ``compose.serve.yaml`` keeps the non-piecewise
    ``--prefill-capture`` default).

Either way the grow lands only a handful of pages while free VRAM still
(a free-VRAM gate) PASSES, and the engine silently serves with a KV pool
too small to hold a single request — correct output, but it
preempts/thrashes on the first prompt and TPOT collapses.

This gate is the page-count analogue the free-VRAM gate lacks: if the
REALIZED servable KV page count (the live growable pool's mapped pages —
:func:`_realized_servable_kv_pages`, NOT the pre-capture forecast the
"constrained" degrade branch leaves stale) is below the VIABILITY floor, it
names the capture pool as the cause + the knobs and REFUSES to boot — by
default, in BOTH vram modes, because a pool this small is not "degraded
serving," it is a dead pool.

The floor is a VIABILITY threshold, not a bare page count: the non-servable
capture-sentinel prefix (``boot_state.kv_growable_prefix_pages`` ==
``max_batch + 3``) PLUS one full-context working KV window
(:func:`arbi_serve.engine.memory_budget.min_useful_kv_pages`, floored at
:data:`_MIN_VIABLE_KV_PAGES`), and never below the legacy
``max(_MIN_VIABLE_KV_PAGES, max_batch)`` count. The prefix term is what
catches the high-max_batch silent degrade: KV stuck at the un-grown prefix
(``num_pages == max_batch + 3``) clears a bare ``max_batch`` count but has
ZERO request-servable pages, so ``prefix + one window`` refuses it while a
healthy pool (KV grown to hundreds of pages) clears it easily.
``ARBI_ALLOW_DEGRADED_CAPTURE=1`` is the explicit escape hatch
(mirrors :func:`enforce_decode_capture_integrity`): it downgrades the
raise to a loud WARN + a registered ``serving_degraded`` entry so the
starvation stays visible instead of silent.

No-op when: cuda_graphs is off (no capture pool to blame), the operator
pinned ``--num-pages`` explicitly (their choice, not a capture surprise),
or the page count is unavailable (non-paged-KV / CPU stub).

What ``--max-batch`` alone bought this boot, against the KV it left.

``None`` when the question does not arise: no KV budget computed yet, no
per-page size, or a model with none of
:data:`~arbi_serve.runtime.pool_taxonomy.CONTEXT_FREE_PER_SEQUENCE_POOLS`
(every all-attention model — its per-sequence state IS the KV pool).

The two sides are commensurable on purpose. ``state_bytes`` is this boot's
own budget rows for the pools that carry one slab row per admissible
sequence and no context term, so they are the same bytes the budget just
subtracted from KV. ``kv_ceiling_bytes`` is the SERVING ceiling, not the
pre-capture cut: the most KV this configuration can ever hold once the
boot's one-time capture transient frees. Comparing against the cut instead
would refuse a boot the post-capture grow goes on to serve.

Largest ``--max-batch`` whose per-sequence state stays inside the KV it leaves.

The two terms trade against ONE pot: what the per-sequence pools do not take
is what the KV ceiling gets, so with ``r`` bytes per slab row and a shared
pot ``P = state + kv_ceiling`` the condition ``(N+1)·r <= P − (N+1)·r``
inverts to ``N <= P / 2r − 1``.

CONSERVATIVE by construction: it holds every other row of the budget at
this boot's size, and several of them (the capture ladder, the penalty and
sampler scratch, the drafter slots) shrink with the batch too — so the
returned batch clears the gate with room, rather than landing on it.

Refuse a boot whose per-sequence state ate the KV pool that serves it.

``--max-batch`` buys TWO things on a hybrid (GDN / Mamba / short-conv)
model, and only one of them is visible as "KV". Every admissible sequence
gets a KV page budget, which is elastic — a narrower ``max_context`` gives
the same pool more sequences. It ALSO gets a recurrent-state slab row,
which is not: those rows are
:data:`~arbi_serve.runtime.pool_taxonomy.CONTEXT_FREE_PER_SEQUENCE_POOLS`,
allocated whole at boot, the same size at any window. Which half is larger
follows the model's layer mix — a checkpoint with more linear-attention
layers than full-attention ones puts most of its per-sequence bytes in the
inelastic half — and that half is subtracted from the budget BEFORE KV is
sized. The boot reports both sides on the refusal below and on the VRAM
ledger's ``state.*`` group.

So a ``--max-batch`` past a certain point spends the card on state for
sequences the remaining KV cannot give context to, and every lever the
other gates name is the wrong one: narrowing the context does not return a
byte of it, and the post-capture grow cannot either (the ceiling this
compares against IS the grow's own target). Refuse HERE, before the pool is
built and long before the capture sweep whose headroom the same bytes
consumed, and name the pools, the ``max_batch`` that sized them, and the
batch that fits.

NOT a clamp. ``--max-batch`` is what the operator asked for; quietly
serving fewer rows than that is the same defect in the other direction.

Where this gate fires, :func:`assert_kv_pages_floor` cannot: its refusal is
for a starved pool with GPU memory sitting FREE past the runtime floor, and
here the memory is spent. Skipped under ``--num-pages``, where the operator
hand-sized the pool and owns the ``(pages, batch)`` tuple — the same
exemption ``_validate_max_batch_fits`` makes.

Fail loud PRE-capture when the KV budget cannot hold one ``max_context`` window.

The pre-capture twin of :func:`assert_kv_pages_floor`. Both enforce the same
viability floor — the KV pool must hold at least one full ``max_context``
window — but that gate runs AFTER the cudagraph capture sweep and reports the
shortfall as "the cudagraph capture pool STARVED the KV cache", offering only
capture/batch remedies (lower ``--mtp-n-draft``, fewer capture buckets, drop
``--cuda-graphs``). That attribution is right when the capture pool is what
consumed the KV headroom. It is WRONG when the pool was never big enough for
one window in the first place — the case an unpinned ``--max-context`` walks
into, since it inherits the model's full trained ceiling (262144 for
Qwen3.6-27B). In that shape, the capture pool is a small fraction of the
gap between the ceiling and the one-window floor, so every knob the
post-capture message names is incapable of closing it while the one
knob that is — ``--max-context`` — goes unmentioned. A refusal that names
the wrong cause and the wrong remedy is the same class of defect as a
silent fallback: it sends the operator to tune a lever that cannot matter.

The REALIZED servable page count is only known post-capture, but a NECESSARY
condition for viability is knowable now: the profiled KV serving ceiling
(``boot_state.kv_serving_ceiling_pages`` — the MAX pages the growable pool
can ever map, ``all_reduce(MIN)``-reduced across the TP group so it is
rank-symmetric) must hold at least ONE full ``max_context`` window
(:func:`arbi_serve.engine.memory_budget.min_useful_kv_pages`) PLUS the
1-page null sentinel the live pool withholds from a request. The realized
pool can only ever be SMALLER than this ceiling, so a ceiling below that floor
GUARANTEES the post-capture gate would refuse — we just
refuse HERE, before a capture sweep whose result would be discarded. The floor used here is
``max(_MIN_VIABLE_KV_PAGES, null_sentinel + one_window)``, which is ``<=``
the post-capture floor (that also maxes in ``max_batch``), so a
refusal here is ALWAYS a refusal the post-capture gate would also reach —
never a false refusal of a config that would otherwise boot.

CALLED TWICE by :func:`arbi_serve.engine.build.build`. The prefix is only
pinned once the boot-strategy phase (``_setup_graph_pool_reserve``) decides
the pool is prefix-backed, which is AFTER the profile. The first call
(post-profile, prefix still 0) is the pure one-window necessary condition and
catches the uncapped-``auto`` case as early as possible; the second (once the
prefix is known, still before the capture sweep) closes the band where a
``max_context`` whose window fits the ceiling but whose window PLUS the
prefix does not would otherwise reach the capture sweep and OOM mid-sizing
instead of getting a clean :class:`MemoryBudgetError`.

AUTO vs EXPLICIT. The outcome when the ceiling does not fit depends on
whether the operator PINNED ``--max-context`` or left it to the engine
(``cfg.cache.max_context_is_auto`` — the one bit that survives
``_resolve_auto_max_context`` overwriting the ``"auto"`` sentinel with the
concrete ceiling):

  * ``auto`` (omitted) → NARROW ``cfg.cache.max_context`` to the largest
    proven-servable window and BOOT. This is CacheConfig's documented auto
    contract ("the model ceiling, clamped down by the realized KV
    capacity"); the narrowing is LOUD (WARNING + ``kv_context_window_auto_narrow``
    counter), and because this gate runs BEFORE capture-bucket sizing /
    admission / verify-buffer sizing — all of which read
    ``cfg.cache.max_context`` live — every downstream consumer sees the
    narrowed value. Called twice, it is idempotent: the second call finds
    the already-narrowed (fitting) ceiling and returns. If even ``auto``
    cannot find one servable block, it fails loud like the explicit case.
  * EXPLICIT ``--max-context N`` → hard refusal (below), no
    ``ARBI_ALLOW_DEGRADED_CAPTURE`` escape: the operator pinned a number
    the card cannot honor and must SEE that, not get a quiet downgrade to a
    window they did not ask for. Bypassing the gate would only run the
    discarded capture sweep and reach the post-capture gate's same refusal
    with the wrong attribution (capture pool, knobs that cannot close the
    gap) — the actionable remedy is to lower ``--max-context``, not to
    force the boot.

No-op when: cuda_graphs is off (mirrors :func:`assert_kv_pages_floor` — this
precedes a capture sweep that only runs under cuda_graphs), the operator
pinned ``--num-pages`` (their deliberate choice — same carve-out as the
post-capture gate), or the serving ceiling / block size is unavailable
(non-paged-KV model / CPU stub).

The longest context the SIZING-time buffers can physically address.

A HARD cap on the served context: admitting a request longer than the block
tables can index is a correctness bug, not a tuning miss. Two families of
pre-resize buffer bound it, and the answer is the smaller:

  * the block tables (piecewise prefill, the captured decode/verify chain,
    the drafter's shared-KV chain). These are sized by
    :func:`arbi_serve._units.block_table_pages`, i.e. at the LARGER of the
    ``max_context`` window and the pool's page ceiling — so on a growable
    pool they already address every page the pool can ever map, which is
    what lets the served context move after the grow without re-capture.
  * the persistent :class:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers`
    block table, sized the same way plus the ``max_k`` speculative tail.
    Absent (no MTP driver) it does not bind.

Returns a token count on a whole-page boundary. Pure arithmetic — no CUDA.

Decide the ADMISSION-time served context from the REALIZED KV pool.

``max_context`` is two different numbers wearing one name:

  1. the SIZING bound — it scales the block tables / drafter slabs that are
     allocated BEFORE the KV pool grows, and those allocations decide how
     much VRAM is left for KV. It must be picked up front, from a FORECAST
     of the pool, because the pool's size depends on it. That is
     :func:`assert_kv_context_window_fits`, unchanged, and its value is
     recorded on ``boot_state.kv_sizing_max_context``.
  2. the SERVED bound — the length a request may be admitted to. This one
     does not have to be known before the pool exists, and deciding it from
     a forecast is what stranded context: a conservative prediction
     permanently capped serving below what the realized pool could hold.

Called immediately after the deferred KV resize, so ``pool_pages_total`` is
final. The served context is::

    min(realized KV capacity, sizing-time addressable cap, model ceiling)

The sizing cap is asserted, not assumed: the pre-resize buffers cannot index
past it. Because that cap is ``sizing_context`` rounded up to a whole page,
any RAISE this function makes leaves ``ceil(max_context / block_size)``
unchanged, so no page-granular consumer — including the capture sweep that
runs after this point — changes shape as a result.

Writes the result to ``cfg.cache.max_context`` (the field admission and the
terminal-length check already read) and mirrors it to
``boot_state.served_max_context``. A LOWER result narrows loudly HERE, where
the pool is known, rather than at the earlier forecast. A HIGHER result is
taken — that is the point of the split.

``pool_is_final`` says whether the pool has reached its final size. Left
unset it is INFERRED: a static pool is final as soon as the deferred resize
built it, while a GROWABLE prefix-backed pool has only its capture-sentinel
prefix mapped at that point and is not realized until the post-capture
``grow_kv_after_capture``. Reading the prefix as the realized pool would
narrow the served context to nothing, so the growable path returns here and
is finalized again — explicitly ``pool_is_final=True`` — after the grow.

No-op for a PINNED ``--max-context`` (the operator's number is never
rewritten; the post-capture floor gate still refuses loud if it does not
fit), for an explicit ``--num-pages`` override, and for a stub/non-paged-KV
engine with no realized pool to read.

Driver-truth (used, total) VRAM bytes for the engine's device, or
``None`` when no CUDA device is available (CPU tests). Split out so the
CPU suite can monkeypatch driver readings.

SETTLES the allocator first (gc + synchronize + empty_cache): right after
a member switch the outgoing member's dropped pools linger until Python GC
runs, so an immediate ``mem_get_info`` still counts them and masks a
no-op'd wake. Admin-swap path only — never the serving hot path — so the
settle cost is irrelevant.

The ACTIVE KV pool's realized capacity vs one ``max_context`` window.

The measurement the live-swap gate (:func:`assert_swap_kv_capacity`)
refuses on, and the payload every successful ``aswap_attention_backend``
now carries — so a harness health-check can assert CAPACITY instead of
probing with a 1-token generation that always fits (the probe that let a
swap serving 0 tokens on every real-length prompt report "serving").

Reads the REALIZED mapped page count (:func:`_realized_servable_kv_pages`
— the same ground truth admission divides by), never a forecast: a
budget-capped wake of a discard-parked member can remap a fraction of the
pool's reserved ``num_pages``.

Returns ``None`` when there is no sized paged-KV pool / no block size
(non-paged model, CPU stub) — the caller treats that as "not applicable",
never as "capacity OK by default".

Fail loud when a live backend/KV-codec swap lands an unservable pool.

The runtime twin of the two boot gates (:func:`assert_kv_pages_floor` /
:func:`assert_kv_context_window_fits`): after a hot-swap the ACTIVE pool
must still hold at least ONE full ``max_context`` window
(:func:`min_useful_kv_pages`, which counts the page-0 null sentinel), or the
engine cannot serve a single real-length
request — it would 429/fault every real prompt while a 1-token health
probe (which always fits) reports "serving". This is exactly what happens
when the canonical single-boot arbi row boots the compressed tkv codec
primary at gmu 0.99 and a live swap to bf16 (``tkv-bypass``, ~4× KV
bytes/token) lands a pool that cannot hold a 16k prompt: the swap's
1-token health-check passes while the whole arm silently serves 0 tokens.

Boot gates run inside ``Engine.build`` and cover the build-required swap
route; this gate runs at the SWAP SEAM (``aswap_attention_backend``) so it
also covers the routes that never re-build:

  * KV slab is remapped under a FREE-VRAM budget cap
    (``wake_growable_kv(budget_bytes=…)``) and can come back a sliver of
    its registered size while ``num_pages`` still reports the reserve,
  * an instant-switch wake whose physical remap silently NO-OPs while the
    page bookkeeping keeps claiming the full slab (the bookkeeping can
    claim a fully-mapped slab while the driver holds only a fraction
    resident for the whole card; the next prefill then segfaults in
    ``cuMemcpyDtoDAsync``) — caught by the driver-truth residency
    cross-check in :func:`swap_kv_capacity_report`, and
  * a build whose post-capture grow degraded below the floor after the
    boot gates passed.

Carve-out (mirrors both boot gates): an explicit ``--num-pages`` override
is the operator's deliberate choice — the e2e smoke lane pins it exactly
so a contended-card swap does not refuse. No degraded-mode escape
hatch otherwise: a swap that cannot serve real traffic must never report
healthy. Raises :class:`MemoryBudgetError`; the swap admin latches a
sticky swap fault on it (serving suspended, ``/health/ready`` 503) so the
refusal is loud end to end. No-op for non-paged / unsized pools.

The clause a KV refusal must carry when its evidence may be an artefact.

THE single rendering of "is this ceiling UNMEASURED or EXCEEDED?", quoted
by every refusal that offers ``profiled_num_pages`` / a serving floor as
evidence that the card is full. Returns ``""`` when the activation profile
measured every shape — the number then really is what the card can hold,
and the caller's own "free VRAM / lower memory pressure" advice stands.

When it is NOT empty the caller must NOT tell the operator the card is too
full. A serving floor inflated by an unmeasured arena split reserves one
forward's activations twice per shape, so the ceiling it produces says
nothing about free VRAM. Two agents were sent to shrink a config that was
never the cause; shrinking eventually "works" only by changing the
allocation count the probe emits, which is not the memory pressure the
message blamed.

Reads ``boot_state.activation_profile_unmeasured_note``, written once at
the profile phase (:func:`~arbi_serve.engine.build_phases_kv.
_build_activation_profile`) from
:meth:`~arbi_serve.runtime.activation_profile.MultiShapeProfile.
unmeasured_arena_note`.

``ceiling_pages`` minus the pages the plan granted against the OSCAR basis.

ONE definition, called by both post-install KV sizing passes — the Phase-4c
deferred resize and the post-capture grow — because a ceiling that lives in
two places is two ceilings, and a cut applied to only one of them is undone
by the other the moment the other one binds.

The plan's ``scratch.attn_codec`` row prices that pool's TQ buffers
(``predict_tkv_scratch_pool_bytes``); the OSCAR basis is resident in the
same pool and has no line at all, so ``compute_kv_budget`` handed KV those
bytes. :func:`arbi_serve.engine.tkv_install.record_attn_codec_basis_bytes`
MEASURES what the install actually made resident, which is why this is a
correction rather than a second forecast.

ORDERING. The install runs inside ``build_active``, which the profile phase
calls before the deferred resize and long before the capture sweep, so both
passes read a measurement that is already complete. A boot that measured
nothing leaves ``attn_codec_basis_bytes`` at 0 and this is the identity —
as is every boot with no rotation configured.

The clause an out-of-memory boot must carry when its budget cache was COLD.

Several rows of the memory plan are MEASURED-or-nothing: the serving
torch-caching overhang, the driver's serving-window growth, the graph pool.
Each is measured where it is observable, persisted under this
configuration's budget-cache identity, and read back by the next boot at
that identity. A boot with none of them held is boot 1 of that
configuration, and it sizes memory with rows that stand on seeds.

That is a legitimate state exactly ONCE. It repeats forever when the cache
directory does not survive the container: the fallback is
``~/.cache/arbi-serve/budget-cache``, which on an ephemeral root filesystem
is gone at every start, so every boot is boot 1 and every boot fails the
same way. ``ARBI_SERVE_BUDGET_CACHE_DIR`` on a persistent mount is what
fixes it, and every committed compose file sets it.

Returns ``""`` when this boot HAD the measurement — the memory pressure is
then real and the caller's own advice stands. When it is not empty the
caller must NOT offer ``--gpu-memory-utilization`` / ``--max-context``: this
configuration has never been measured on this card, so shrinking it tunes a
plan that was never the constraint.

Names of non-zero budget rows reporting ``SEEDED`` where a measurement belongs.

``rows`` is ``((name, bytes, note, provenance), ...)``. A row the taxonomy
itself declares ``SEEDED`` is left out: ``unclaimed.allocator_slack`` is a
permanent bootstrap constant (a chosen round-up, never measured) and
``transient.serving_step.token2wav`` an operator-set reserve, so on those a
seed is the design and not a failed measurement — the same rule the
operator card's ``unexpected_provenance`` applies.

Refuse (bench) or warn (production) a serving budget that stands on seeds.

Reads the itemised serving-step reserve
(:attr:`~arbi_serve.engine.boot_state.EngineBootState.serving_step_terms`)
once the post-capture KV grow has priced it and before the freeze persists
this boot's measurements. A ``SEEDED`` row there is a bootstrap constant
standing where a measurement belongs, so a number benchmarked on this boot
is not a measurement of this configuration.

Boot 1 of a configuration has nothing persisted to read and seeds by design
(:func:`cold_budget_cache_note`); every mode warns and proceeds, and the
freeze persists what it measures. A seed on a boot whose cache already held
this configuration means the persistence or its key is broken, and
``vram_mode == "bench"`` refuses to report numbers from it.

``measured_advice`` when the profile is trustworthy, the cause note when not.

Keeps the two facts from being stated at once: an operator either gets the
"the card is too full — free VRAM" advice (because the ceiling was
measured) or the unmeasured-probe cause (because it was not), never both.

THE reserve for a boot that has NO measured capture pool, and says so.

One seam for the whole class "Phase 4a did not come back with a measurement",
whatever prevented it: no cuMem driver, the flag off, a truncated sweep, or
the measurement pass itself hitting an allocation failure. Those all leave
the boot holding the same object — the closed-form cold ESTIMATE, which
:mod:`arbi_serve.engine.capture_sizing` documents as UNDER-COUNTING the true
captured pool — so they must all leave it holding the same NUMBER. Two
different fallbacks for one fact is how the smaller of them ends up on the
path that needs the larger: Phase 4c's holdout is additive in this reserve
(:func:`~arbi_serve.engine.memory_budget.deferred_kv_reserve_bytes`), so a
reserve short by X sizes KV over by X and the Phase-4c ``build_active`` — or
the Phase-5 sweep behind it — is the one that runs out of memory.

Returns ``max(estimate × MULT, FLOOR)``: strictly above the estimate, because
the estimate is a documented lower bound and nothing on this path is going to
correct it before capture. Also publishes the fact — counter, warning, and
``boot_state.graph_pool_reserve_unmeasured_note`` for the boot-ready banner —
so an unmeasured reserve is never indistinguishable from a measured one.

Publish WHY the activation profile produced no measurement, if it did not.

One writer for ``boot_state.activation_profile_unmeasured_note``, so the two
paths that can leave the profile without a measurement — a shape whose arena
split could not be replayed, and a profile that did not run at all — cannot
describe the same fact differently. The KV refusals read it through
:func:`~arbi_serve.engine.build_memory_sizing.scarcity_advice`.

Tolerates an engine double with no ``boot_state``: the note is diagnostic,
and losing it must never turn a working boot into a crash.

Bytes the activation arena must hold for one layer's transients.

The arena rewinds per LAYER (the eager dispatch seam resets it before
allocating), so its capacity is a MAX over layers, not a sum: the widest
attention output buffer any layer can ask for at the engine's widest step.
That is ``max_batched_tokens`` rows of ``num_heads * head_dim`` in the
model's working dtype, taken over every attention layer because a
hetero-head_dim arch's layers differ.

NOT the profiled activation peak. The peak covers every transient of a
whole step on the heap; the arena holds one buffer at a time, and sizing it
at the peak would reserve — permanently, since a cuMem pool keeps a freed
block mapped — VRAM that the arena can never fill.

Returns 0 when the geometry cannot be read, which leaves the arena unbuilt
rather than guessed at.

Construct the activation arena, AFTER the boot's KV sizing has settled.

Deliberately late. The arena's buffer lives in ``scratch.forward_arena``,
and :func:`~arbi_serve.engine.inprocess_capture.
release_idle_forward_arena` refuses to release that pool while ANY storage
in it is live — a captured graph could have baked the address of anything
that survived the step boundary, and the arena buffer survives by
construction. Building it here, after the KV sizing has settled, keeps the
engine's own idle re-tightening reachable for the free physical this pool
strands; the arena's capacity is one layer's transients, a small fraction
of what the pool holds.

Runs on every boot path (the capture sweep is optional); a no-op unless
the flag is on and CUDA is present. The arch refusal already fired in
:func:`_build_activation_profile`, so reaching a build here means
the model threads it. Failure is logged and leaves the engine arena-less,
which is the default-off behaviour.

Device VRAM the centroid fit needs for its WIDEST (layer, side) pool.

The ladder shares one pool per side and pools are built and dropped one at
a time, so the peak is one pool: every pooled token, for one layer's KV
heads at one head_dim, at :data:`_FIT_BYTES_PER_SAMPLE`.

This is the term the calibration-scoped cap was missing. The cap's own
docstring names the centroid fit as a reason the pool must shrink, but it
sized only the capture BASKET and let whatever fell out be the fit's
budget. Measured on a 4090 serving Qwen3.8-27B: the cap freed 0.31 GiB and
the fit's first pool wanted 1.5 GiB, so the boot died after a 30-minute
capture -- as a CUDA OOM with cudagraphs off, and as an out-of-bounds
gather in a replayed MTP-verify graph with them on.

Returns 0 when the topology cannot be read, which leaves the cap exactly
as it was rather than sizing on a guess.

Cap the boot KV pool when boot DEFERRED calibration to this engine.

A deferred-calibration boot is not a serving boot: it comes up in
``tkv-bypass`` purely so in-process calibration can capture raw K/V, and
then hot-swaps to the real codec. Sizing its pool the serving way — fill
the ``gpu_memory_utilization`` budget — is actively wrong twice over:

  * calibration needs WORKING VRAM the pool has just eaten (the capture's
    per-prefill clones, the centroid fit, the drift reference forward), and
  * the swap builds the codec backend as a prepared member BESIDE the
    active one, so it needs headroom too.

A serving-sized pool can OOM after capture even though the calibration
basket's live set is only ONE ``max_context`` prefill's worth of pages.

So size for the calibration basket instead: the concurrent sequences the
calibration traffic actually keeps live, of ``per_seq`` tokens (the short
capture/drift corpus rows, NOT ``max_context``), times a slack multiple.

THIS POOL ONLY. The cap describes the pool the calibration traffic runs
against, and reads a BOOT-SEQUENCE instruction to recognise it, so it can
only stay correct while no other pool is built from that same config. The
swap that follows builds one — the member that will SERVE — and it sizes
from the full budget because
:func:`~arbi_serve.config_overrides._serving_variant_of` strips the
instruction out of every variant config, leaving ``ic`` falsy here. Without
that strip this cap bites a second time, on the served pool, and hands
serving the calibration basket's page count. No-op on a normal boot.

The concurrency term differs by what happens to this pool AFTER calibration:

  * ``--calibrate-only`` writes the bundle and hard-exits, so nothing ever
    serves off this pool. Its only reader is the calibration traffic, whose
    concurrency is :data:`~arbi_serve.calibration.inprocess.CALIBRATION_LIVE_SEQS`
    — every stage awaits one row before submitting the next.
  * the serve path swaps to the real codec afterwards and that swap FAILS
    SAFE: on a bad bundle the bypass backend keeps serving on THIS pool, at
    ``max_batch`` concurrency. So the serve path holds the serving-shaped
    ``min(n_seqs, max_batch)`` term — sizing it down to the calibration
    traffic would turn a fail-safe into an unservable engine.

Set ``mtp_enabled=True`` on every paged-KV backend when MTP is on.

The ``MTPFusedAttend`` wrapper — and thus ``core._mtp_attend``, the
split-K verify kernel — is only constructed when
``TkvBackend.mtp_enabled`` is True (tkv core_factory). The
engine-boot backend constructor ``_parse_tkv`` builds the backend
with no knowledge of ``cfg.mtp``, so without this the (B, K+1)
verify shape silently falls through to the heavier Turbo prefill path
instead of the MTP kernel. Must run BEFORE the cores are built
(``ensure_tkv_cores_built``). ``mtp_block_m`` self-corrects per
step via the verify pass's ``set_mtp_block_m``, so only the enable
flag is threaded here.

Backend-agnostic by attribute: any paged backend that exposes a
writable ``mtp_enabled`` (tkv codec backend; tkv-bypass) is flipped.
tkv-bypass has no bf16-bypass split-K MTP kernel in turbo-attn, so
its verify shape runs through the Turbo prefill path (correct,
eager-by-design); threading the flag keeps the per-step verify scope
(``set_mtp_block_m``) symmetric with tkv.

Raise every build refusal that is decidable from CONFIG ALONE.

The ORDER of a refusal is part of its contract. A build phase runs after
the load has bound weights and, on a donor-share member, after the
compaction release hooks have torn down EXL3 kernel descriptors that the
DONOR — the member still serving — is addressed by
(``release_secondary_refs_for_compaction`` deregisters the linear's uid
from the process-global inner registry, and the donor-shared head modules
are the SAME objects on both graphs). A refusal raised there aborts before
the matching rebind, so the config that was still serving is left unable to
forward. Nothing about the levers gathered here needs a model to decide, so
they are decided before anything shared can move.

Called early in :func:`_preflight`, before the device is touched (so a boot
fails in seconds rather than after a full weight load), and — for a live
config override — before the park/build sequence is entered at all
(:func:`arbi_serve.engine.config_variant._aapply_overrides_locked`).

Raises whatever the individual lever raises — the refusal type and message
are the lever's, unchanged, so an operator reads the same text wherever it
fires.

A lever belongs here only when hoisting it changes NOTHING but the moment
of the refusal. ``lm_head_quant``'s own MTP gate looks eligible and is not:
:func:`_apply_head_quant_phase` deliberately BYPASSES it on a donor-share
member, because an MTP-off member of a ``full-*`` donor serves the donor's
shared quantized head and is a supported config. Hoisting that check would
refuse a boot that works today.

Point the EXL3 descriptor's bsz-1 input row at THIS member's pool.

The row is allocated by ``LinearEXL3.__init__`` during the weight load, so
the pool has to be armed BEFORE the load rather than reserved after it like
the reconstruct scratch — a descriptor's buffer cannot be re-homed
afterwards without rebuilding the descriptor, and the captured decode
graphs bake its address.

``capture.io_buffers`` is the home and
:mod:`arbi_serve.weight_quant.exl3.kernel_scratch` carries the argument.
Armed unconditionally (this costs a module-global write) rather than gated
on the checkpoint being EXL3: a non-EXL3 boot never constructs a
descriptor, so nothing is ever allocated through it.

Disarms when the pool is absent (CPU/test engines): the descriptor row
then comes from the default allocator, and the owner walk names it.

The modules whose dense weights a LATER boot phase may replace and free.

``_apply_head_quant_phase``'s ``full-*`` modes swap ``lm_head`` and
``mtp.fc`` for quantized ones, and ``_apply_embed_quant_phase`` swaps the
input embedding, in both cases to hand the dense shard back before the KV
budget is sized. Those bytes must therefore stay in the allocator that can
return them: a cuMem-backed pool never releases a freed block
(pytorch#145168), so folding them in would convert a reclaim worth over a
gigabyte on a 27B into permanently stranded memory.

Named by MODULE rather than by mode because the mode is not resolved yet —
it resolves against the loaded model in the head-quant phase itself. Listing
them unconditionally costs the fold only these few tensors, which are the
largest in the model and the least worth folding anyway.

Route the sampler's penalty accumulator into ``scratch.penalty_accum``.

The accumulator is a ``(max_batch, vocab)`` device buffer the presence /
frequency penalty path subtracts from the logits every decode step.

``predict_penalty_accum_pool_bytes`` prices this pool from the same
``max_batch`` in :func:`~arbi_serve.engine.profile._predict_per_pool_bytes`,
so the prediction and the allocation agree exactly. The allocation lands
after the KV pool is sized but before the Phase 4c deferred resize measures
free VRAM, so the bytes are counted once.

Allocating OUTSIDE the pool would land the buffer in
``unpooled.torch_default_pool``, which the VRAM ledger reports as
unbudgeted by construction — hence the ``pool.use()`` in ``_alloc``.

Without this hook the accumulator stays DISENGAGED on CUDA and the
penalty path takes the per-step counts matrix — correct, just slower.

Route the sampler mask kernel's device scratch into a named pool.

The scratch is a ``(programs, vocab)`` fp32 buffer plus two fp32 lookup
tables, allocated lazily on the first sampler call. Without this hook they
land in ``unpooled.torch_default_pool`` — which the VRAM ledger reports as
unbudgeted BY CONSTRUCTION, and whose bytes KV sizing has already spent.

Size is not trivial and is not fixed: program count is
``min(SM count, rows)``, and an MTP verify slate presents ``(K+1) x B``

Unlike the penalty accumulator this scratch is NOT optional — the mask
kernel cannot run without it — so a missing pool falls back to the default
allocator (exactly today's behaviour) rather than refusing. Installing the
hook is therefore a pure ACCOUNTING change: the same bytes, in a pool that
names them.

Phase 1b: load-time RTN head quantization (``cfg.lm_head_quant``).

Runs AFTER the full weight load (the dense lm_head shard is the RTN
input) and BEFORE ``_profile_and_size_kv`` (so the freed bf16 shard
in a ``full-*`` mode — or the added quantized copy in a ``drafter-*``
mode — is visible to the KV budget) and BEFORE ``build_mtp_driver``
(whose ``head._lm_head is model.lm_head`` identity assert the
full-mode swap preserves via the tied-module rebind).

The default ``"auto"`` resolves HERE (the loaded model is the input:
:func:`~arbi_serve.weight_quant.head_quant.resolve_head_quant_mode`) —
``drafter-int4`` where the dense-head lever exists, ``off`` with a
logged + counted reason where it does not. The RESOLVED mode is
written back to ``eng.cfg`` so every downstream reader (KV-floor
refusal remedies, flat-dump gating, ``/v1/admin`` config surfaces,
bench boot-fact gates) sees the truth the engine actually serves,
never the unresolved sentinel. Explicit modes pass through and every
misconfiguration fails loud inside
:func:`~arbi_serve.weight_quant.head_quant.apply_head_quant`.

Phase 1b2: the drafter vocabulary-prefix head (``cfg.draft_vocab_prefix``).

Runs AFTER :func:`_apply_head_quant_phase` — the two install through the
SAME drafter seam (``set_draft_lm_head_override``), and this order makes
the collision a loud refusal from the lever the operator named last rather
than a silent overwrite — and BEFORE ``_profile_and_size_kv``, so the added
head is visible to the KV budget.

Four settings, and none of them picks a width for the operator:

``"full"`` (DEFAULT)
    The drafter keeps the whole lm_head. Nothing is built, no path
    changes. The lever is a BANDWIDTH trade that COSTS resident VRAM, so
    it is opt-in.
``"half"`` / ``"quarter"``
    That fraction of THIS checkpoint's lm_head, rounded down to a whole
    128-block. Expressed as a fraction so the cut moves with the model
    rather than asserting a constant that means different things on
    different heads.
``N``
    An assertion that this width is servable. Every ineligible model /
    flag combination FAILS LOUD; nothing is clamped, because a drafter's
    reach the operator cannot read off their own config is not a lever, it
    is a surprise.

A requested cut is measured against the checkpoint's id ordering before it
installs (:func:`~arbi_serve.draft_vocab_prefix.probe_prefix_order`) and
the per-script table is logged AT THAT RUNG — a quarter strands far more
than a half, and the operator has to see the rung they actually chose.

Read the drafter's trailing window off THIS checkpoint's added tokens.

The ids the window has to reach are a tokenizer fact, so the width that
reaches them is derivable and must not be a constant an operator has to
know: one that is never typed leaves a measured acceptance win inert, and
one that is copied to the next checkpoint reaches the wrong ids.

A shape the derivation cannot describe REFUSES by name and falls back to a
plain prefix (``tail == 0``) — never to a window chosen anyway. A window
that reaches none of what it was chosen for still widens the drafter's row,
and nothing on the served path would report that it bought nothing.

Resolve ``half`` / ``quarter`` against THIS checkpoint's lm_head.

Returns the width to install, or 0 (with a named refusal) when the head
cannot be read or the fraction leaves nothing to slice. A fraction is a
request to cut the model in front of us, not an assertion about a number,
so a head that cannot be measured disables the lever instead of failing
the boot — the operator named no width for the failure to be about.

Measure the checkpoint's id ordering at ``n_keep`` and report it.

A GUARD, never a chooser: the operator named this width. What the probe
can say is whether an id prefix is a SHORTLIST on this checkpoint at all
— the frequent tokens sitting at low ids — and that is a checkpoint fact.
Which languages the rung strands is a TRAFFIC fact and belongs to the
operator, so the per-script table is printed AT THE CHOSEN RUNG rather
than gated on.

Best-effort throughout: the probe is instrumentation, and losing it must
not lose the install the operator asked for.

Phase 1c: load-time fp8 quantization of the input embedding
(``cfg.embed_quant``).

Runs AFTER the full weight load (the dense embedding is the RTN input) and
BEFORE ``_profile_and_size_kv``, so the bf16 table freed by the ``fp8`` mode
is visible to the KV budget — the whole point of the mode is that KV grows
into the ~half-table it reclaims. ``off`` is a pure no-op.

Mirrors :func:`_apply_head_quant_phase`: a donor-share member inherits the
(possibly fp8) embedding module from the donor, so re-running the transform
would double-quantize — record the mirrored mode + zero new bytes and
return. ``cfg.embed_quant`` is settled by
:func:`~arbi_serve.weight_quant.embed_quant.resolve_embed_mode` (the same
decision the load seams took); explicit-mode misconfiguration fails loud
there.

Construct the AuthManager from CLI flags.

Exactly one of ``auth_tokens_file`` / ``auth_verifier_module`` may
be set. When neither is set, the engine runs in OPEN mode
(verifier=None, AuthManager.enabled == False) and the CLI already
printed the "no auth configured" warning.

Publish the DENSE flat-weight cache at the pre-head-quant seam.

Runs on EVERY rank right after ``_load_model_phase`` (dense weights are
loaded, unpacked to the runtime weight form (exl3 trellis-decode / AWQ→Marlin / NVFP4), compacted into the weights slab, and the
derived RMSNorm-fold / GDN-dequant buffers are built) and BEFORE
:func:`_apply_head_quant_phase`. The captured ``state_dict()`` is therefore
the pure DENSE serving set — with NONE of the head-quant-derived tensors
(the INT4 drafter lm_head copy, the INT4 ``mtp.fc``), because the transform
has not run yet. That is exactly the tensor set a warm boot's freshly-built
``skip_weight_load`` graph carries at this same seam, so the dump
round-trips; ``_apply_head_quant_phase`` then re-runs on the warm boot too
(it is unconditional in :func:`build`), reproducing the INT4 heads
byte-for-byte on top of the DMA-filled dense weights.

The cache key (:func:`default_flat_cache_dir`) folds in model content +
parallel topology + rank but NOT the head-quant mode, so an
``--lm-head-quant off`` boot and the default ``drafter-int4`` boot SHARE one
dense cache: the expensive cold weight-unpack is cached once and the
cheap head-quant transform is recomputed every boot.

This is the only seam where the dump can capture the pure dense set:
``drafter-int4`` REPLACES ``mtp.fc`` with an INT4 module and ``full-int4``
frees the dense bf16 lm_head shard, so a post-transform It
runs per-rank in the boot process, matching the per-rank warm LOAD, and
is guaranteed complete before serving, so no idle-gate abort can leave
cold boots permanently slow.

Cold-boot cost only (one write per fingerprint); warm boots already have
the blob and skip. Best-effort by design: a dump error logs LOUD and leaves
the boot cold — it never blocks serving and never risks a wrong graph (the
warm LOADER fail-loud-checks every filled byte against the dump signature).

Load the serving-path triton kernels BEFORE the deferred KV grow.

Several serving kernels compile/load LAZILY on a code path that runs AFTER
the deferred KV grow has filled the card and AFTER ``freeze_for_serving``
snapshotted the torch-caching overhang: the readiness gate's single-row
small-token prefill (warmed by nothing else on the eager path —
``_compile_warmup_prefill`` only warms multi-row ``B=2`` prefills at
``{128, 1024, cap}``), and the ``serve-kernel warmup`` phase's sampler /
TKV-module / GDN-chunk-map kernels. On a fresh boot those are already warm
from the pre-grow compile pass, so the post-grow phase is a cache-hit
no-op; but a cross-model member SWAP drops the incoming member's kernel
blind to them — it under-counts on EVERY boot (never self-corrects) — and

Loading them HERE — before the grow — makes them resident while free VRAM
still exists: the grow's ``mem_get_info`` read accounts for them, the
freeze measures them into the persisted overhang, and the later gate /
serve-kernel phase replay resident kernels allocating zero new bytes. All
ranks run identically (lockstep-safe). Best-effort throughout: a failure
here costs the first matching request its compile latency, never a wrong
result — and the readiness gate still guards.

`` (N servable tokens)`` for a diagnostic, or ``""`` if unknown.

Renders NOTHING rather than a guessed token count: this runs on a boot
that is already failing, and a plausible-looking wrong capacity is worse
than none. Never raises.

Materialize every lazily-fused derived weight, inside ``pool``.

Returns the number of fused tensors built.

A fused cache is a WEIGHT: persistent for the process life, read by the
compiled/captured forward, and therefore address-bound once a cudagraph
records it. Building it here — in the load path, inside the model's weight
pool — keeps it out of the per-step forward arena, which is a scratch pool
the boot may destroy to return its idle physical. A tensor a captured graph
reads must never live in a pool that can be unmapped or re-homed under it.

``pool`` is a :class:`~arbi_serve.runtime.named_pool.NamedMemPool` (or
``None`` for the ambient allocator, which is what the CPU tests and any
pool-less engine stub use).

Build every per-module derived weight that the compiled forward
reads, ONCE, in the load path (post weight-load + compaction, pre
compile/capture).

Two derived weights move here off the per-forward lazy path so the
fullgraph forward reads a stable tensor unconditionally (a lazy cache
with a Python ``is None`` / device / dtype branch inside the captured
region forces a Dynamo recompile → ``FailOnRecompileLimitHit``):

  * :class:`Qwen3_5RMSNorm` — the ``1.0 + weight.float()`` fold
    (``prepare_for_compile`` → ``folded_weight`` buffer).
  * :class:`GDNBlock` — the fused dense ``[B|A]`` AWQ-dequant weight
    (``prepare_fused_ba_dequant`` → ``_fused_ba_dequant`` attribute).
  * the lazily-fused projection weights (:data:`_FUSED_WEIGHT_BUILDERS`)
    — ``[Q|gate|K|V]``, ``[Q|K|V|Z]``, ``[B|A]``, ``[gate|up]``. Built in
    the ``model.weights`` pool so a persistent, graph-referenced weight
    never lands in the per-step forward arena.
  * :class:`FP8LinearBase` — the RowWise ``scale_b`` (1, N) fp32 vector
    the ``torch._scaled_mm`` decode path reads
    (``prepare_scaled_mm_scale`` → ``_scale_b_rowwise`` buffer). Rebuilt
    here post-device-placement so a warm flat-dump boot (which restores
    the persistent ``weight_scale`` without re-binding via ``fp8_load``)
    repopulates the NON-persistent derived buffer before capture.

Fails loud — no silent fallback — if a module raises while building
its derived weight; an unexpected un-buildable module is a bug, not a
survivable condition.

Run ONE real eager forward over a synthetic batch on the REAL engine
state (real model, real pool, real metadata builders, real attn ops).

Shared by the recurrent-kernel warmup and the boot readiness gate. The
synthetic rows index page 0 / the low recurrent rows only (same contract
as the capture sweep), so callers must restore the boot-zero invariant
afterwards via ``_zero_recurrent_slabs_after_capture``.

DRAIN BEFORE BUILD, for the same reason the serving path drains in
``forward_exec.build_metadata`` and the capture sweep drains in
``capture_admin.drain_recurrent_flushes_before_capture``:
``GdnMetadataBuilder._finalize`` REFUSES a metadata build with recurrent
zero-clears still pending, because the GDN prefill path gathers ``h0``
unmasked and would read a freed request's stale state.

LATENT, NOT OBSERVED. Both branches of the build below reach a metadata
build (``_run_metadata_builders`` on the piecewise path, a direct
``builder.build`` loop on the eager one) and neither sat downstream of a
drain. Instrumenting ``has_pending_zero_clears()`` at entry on a GDN
hybrid at ``prefill_capture=full`` read False on both entries and the boot
succeeded, so nothing has been shown to queue a clear this early on a cold
boot — this closes a gap in the invariant, it does not fix a reproduced
failure. It is worth closing anyway because the queue is a property of
admission rather than of this function, and every other build site in the
tree already honours the contract.

Unconditional rather than gated on the model having recurrent layers: the
drain is idempotent and costs one attribute probe when nothing is queued.

WHERE IT ALLOCATES: ``scratch.forward_arena``, the pool the step this
stands in for allocates from. ``forward_exec`` runs the served metadata
build and ``model.forward`` inside ``_activation_arena_pool_ctx`` and the
activation profile probes that same pool, so every reserve the KV sizing
states over a serving step is stated over it. A forward here that ran
outside it would be a tenant of ``unpooled.torch_default_pool`` — ours,
live, and claimed by no budget — and the two pools are separate private
allocators, so an idle block in one cannot serve an allocation aimed at the
other. That matters at exactly this seam: the KV sizing's settle hands the
default pool's free blocks back to the driver (``empty_cache`` reaches that
pool and skips a private ``MemPool``) and the grow gives them to KV, so a
forward that runs after the grow must find those bytes on a card the grow
has filled, while the arena's own idle blocks — which nothing else can
spend — sit beside it. The scope is the whole build-and-forward rather than
any one allocation site: an inductor extern call (``extern_kernels.mm``) is
not a site this tree can wrap, so homing the region is what keeps the next
one accounted too.

Build the synthetic batch and run the forward, in the caller's pool.

Split from :func:`_run_synthetic_forward` so the pool context wraps the
batch build, the metadata build and the forward as ONE scope. The served
step homes the same three (``runtime._batch_build_materialize`` for the
fresh per-step tensors, ``forward_exec`` for the builders and the forward),
and a boot forward that homed only some of them would leave the rest in the
class this routing exists to empty.

Run a few synthetic forwards to warm FLA Triton kernels.

Hybrid models (GDN / Mamba / ShortConv) compile Triton kernels on
first invocation. If the first live request triggers compilation,
the compilation latency and non-deterministic intermediate state
can cause the first-request output to diverge ("req[0] divergence").
Running 2 tiny forwards here — a small prefill and a single-token
decode — forces kernel compilation at boot so the first live
request gets a steady-state kernel.

No-op when the pool has no recurrent views.

Is ``exc`` an allocator failure, i.e. does a memory verdict fit it?

Walks the ``__cause__``/``__context__`` chain because the allocator
failure is frequently wrapped — a cuBLAS handle creation that dies with
``CUBLAS_STATUS_ALLOC_FAILED`` surfaces as a ``RuntimeError`` whose text
is the only evidence, and a torch ``OutOfMemoryError`` raised inside a
custom op arrives wrapped by the dispatcher.

Deliberately NARROW. The default answer is False, so an unrecognised
failure keeps its own type and its own traceback: the failure mode this
guards against is a correctness refusal reported as a budget problem,
and that mode is only reachable by classifying too broadly. A memory
failure misclassified as "other" still fails the boot loudly with the
same numbers attached — it just does not claim to know the remedy.

Boot readiness gate: prove ONE real eager forward works BEFORE the
server declares "serving now".

Runs a small PREFILL forward through the REAL model / pool / metadata
builders / attn ops — the FIRST thing a live request does (prompt
prefill, before any decode step). If it raises — e.g. a capture-OOM
cascade left the card with ~0 B free, so the first forward dies at
``cublasCreate`` with ``CUBLAS_STATUS_ALLOC_FAILED`` and every request
would error forever — the boot FAILS LOUD with the VRAM/budget
breakdown and remediation instead of serving a guaranteed total
outage. This gate is NOT best-effort and is NOT relaxed by
``vram_mode``: a server that cannot run a single forward must never
report ready.

Why PREFILL, not a single-token decode: a direct ``eng.model.forward``
of a DECODE shape replays the torch.compile cudagraph for that shape
without the ``execute`` path's persistent-buffer ``copy_()``-refresh +
serialization, so on the compiled dense-``qwen3`` decode path it races
the capture-time buffers → an async out-of-bounds index (device-side
assert at the last-token gather) that ONLY the gate hits — live decode
goes through ``model_runner.execute`` (captured-graph replay) and does
not hit this race. A prefill forward is both more faithful (the first live request
prefills first) and dodges the direct-forward decode-replay hazard,
while still proving weights load + KV pool + attn op + cuBLAS + the
activation arena all work (it is the strictly heavier allocation path,
so an OOM that would sink decode sinks prefill here first).

Informational probe of free GPU memory at boot.

The deterministic budget (:func:`profile_and_size_kv_pool`) is the
tight gate; this just logs free vs total VRAM so operators can
spot "obviously wrong card" misconfiguration before paying the
weight DMA cost.

Return the bytes THIS process holds on ``device_index`` per NVML.

The same counter ``nvidia-smi`` shows in its per-process pmon view,
including CUDA context, NCCL buffers, cuBLAS workspaces, AND
torch's reserved segments. Returns 0 when NVML cannot say which of the
card's processes is this one — see
:func:`~arbi_serve.engine.boot_vram_drain._self_used_bytes`, which owns
that identification for every reader of it.

Wrap a boot region; on exit, log how much GPU memory it added.

Reads ``nvidia-smi``'s per-process counter via NVML, so the delta
includes EVERYTHING the region cudaMalloc'd. Failures are swallowed
so the wrapped block always runs.

Zero every recurrent / short-conv slab row after the capture sweep.

The boot-time cudagraph capture sweep (decode + prefill + piecewise)
runs synthetic forwards through the recurrent blocks, leaving
non-zero garbage in rows ``0..max_num_seqs-1`` of every
per-layer slab via the captured ``index_copy_`` write-back. The
canonical ``has_initial_state_for(batch) = seq_lens > 1``
derivation flags every non-trivial prefill row as "carrying
prior state", so the first live prefill picks up the capture-
sweep garbage as its recurrent initial state and decode loops
into degenerate-token output (the "Mr. G. G. G." failure on
Qwen3.5-0.8B).

Dispatches to :meth:`MultiStatePool.zero_all_slot_rows`, which reaches
every slot-state kind present. No-op on pure-attention pools. Cheap:
each call is one ``zero_()`` per slab, batched within a single CUDA
stream.

Boot-path only. A stable-VA sleep/resume cycle runs no capture
sweep — it remaps the same VAs and restores the slabs' offloaded
contents — so there is no capture residue to clean up on wake, and
re-zeroing there would discard the recurrent state of sequences the
cycle is designed to preserve.

capture_admin: cudagraph capture orchestration, split by phase.

Re-exports every public and underscore symbol so
``from arbi_serve.engine.capture_admin import <name>`` resolves the whole
family. The legacy module ``arbi_serve.engine.cudagraph_admin`` re-exports
this package for back-compat.

Shared capture-admin primitives: logging, sweep counters, failure /
degrade helpers — plus a re-export surface for helper groups defined in
sibling modules:

  * ``_shapes``    — capture-shape / ladder / bucket derivations
  * ``_mtp_fill``  — MTP drafter-KV fill predicates
  * ``_plans``     — config-derived prefill / piecewise capture plans
  * ``_synthetic`` — synthetic-batch build + the piecewise bucket capture

``_common`` re-exports every name from those modules at the bottom of
this file, so ``from arbi_serve.engine.capture_admin._common import
<name>`` resolves for the family modules (decode / mixed / prefill /
drafter_chain / layer / compile_warmup) and the package ``__init__``.
``logger`` is defined here before that re-export block runs, so the
siblings can import it from ``_common`` without an import cycle.

Poison every capture pool after a cudagraph capture RAISED.

A capture that raises mid-``torch.cuda.graph(...)`` can leave a PARTIALLY-
captured ``torch.cuda.CUDAGraph`` holding DEVICE pointers into the graph
pool (``capture.cudagraphs``) and the baked I/O buffers (``capture.
io_buffers``) — and that partial graph is typically NOT stored in the
engine's capture-pool dict, so ``clear_all_cudagraphs`` never ``reset()``s
it. Latching :meth:`NamedMemPool.mark_capture_failed` on these pools makes
teardown park them permanently (``free_all`` never unmaps their physical;
``reap_leaked_pools`` never destructs them), so a later finalization of the
partial graph cannot fault on freed device memory. Best-effort: a stub
engine / missing pool is a clean no-op.

A capture failed in a way that leaves the capture pools unusable.

Raised by :func:`_note_capture_failure` so the sweep that was walking a
ladder STOPS instead of queueing the next shape behind a half-recorded
graph. The distinction it encodes is the one
:func:`_mark_capture_pools_failed` already documents as its trigger — a
capture that raised mid-``torch.cuda.graph(...)`` can strand a PARTIALLY-
captured ``CUDAGraph`` holding device pointers into the capture pools —
and the next capture into those pools inherits the corruption, surfacing
as a ``cudaErrorIllegalAddress`` several shapes later with nothing in the
log pointing back at the shape that actually broke.

True when ``exc`` says the CARD ran out, not that the capture is broken.

Two channels, because a capture allocation can fail through either:
torch's caching allocator (``is_oom_error`` — the type, plus the driver /
cuBLAS message markers), or the cuMem / CUDA driver directly, which
allocates outside the caching allocator and can surface a genuinely-full
card as a capture-invalidation error whose message never says "memory".
The authoritative fallback for that second channel is the card itself: if
free VRAM at the failure point is at or below the runtime floor, the card
was full regardless of which error string the driver chose.

Shared by :func:`_raise_on_tp_capture_failure` (which picks the operator
message) and :func:`_note_capture_failure` (which decides whether the
sweep may continue), so the two cannot answer the same question
differently.

A graph that recorded the WRONG ARITHMETIC is not a capture failure.

Every capture sweep's fallback -- note it, skip the bucket, serve that
shape eager -- is the right answer to a bucket that COULD NOT be captured:
an unsupported op, an allocator denial, a shape the backend refuses. It is
the wrong answer to a bucket that captured perfectly well and recorded
numerics the build did not declare.

Two things go wrong if such a refusal is caught. The operator asked for
captured prefill and silently gets none of it, at a WARNING. And the
warning names capture as the problem when the problem is a forward that
reached the op without declaring its phase -- a wrong diagnosis, which is
worse than a crash because it is actionable in the wrong direction.

Walks the cause chain: the refusal is raised from a context manager's exit
and can reach a sweep wrapped by another teardown's own error. Terminates
on a cycle, and is a no-op for a build with no EXL3 in it rather than
importing the module to ask.

Fail fast when a per-rank capture failure cannot be skipped.

Under TP>1 the capture sweeps run in rank-lockstep and every captured
forward contains NCCL collectives. A rank that fails (e.g. OOM) and
skips the shape leaves its peer blocked inside the shape's collective,
and the boot then dies in the NCCL watchdog with the real cause
buried. There is no rank-symmetric recovery from an asymmetric skip,
so raise immediately and loudly instead.

Tally a capture failure, and REFUSE to continue the sweep when it
poisoned the capture pools.

Every per-shape capture except-handler routes through here, so this is
where the sweep's two very different failure modes are separated:

**A shape that could not be captured.** A pre-flight refusal, a config
error, a scratch pool sized too narrow, a card too small — the failure
happened before the ``torch.cuda.graph(...)`` record region opened, or it
was a memory shortfall. Nothing is half-written; the shape runs eager and
the sweep walks on. A CUDA OOM additionally increments
``_capture_oom_failures``, which the post-capture VRAM gate refuses the
boot on, and poisons the pools so teardown never unmaps physical a partial
graph still holds.

**A capture that aborted MID-RECORD.** The record region was open and CUDA
tore it down underneath the sweep — a pageable H2D, a host sync, an
unsupported op. This is NOT a skipped bucket: it strands a partially-
captured ``CUDAGraph`` holding device pointers into the capture pools
(:func:`_mark_capture_pools_failed` documents exactly this trigger), and
the NEXT capture into those pools inherits the corruption and dies with
``cudaErrorIllegalAddress`` — one shape later, with nothing in the log
naming the shape that actually broke. Observed on the mixed sweep, where
the continuation family aborted the record and the first-chunk family
queued behind it took the boot down at ``grow_kv_after_capture``.
:exc:`CaptureContextPoisoned` stops the sweep at the shape that broke.

The discriminator is the CUDA stream-capture state at unwind, read by
:func:`~arbi_serve.runtime.capture._common._force_end_stream_capture` from
``cudaStreamEndCapture`` and carried on the exception. Not "it threw
during capture", which cannot see that a config refusal never opened the
region; not "it was an OOM", which is one cause of a mid-record abort and
not the interesting one. A device-side assert / illegal address counts
too, wherever it was raised: it corrupts the context whether or not a
record was open.

A memory shortfall is carved OUT even when it aborted the record. It is
the one mid-record abort the boot already prices — ``_capture_oom_failures``
for the VRAM gate, ``boot_substitutions`` for the persist guard — and a
card too small for a bucket must still degrade to eager rather than refuse
a boot.

When ``shape_desc`` is given, the failure is ALSO recorded (with the
exception object) on ``eng._capture_failures`` — the attempted-and-failed
bucket ledger :func:`enforce_decode_capture_integrity` fails the boot on.
Only the DECODE/VERIFY sweep and ``precapture_compile_warmup`` pass a
``shape_desc``: those buckets are implied by the serving config, so a
failure there is a deterministic boot failure, not a best-effort skip.

Boot-fail if a verify shape that was expected captured is missing.

``expected_verify_keys`` holds ``(B, S, lora_bucket, kv_pages_bucket)``
for every MTP verify shape that PASSED the verify pre-flight during
the sweep. Each must resolve in ``eng.captured_graphs``; a miss means
the verify forward will fall through to the eager (uncaptured) path
for that shape while the engine pretends to serve at full speed.
cudagraph is the serving config — refuse to boot, naming the shapes,
rather than serve degraded.

Pre-flight REFUSALS (EXL3 / MLA / out-of-template-K) are NOT in the
set — those are legitimately eager-by-design and were already logged
loud at WARNING.

Boot-fail when a SCHEDULABLE decode/verify shape has no captured route.

The capture ladder is derived from this boot's own limits — the B rungs
from ``cfg.batch.max_batch``, the verify width from the driver's served
``K`` — and both derivations are free to be sparse, because an off-rung
``B`` pads up to the next captured rung rather than running eager. That
makes the ladder's coverage an EMERGENT property of three separate pieces
of code (the rung generator, the pad-up lookup, and admission), and a
change to any one of them can open a hole that nothing else reports: a
shape with no route runs the live per-layer forward on every step of every
request that hits it, at unchanged output and unchanged acceptance length.

This gate closes that. :func:`schedulable_decode_shapes` enumerates the
``(B, S)`` space the scheduler can present (see its docstring for the
admission proof), and
:func:`uncovered_schedulable_decode_shapes` reports any that resolve to
neither an exact rung nor a pad-up target — checked against the pool the
runtime will actually resolve against, once per captured
``(lora_bucket, kv_pages_bucket)`` class.

``ARBI_ALLOW_DEGRADED_CAPTURE=1`` downgrades to a registered degradation,
the same contract as :func:`_assert_required_verify_captured`. So does a
disarmed ``ARBI_DECODE_PAD_CUDAGRAPH``: that removes the pad-up route the
sparse ladder is designed around, so the resulting gap is an operator
choice about an unchanged ladder rather than a derivation bug.

MTP drafter-KV fill predicates for capture-variant selection.

The boot-constant decisions on whether a captured decode / prefill graph
bakes the MTP drafter-KV fill, plus the fail-loud audit that every
captured graph's baked variant matches live routing. No sibling imports.

Boot-constant decision: bake the MTP drafter-KV fill into a
decode-family capture at ``seq_len``, or elide it.

Decode has NO replay-side fallback (cudagraph is the only decode
path — eager decode is a severe regression), so the captured variant
must equal what the live batches that replay this shape carry:

  * ``seq_len == 1`` — plain decode. NOTHING consumes an MTP
    drafter-KV write at S=1, so the fill is dead compute on every
    decode step of every server whose checkpoint bundles an
    ``mtp_head``. The drafter-KV cells a draft chain attends to are
    written by the two paths that DO carry the fill: the verify
    forward (``seq_len > 1`` below, where it is load-bearing) and
    the eager MTP seed forward on the tick a row finishes prefill.
    A decode batch may still arrive carrying
    ``mtp_fill_enabled=True`` — the SPMD loop derives that flag
    slate-wide from "any row opted into MTP" — and replaying the
    fill-less graph is exactly right for it: the flag says the
    slate has MTP traffic, not that this S=1 forward owes anyone a
    KV write.
  * ``seq_len > 1`` — MTP verify shape. The verify batch always
    carries ``mtp_fill_enabled=True``
    (``batch_build.reconstruct_verify_batch``) and the fill is
    LOAD-BEARING: it writes the K+1 accepted-token KV cells the
    next draft chain attends to. True whenever the model can run
    the fill (bundled head, no DFlash tap) — under DFlash the
    model-side tap gate skips the fill in eager too, so False
    stays faithful there.

:func:`assert_captured_variant_policy` re-derives this expectation
per captured graph after every sweep and hard-fails the boot on
disagreement — no silent variant mismatch can reach serving.

The variant the live routing will present to ``graph``'s replays.

Shape-classes (mirrors the pool key semantics):

  * MIXED (``is_prefill=True`` with ``flat_tokens>0``) — False:
    ``split_mixed`` pins ``mtp_fill_enabled=False`` on the batch it
    builds, so the routing genuinely presents False here.
  * pure prefill (``is_prefill=True``, ``flat_tokens=0``) —
    :func:`_prefill_capture_mtp_fill`, the same predicate
    ``lookup_captured_graph_for_prefill`` evaluates against the live
    batch, so the audit compares the capture decision against actual
    routing rather than against a constant that mirrors the capture
    decision itself.
  * decode family — :func:`_decode_capture_mtp_fill` at the
    graph's ``seq_len``.

Boot-time fail-loud audit: every captured graph's baked MTP-fill
state must match what the live routing will carry to it.

Decode cannot refuse-to-eager at replay (cudagraph is the only
decode path), so a wrong variant would either silently re-bake the
dead fill (perf lie) or starve the verify drafter of its KV writes
(accept collapse). Walked once per capture sweep at boot — never on
the serving hot path.

Will a prefill REPLAY caller ask this graph for ``hidden_out``?

This is the capture-side half of the lookup's ``require_hidden`` gate
(:func:`arbi_serve.runtime.captured_lookup_gates.lookup_captured_graph_for_prefill`).
One replay caller passes it: ``EagerModelRunner._seed_forward_or_replay``,
the MTP-opted final prefill chunk that seeds the drafter, and it passes it
UNCONDITIONALLY. So the question a capture must answer is not "is the fill
baked" but "does this configuration run a drafter seeded off a prefill
chunk" — true for the bundled ``mtp_head`` AND for a DFlash tap.

It used to be answered with ``retain_hidden = mtp_fill``, on the premise
that a hidden-consuming replay is "reachable exactly when the fill variant
is baked (bundled head, no tap)". The tap breaks that premise in one
direction only, and silently: a tap owns the MTP slab, so
:func:`_prefill_capture_mtp_fill` correctly declines the fill, no graph
retains a hidden, and the seed forward — which still wants one — is refused
on every single chunk while the boot reports a successful capture. Measured
on 27B + DFlash2 at 16k: ``prefill_graph_lookup`` fired 0 / refused 668, of
which 658 were ``hidden not captured``, over a pool that cost 916 MB and
130 s of boot (#2233).

Retention is a ``(B, H)`` last-slot gather, kilobytes — see the note in
``capture_prefill`` on why the full ``(N, H)`` is not retained. The thing
the old coupling conserved was that; the thing it lost was every replay.

Boot-constant decision: bake the MTP drafter-KV fill into the
captured prefill graphs, or elide it.

Mirrors the ``fill_would_run`` predicate the lookup's variant gate
evaluates (:func:`lookup_captured_graph_for_prefill`): bundled
``mtp_head`` present and no DFlash tap owns the MTP slab. The two must
agree — a captured variant the live routing never asks for is refused on
every replay ("mtp-fill variant mismatch") and its capture VRAM is dead.

The earlier chunks of an MTP-opted request sample nothing, take the
normal ``return_hidden_state=False`` path, reach the lookup, and need
the fill baked. The final prefill chunk (the one that seeds the drafter
and needs ``return_hidden_state=True``) is also a consumer: fill
variants are captured with ``hidden_out`` retained (``capture_prefill``)
and the seed forward replays them through
``EagerModelRunner._seed_forward_or_replay``.

Baking the fill for them is not the DFlash dead-fill class: the fill is
exactly the work those chunks must do, and the eager path does it anyway.
Non-MTP slates (``mtp_fill_enabled=False``) take the mismatch side and
fall back to eager — correct in both directions, never corruption.

Config-derived capture plans (whole-forward prefill + piecewise).

The resolved dataclasses + resolvers consumed by both the boot
orchestrator and the boot banner, so what the banner prints is what the
sweep captures. Imports only ``logger`` from ``_common``.

Resolved whole-forward prefill capture decision, derived from config.

The single source consumed by both :func:`precapture_prefill_graphs`
(which sweeps ``eligible``) and the boot banner
(``engine.build._log_boot_config``), so what the banner prints is the
ladder the sweep captures by construction. ``skip_reason`` is non-None
when no sweep will run at all; the model pre-flight
(:func:`_can_capture_prefill`) is deliberately NOT part of the plan —
it needs the loaded model and is attributed on the
``prefill_capture_sweep`` counter at sweep time.

Derive the prefill capture plan from config (no model needed).

Mirrors the gate order of :func:`precapture_prefill_graphs`:
``cuda_graphs`` → ``prefill_capture == "full"`` → configured buckets →
``chunk_prefill`` auto-extension →
``prefill_cap = min(max_batched_tokens, max_context)`` clamp. The plan is
world-size-agnostic: whole-forward prefill capture is governed by the
configured buckets at every world size (TP>1 included), the same as the
piecewise per-layer sweep. TP>1 correctness is boot-enforced by the tkv
``_preflight_page_metadata_stride_contract``. ``world_size`` is retained
in the signature for the boot banner / caller symmetry.

Name the configured rungs the ``prefill_cap`` clamp drops.

Level policy: WARNING when the operator explicitly configured the
ladder (a user-set rung that silently doesn't exist is a config lie);
INFO for the compiled-in default set (its top rung is sized for
larger-``max_batched_tokens`` configs and trimming it to the config
is the designed derive-from-config behaviour). Either way the rungs
are named in the log line — the configured set and the captured set
must never diverge silently.

Resolved piecewise per-layer prefill capture decision, config-derived.

The single source consumed by both the boot orchestrator's call-site
condition context and the boot banner (``engine.build._log_boot_config``),
so the banner names the state the boot actually resolves — including the
TP>1 resolution, which is otherwise invisible until the sweep logs.
Model-dependent refusals (CUDA availability, the recurrent-kind
pre-flight, capture-unsafe attention) are deliberately NOT part of the
plan — they need the loaded model and are attributed on the
``piecewise_sweep`` counter at sweep time.

Derive the piecewise prefill capture plan from config (no model needed).

Mirrors the config-level gate order of the boot orchestrator
(``engine.build.build``): ``cuda_graphs`` → ``prefill_capture == "piecewise"``
→ TP resolution. TP>1 does not disable the sweep and does not require
``ARBI_DFLASH_TP_SHARD`` — see the lockstep-safety invariant stated in
:func:`precapture_layer_graphs`.

Flag-truth refusal reason for ``prefill_capture_sweep`` when the
whole-forward sweep's call site (``engine.build._build_impl``, inside
its ``eng.cfg.cuda_graphs`` region) does not even invoke
:func:`precapture_prefill_graphs` this boot — or ``None`` when it does
(the sweep itself then owns any further refusal, e.g. the model
pre-flight or "no buckets fit prefill_cap=...").

Mirrors the two checks the call site's OUTER gate makes BEFORE calling
:func:`precapture_prefill_graphs`:

  - ``prefill_capture != "full"``: the mode wasn't selected — a config
    decision, not a broken sweep. This is also why flipping
    ``gdn_prefill_capture`` alone (a ``RuntimeFlags`` field, scope
    "backend" — it DOES trigger a member rebuild) can never move this
    counter: the flag only matters INSIDE the sweep (the
    ``_can_capture_prefill`` GDN-hybrid opt-in), and the sweep is gated
    off here whenever ``prefill_capture`` is not ``"full"`` — a
    separate, coarser config knob the flip does not touch.
  - ``prefill_capture == "full"`` but no buckets configured: the raw
    (unclamped) ``prefill_cudagraph_buckets`` tuple is empty, so the
    call site's own ``and cfg.prefill_cudagraph_buckets`` short-circuits
    BEFORE :func:`precapture_prefill_graphs`'s identical internal check
    (the ``if not buckets`` refusal) ever runs — making that internal
    check dead code from this call site. Marked here instead so it
    isn't silently unreachable.

Without this, ``prefill_capture_sweep`` reads a flat zero whenever
either of these two config states holds — indistinguishable from "the
sweep should have fired and silently didn't". Mirrors the DORMANCY vs
BREAKAGE marker the boot orchestrator already applies to the sibling
piecewise sweep's call site for the identical reason.

Auto-extend the prefill bucket ladder UP to cover ``chunk_prefill``.

A prefill chunk of ``chunk_prefill`` tokens exceeding the top captured
bucket finds no graph and runs eager every step — the same silent-footgun
class as the decode ladder vs ``max_batch``. Adding ``chunk_prefill`` as a
rung keeps the modal chunk on a captured graph. No-op when the ladder
already covers it; the caller's ``prefill_cap`` clamp still bounds the
result to ``min(max_batched_tokens, max_context)``.

Shared capture-shape derivations (ladders, buckets, cross-products).

Pure, config-driven helpers the decode / verify / drafter / piecewise
sweeps all derive their shapes from — one source of truth so a drift can
never silently drop a captured shape. Imports only ``logger`` from
``_common`` (defined before ``_common``'s re-export block runs), so there
is no import cycle.

Effective dense-ladder ceiling for ``max_batch``.

``ARBI_DENSE_B_LADDER_MAX`` (when set to a positive int) is an explicit
override. Otherwise the dense band is the FIXED low band ``[1, 8]``
(clamped at ``max_batch``): every B up to 8 is an exact captured rung, and
everything above is served by the geometric+stride rungs plus pad-up.

Fixed, not proportional to ``max_batch``. A ceiling that tracks the
served concurrency (``min(max_batch, 64)``) is a trap; see
``_DENSE_AUTO_MAX_BATCH`` for what it costs (at ``--max-batch 64``
with MTP K=5 on TP2 it makes the
boot REFUSE) and what the dense rungs buy (nothing measurable — pad-up
at 78% pad overhead is free to within a percentage point).

Still monotone in ``max_batch``: raising max_batch never REMOVES a dense
rung, it only clamps the band on a launch narrower than 8.

Generate the plain-decode (S=1) batch-size ladder for ``max_batch``.

The ladder is a DENSE low band ``[1, 8]``, then the geometric base
``{16, 24, 32}``, then a ``+8`` stride up to and including ``max_batch``:

  * ``max_batch=8``   → ``{1,2,3,4,5,6,7,8}`` (all dense),
  * ``max_batch=16``  → ``{1..8, 16}``,
  * ``max_batch=32``  → ``{1..8, 16,24,32}``,
  * ``max_batch=64``  → ``{1..8, 16,24,32,40,48,56,64}``,
  * ``max_batch=128`` → ``{…,64,72,80,…,128}`` — auto-extended, no
    literal to bump.

Every rung is ``<= max_batch`` and ``>= 1``; ``max_batch`` itself is
ALWAYS the final rung. That last property is the coverage guarantee, not
density: with ``cfg.decode_pad_cudagraph`` armed an off-rung B pads up to
the smallest captured rung and replays THAT graph
(:func:`~arbi_serve.runtime.captured_lookup.lookup_captured_graph_for_forward_padded`),
and a pad target exists for every admissible B precisely because the
ladder terminates on ``max_batch``. One shape per ``(S, lora, kv)`` class
would suffice for coverage; the rungs are purely a pad-waste/perf choice,
and what they cost and buy is measured in ``_DENSE_AUTO_MAX_BATCH``.

MTP verify-shape ``S = K + 1`` set, bounded by the served verify
block_m ceiling ``served_bm``.

``full_ladder=False`` (the default) returns only the served depth
``{max_k + 1}`` — the single verify width a homogeneous deployment
replays (every request at ``sampling.mtp_k == n_draft``; the ``S=1``
collapse is served by the plain-decode graph). The verify sweep is
then ``|B-ladder| × 1`` instead of ``|B-ladder| × max_k`` graphs, and
each captured graph retains its own working set in the shared
``graph_pool`` (so the pool scales with the graph count).

``full_ladder=True`` returns the full admissible ladder ``{k + 1 : k ∈
[1, max_k], k + 1 ≤ served_bm}`` — every per-request ``K`` a client may
opt into has a captured graph, at the cost of the ``max_k×`` larger
pool. Gated by ``cfg.mtp.capture_full_k_ladder``.

``max_k + 1`` is clamped at ``served_bm`` in both modes (admission
rejects ``K`` whose ``K+1`` exceeds the served ceiling, so a clamped-out
top width is never replayed).

A TREE has no ``K`` axis to ladder over, so ``full_ladder`` is ignored
under one. ``tree_step_k`` resolves EVERY chain depth to the same node
count, and ``active_verify_block_tables`` hands out tree geometry only
for a block whose length is that count — so a rung at any other width
captures the CHAIN forward at a width no tree step ever presents, and
no tree step can replay it. Walking the ladder anyway reads the node
count the callers already substituted for ``max_k`` as though it were
a chain ceiling, which is what turns a 1-rung sweep into an N-rung one
and prices the difference straight out of the KV pool (the cold-boot
graph-pool bound sums the per-shape working set).

Captured KV-page bucket ladder (int page counts), or ``()`` when disabled.

Clamps ``cfg.cudagraph_kv_pages_buckets`` at ``max_context_pages``, dedupes,
and pins the top bucket to ``max_context_pages`` so long-context decode never
falls to eager. ``()`` (no buckets configured) is the no-bucket path — the
decode sweep maps it to the single ``None`` sentinel; the count helper maps
it to 1.

LoRA capture bucket ladder.

``(0,)`` (the no-LoRA baseline) unless a LoRA capture pool is wired AND
``cfg.lora_max_loras > 0``, in which case it is the ``ACTIVE_LORA_BUCKETS``
rungs ``<= lora_max_loras``. Every captured ``(B, S)`` shape is replayed once
per bucket, so this multiplies the graph-pool reserve.

Resolve the full ``(B, S)`` decode-capture sweep for boot.

The plain-decode (``S=1``) B-rung ladder is generated from
``cfg.batch.max_batch`` via :func:`_decode_b_ladder` (not read off
the literal ``cfg.cudagraph_shapes`` tuple), so raising max_batch
auto-extends the rungs rather than being bounded by a fixed tuple.
The operator may still pin extra shapes in ``cfg.cudagraph_shapes``
(verify ``(B, K+1)`` rows or off-stride B) — those are unioned on
top of the generated ladder.

When the engine is MTP-enabled with a driver attached, the sweep is
further auto-extended to cover every MTP-verify shape ``(B, K + 1)``
for ``B`` taken from the configured ``B`` set and
``K + 1 ∈ [2 .. min(tkv_max_block_m(), max_k + 1)]``. The
auto-extension covers the ``(B, K)`` cross-product the verify pass
actually walks under the uniform-K invariant (``step_k =
min(slate.mtp_k)``; every admitted row already has ``mtp_k <=
driver.max_k`` because admission rejects a larger K rather than
clamping it — see ``engine.submission`` / ``effective_row_k``), so
every step finds a matching captured graph rather than falling
through to the live forward.

Deduplicated and order-stable: the generated S=1 ladder comes
first, then operator-pinned extra shapes, then the MTP-verify
auto-extension appends new shapes at the end, so a diff against a
known sweep sees a stable derivation. This is the DERIVATION order,
not the capture order — :func:`precapture_decode_graphs` walks the
result largest-working-set-first (:func:`_capture_largest_first`), so
the per-shape capture log is descending.

Returns:
    ``[(B, S), ...]`` — every entry has ``S ∈ [1, tkv_max_block_m()]``
    and ``B >= 1``. Out-of-range entries from the operator config
    are filtered out and warned about by the caller's per-shape
    loop (``S`` outside ``[1, tkv_max_block_m()]`` skipped at the
    capture site).

Order a capture sweep LARGEST-WORKING-SET-FIRST.

Every capture in a sweep records into the one shared
``capture.cudagraphs`` mem pool and overlays its transient working set
on blocks earlier captures already freed. Capturing the biggest shape
first therefore reserves the pool's true peak ONCE, and every later
(smaller) shape's transients fit inside blocks that are already mapped;
ascending order instead walks the high-water mark upwards shape by
shape, mapping fresh physical at each new peak. This is the rationale
the piecewise sweep documents and that vLLM's capture loop states as
"memory allocated for smaller graphs can reuse the pool blocks of
larger ones".

``shapes`` are ``(B, extent)`` pairs — ``(B, S)`` for the decode /
MTP-verify sweep, ``(B, K)`` for the drafter-chain sweep. The sort key
is the pair's PRODUCT first: a capture's transient is dominated by the
tokens (decode/verify) or chain steps (drafter) it materialises, so a
narrow-but-deep shape can outweigh a wider shallow one — ``B`` alone is
not the cost. ``B`` breaks ties, because the per-ROW terms (the paged-KV
gather, the persistent per-row buffers) are the part that does not
follow the product.

Scope of the saving: it is bounded by the pool's HELD-FREE bytes after
the sweep — exactly what an ascending walk would strand. Where a
forward's transients come out of a pre-sized arena rather than the
capture pool, the pool ends the sweep fully allocated and the order is a
structural invariant (nothing to strand) rather than a saving. The
ordering still belongs here so a capture family that DOES allocate into
the pool cannot re-introduce the stranding.

Ordering only: the returned set equals the input set.

Derive the ``(B, K)`` bucket cross-product for drafter capture.

Pulls batch sizes from the same generated plain-decode B-rung
ladder the decode capture uses (:func:`_decode_b_ladder`, not the
literal ``cfg.cudagraph_shapes`` tuple) so the drafter capture
lines up with the shapes the verify pass replays and auto-extends
with ``max_batch`` for free. ``K`` set is the served depths ``{1,
n_draft}`` by default (mirrors :func:`_verify_s_set` — a homogeneous
server only ever drafts at K=n_draft, and K=1 doubles as the seed-
drafter call after prefill), or the full ladder ``{1, ..., n_draft}``
when ``cfg.mtp.capture_full_k_ladder`` is set. ``pool.get(B, k)`` is
exact-match on K, so an uncaptured middle K runs the live eager chain
every step (per-collective rank-desync waits under TP>1); under the
default capped ladder, admission rejects an opt-in middle K at EVERY
world size (the uncaptured-verify-shape gate in
``engine.request_factory``), so neither the eager-collective deadlock
nor a TP1 eager verify step can form.

By default (``cfg.cudagraph_max_drafter_shapes`` unset / ``0``) the
full cross-product is returned — every ``(B, K)`` the scheduler can
produce is captured, the pool is unbounded (sized by the graph-pool
budget), and no shape is ever dropped: the fast path engages for
every served request because every served ``(B, K)`` has a captured
graph. The shape count is derived from ``resolved_max_batch`` (the B
ladder) and ``mtp.n_draft`` (the K ladder), never a magic constant.

A non-zero ``cfg.cudagraph_max_drafter_shapes`` is an experiment-only
override ceiling: the cross-product is pre-trimmed to it (matching
the LRU cap :class:`DrafterChainGraphPool` then enforces) so the
sweep never captures a graph just to evict it. Trim priority keeps
every ``B=1`` (c=1) shape first, then ``K=1`` / ``K=n_draft`` for
higher B, then middle Ks descending. A trim is logged with the
dropped buckets and the config knob; a ceiling too small to hold
every B=1 shape is refused.

Number of captured KV-page buckets in the decode/verify cross-product.

Mirrors the bucket derivation in :func:`precapture_decode_graphs`:
``cfg.cudagraph_kv_pages_buckets`` clamped at ``max_context_pages`` (and
the top bucket pinned to ``max_context_pages``), deduplicated. ``()`` (the
no-bucket path) is one bucket (the ``None`` sentinel). Every captured ``(B,
S)`` shape is replayed once per KV-page bucket, so the captured graph count
— and the private graph-pool VRAM it reserves — is multiplied by this
factor. The cold-boot graph-pool estimate must include it or it under-counts
the pool by ``len(kv_buckets)×``, which can OOM a high-K verify capture at
mb≥2.

Number of captured LoRA buckets in the decode/verify cross-product.

Mirrors :func:`precapture_decode_graphs`: the 0-bucket (no-LoRA) is always
captured; an engine with a LoRA capture pool wired adds one bucket per
active rung up to ``cfg.lora_max_loras``. Every ``(B, S)`` shape is captured
once per LoRA bucket, so this multiplies the graph-pool reserve. Q36
(``lora_max_loras=0``) is 1.

The MTP verify ``(B, S=K+1)`` capture shapes derived from config alone.

:func:`_decode_capture_shapes` auto-extends the decode sweep with the verify
cross-product only when the MTP driver + verify buffers are already attached
to the engine. But the cold-boot graph-pool estimate
(:func:`arbi_serve.engine.build.cold_boot_graph_pool_upper_bound_bytes`)
runs in ``profile_and_size_kv_pool``, before ``build_mtp_driver`` attaches
them — so at estimate time ``_decode_capture_shapes`` sees no verify shapes
and the estimate omits the entire verify-graph pool. At high K / mb≥2 that
omission lets KV size too large and the real verify capture can OOM the
card.

This helper reconstructs the same verify ``(B, S)`` set the capture sweep
will produce, from config + the served verify block_m ceiling, with no
dependency on a driver being attached:

  * ``B`` set = the generated decode B-rung ladder (:func:`_decode_b_ladder`)
    ∪ operator-pinned ``cudagraph_shapes`` B ∪ {1}, capped at ``max_batch``.
  * ``S`` set = ``{max_k + 1}`` (the served depth) by default, or the full
    ``{K+1 : K∈[1..max_k], K+1 ≤ tkv_max_verify_block_m()}`` ladder when
    ``cfg.mtp.capture_full_k_ladder`` is set — see :func:`_verify_s_set`.

``max_k`` is the served draft depth (``cfg.mtp.n_draft`` for the bundled
head; the DFlash ``block_size-1`` for the DFlash drafter). Returns ``[]``
when MTP is off or ``max_k < 1``.

True if the model carries any MLA_SHARED layers (DeepSeek V2/V3).

``arbi_serve::mla_attention`` IS Dynamo-traceable: it is declared
through ``torch.library.custom_op`` with ``mutates_args=("output",
"kv_cache")`` and a registered fake, and every per-step value it
routes on is an op argument rather than a Python side-channel
(regression-locked by ``tests/test_compile_op_fakes.py`` and
``tests/test_backends_through_custom_ops.py``). ``device_types="cuda"``
restricts the REAL impl, not the fake, so a trace never reaches it.

Block-level ``torch.compile`` of the MLA decoder layers is a separate,
unvalidated question — the DeepSeek layer classes carry
``@support_torch_compile(enable_if=False)`` and therefore present no
compile-eligible root at all.

Checks ``model.layer_specs`` (outer DeepseekV3Model) and
``model.model.layer_specs`` (inner _DeepseekV3Inner, in case ``model``
is wrapped or the layer_specs are nested).

The ``num_tokens`` bucket ladder the piecewise prefill capture sweep
builds — the single source of truth shared by
:func:`precapture_layer_graphs` (which captures each rung) and
:func:`_compile_warmup_piecewise_prefill` (which pre-warms the eager
kernel surface for each rung on every rank).

Deriving both from one deterministic function guarantees the warm pass
and the capture sweep agree rung-for-rung on every rank — the TP>1
lockstep invariant: a rung the warm pass misses would fire a fresh
(per-rank-variable-latency) FLA / Triton autotune inside the capture
collectives and desync the ranks.

The dense low rungs (:data:`DEFAULT_PIECEWISE_BUCKETS`, <= 512) plus
the high rungs (:data:`PIECEWISE_HIGH_RUNGS`) up to the served prefill
chunk size (``cfg.batch.chunk_prefill``, so full chunks replay instead
of running eager), each capped at ``cfg.batch.max_batched_tokens``,
plus — when the per-step budget sits just above the chunk — one rung at
``max_batched_tokens`` for the mixed prefill+decode step (see below).
``ARBI_PIECEWISE_TOP`` overrides the chunk-derived top and suppresses
the mixed rung (an A/B pin means exactly what it says);
``ARBI_PIECEWISE_SPARSE_TOP`` keeps only the dense low ladder + one big
rung (sub-rung prefills run eager via the pad-waste floor rather than
inflating onto the big rung). Returns ``[]`` when nothing fits.

Bytes the piecewise per-rung gate must hold back for post-capture serving.

The post-capture KV grow
(:func:`arbi_serve.engine.inprocess_capture.grow_kv_after_capture`) sizes
the KV pool into ``free − serving_floor``, so every byte the sweep
captures beyond this reserve comes straight out of servable KV. Reserving
the serving floor here turns a KV starvation risk into graceful
degradation: the sweep stops adding rungs while KV can still grow
healthy, and the dropped rungs run eager via the dispatcher's pad-up /
passthrough (capture what fits, run the rest eager).

Three additive terms:

  * ``serving_floor_for_grow`` — the free-VRAM floor the post-capture grow
    holds out (decode-class activation peak + MTP verify tail + …). The
    grow leaves exactly this much free, so the sweep must leave it plus
    everything below for KV to land above the floor at all.
  * a healthy-KV page tranche — the same per-config floor
    :func:`assert_kv_pages_floor` is held to, in the SHORTFALL form the
    gate's currency requires (see below). This is the minimum the grow must
    be able to map, not a capacity target; the rung-vs-KV tradeoff beyond it
    belongs to the operator's max_batch / max_context / prefill-capture
    knobs.
  * the boot-measured single-shape (MBT-wide prefill) peak activation
    transient — the eager forward transient the next rung's capture (and
    any post-drop eager prefill) re-allocates from the default allocator.
    The rung-cost estimate models only graph bytes, so the capture-time
    margin does not otherwise cover a multi-GiB eager prefill workspace;
    mid-sweep ``empty_cache`` (a dropped high rung) can have returned
    those segments to the driver, so the gate's ``free`` reading would
    otherwise double-promise them.

TWO CORRECTIONS THIS FUNCTION EXISTS TO NOT REPEAT (both measured, 27B/TP1
exl3-4.0bpw + DFlash2, one 24 GiB RTX 4090, ``--max-context auto``,
``--prefill-capture piecewise``, gmu 1.0 — the canonical served shape):

1. THE BASIS. ``cfg.cache.max_context`` at sweep time is the PROVISIONAL
   SIZING bound, not the served context: on an engine-flexed boot
   (``--max-context auto`` or a recipe default) it is the model ceiling
   clamped by reserved VA, and the boot's own log says so
   (``max_context=auto: PROVISIONAL SIZING bound``) because
   :func:`finalize_served_max_context` decides the served window later,
   from the realized pool. Reserving one full window of THAT is not a
   conservative reserve, it is an unsatisfiable one: the window was
   4356 MiB against 3122 MiB free at the gate, so ``need`` exceeded
   ``free`` by 2758 MiB at EVERY rung — including ``num_tokens=1``, whose
   rung cost estimate is 0 — and the sweep captured 0 of 12 rungs on every
   boot while the flag reported ``piecewise``. Growing the card grows the
   provisional bound with it, so no card size fixes it.
   :func:`kv_pages_floor_for_flexed_context` is the floor that applies when
   the context is an OUTPUT; the full-window floor applies when, and only
   when, the operator PINNED one (``_context_may_narrow`` is the single bit
   every narrow site reads, so this gate cannot disagree with them).
2. THE CURRENCY. The floor is a POOL-SIZE floor (pages the pool must HOLD)
   and this gate spends it against FREE VRAM — which is only what the grow
   can still ADD. The pages the pool has already mapped are not in ``free``
   and the sweep cannot take them back, so charging them again is a double
   count. Small on a prefix-backed pool at the gate (11 of the 1025 floor
   pages, 47 MiB, on the boot above — the sweep runs before the grow, so
   most of the pool is not mapped yet) and large on a pool the deferred
   resize already sized. Reserve the SHORTFALL either way.

Deliberately conservative: an over-estimate only truncates the captured
ladder earlier (dropped rungs run eager — correctness-neutral, and prefill
capture is throughput-neutral on decode-dominant traffic), while an
under-estimate re-opens the silent degraded boot. Returns 0 when the KV
pool is not growable (a pre-sized static pool cannot be starved by the
sweep).

True when the model carries at least one split-attn / split-gdn
decoder block — the layers :func:`precapture_layer_graphs` captures
piecewise (``_pre_attn`` / ``_post_attn`` or ``_pre_gdn`` / ``_post_gdn``
into per-bucket cuda graphs). A model with no split layers takes the
whole-forward capture path only, so the piecewise-prefill warm has
nothing to pre-warm.

True when the engine's KV pool carries a recurrent (GDN / Mamba /
ShortConv) state view — i.e. the model has linear-attention layers whose
per-token decode kernel JIT-compiles per batch width.

The decode warm pass is a no-op on a pure-attention model (nothing
recurrent to compile-per-width).

Every ``(B, S)`` a decode-flavour forward of this boot can present.

The REACHABILITY model the capture ladder is sized against, in one place,
derived from the boot's own admission/step limits rather than a heuristic:

  * ``B ∈ [1, cfg.batch.max_batch]``. The scheduler seeds
    ``batch_left = max_batch`` and decrements it once per appended slate
    row, re-testing before every append, so a slate can never exceed it;
    the MTP strategy only partitions that slate, and
    :class:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers` raises
    on ``B > max_num_seqs = max_batch``. Every intermediate ``B`` is
    reachable because rows arrive and finish one at a time.
  * ``S ∈ {1} ∪ verify_s``. A plain decode row is structurally exactly one
    token wide (``resolve_row_shape`` raises otherwise). A verify step
    coerces the slate to one uniform ``step_k`` and emits
    ``S = step_k + 1``; ``step_k`` collapses to 0 (⇒ ``S = 1``) whenever
    the draft is unavailable, and admission REJECTS — never clamps — a
    per-request ``mtp_k`` outside the captured depth set, so ``verify_s``
    is exactly :func:`_verify_s_set`.

``B × S`` carries no independent cap: the verify expansion is not charged
to ``max_batched_tokens`` (a decode row costs one token in the slate
budget regardless of ``K``), so no ``(B, S)`` in this product is excluded
by the token budget.

Args:
    max_batch: the served ``cfg.batch.max_batch``.
    verify_s: the served verify widths (:func:`_verify_s_set`); empty when
        the boot captures no verify shapes.

Returns:
    ``[(B, S), ...]`` sorted by ``S`` then ``B``.

Schedulable ``(B, S)`` shapes with NO route to a captured graph.

A shape is covered when the captured rungs at its ``S`` contain ``B``
itself, or — with decode pad-up armed (``cfg.decode_pad_cudagraph``) — a
rung ``B' > B`` with ``B' <= max_batch``, which
:func:`~arbi_serve.runtime.captured_lookup.lookup_captured_graph_for_forward_padded`
replays after padding the slate with benign scratch rows. Anything else
runs the live eager forward at that shape for the life of the boot.

This is the invariant a ladder derivation must preserve: pruning a rung is
only safe while every shape :func:`schedulable_decode_shapes` enumerates
keeps a route. Returns the violations, sorted, so a caller can name them.

Args:
    captured_b_by_s: ``S -> captured B rungs`` for ONE
        ``(lora_bucket, kv_pages_bucket)`` class of the captured pool.
    max_batch: the served ``cfg.batch.max_batch``.
    verify_s: the served verify widths (:func:`_verify_s_set`).
    pad_up: whether decode-cudagraph pad-up is armed.

Synthetic-batch construction + the piecewise per-bucket capture driver.

The hoisted helpers that build a self-consistent synthetic capture batch
over the persistent piecewise buffers and drive the per-layer capture
sweep for one ``num_tokens`` bucket. Imports the paged-KV sliding-window
helper from ``_shapes``.

Drain pending recurrent zero-clears + savepoint resumes before a
capture sweep builds recurrent (GDN / Mamba / ShortConv) metadata.

Settles the pool at the sweep boundary, so the whole sweep — buffer
fills, metadata builds and the capture forwards they feed — runs
against rows whose queued writes have already landed, rather than
against a pool that settles partway through.

This is a convenience at the boundary, not the guarantee: the
invariant itself is enforced one level down, where it cannot be
skipped by a caller, in
:meth:`RecurrentMetadataBuilder.build`. Both drain through the same
primitive so the ordering (zero-clears first, savepoint resumes
second) has one definition.

Populate the persistent piecewise buffers for a synthetic capture batch.

``rows`` is the per-row token count (``B = len(rows)``, ``N = sum(rows)``).
Writes a self-consistent batch into the persistent ``pb`` slices: zeroed
ids/slot_mapping, per-row-restart positions, cu_seqlens from the row cumsum,
``seq_lens = rows``, identity recurrent state indices (int32 + the int64
alias the recurrent block forwards consume) and cleared
``has_initial_state``. The live ``_build_batch`` overwrites these every step;
the values only need to be capture-shape-consistent.

Every ``cu_seqlens_q_cpu`` state a LIVE batch at ``B`` rows can carry.

The compiled GDN layer branches on this field's None-ness inside the
``fullgraph=True`` region (``_gdn_fla_forward._resolve_cu_seqlens``), so
it is a Dynamo GUARD, not a value: a boot that warms one state while
serving produces the other leaves the first live request to compile the
layer. The axis is BOUNDED (two states), so warmup enumerates it rather
than betting on the boot's resolution — the admin console can flip
``prefill_capture`` at runtime, which flips this field under traffic.

The boot-resolved state is returned FIRST, so a caller that warms only
the head of the tuple still covers today's traffic.

Build a ``ScheduledBatch`` over the persistent ``pb`` slices.

Reads the buffers :func:`_fill_synthetic_prefill_buffers` populated for the
same ``rows`` (``B = len(rows)``, ``N = sum(rows)``). Used by every synthetic
capture / compile-warmup forward so they share one batch-shape contract.

``cu_seqlens_q_cpu`` is NOT hand-set here: pass ``eng`` and the field is
resolved by :func:`~arbi_serve.runtime._batch_build_helpers.
host_cu_seqlens_q_twin` — the same call
:func:`~arbi_serve.runtime.batch_build.build_batch` makes on the serving
path — so the two constructions cannot disagree on a field the compiled
GDN layer guards on. Callers enumerating BOTH None-ness states (see
:func:`synthetic_host_twin_states`) pass the state explicitly instead.
Omitting both keeps the historical ``None`` (the capture sweeps, which
must not bake host lengths into a recorded graph).

Run each per-kind metadata builder and bind the meta onto ``batch``.

PAGED_KV / MLA metas bind directly; recurrent (MAMBA / GDN / SHORT_CONV)
metas get ``has_initial_state`` rebound to the persistent
``rec_has_initial_state`` slice so capture + replay flow through a stable
``data_ptr`` (the default builder derives ``seq_lens > 1`` into a fresh
tensor, which is capture-incompatible).

DRAINS FIRST, making this the boot-side twin of
``forward_exec.build_metadata``'s seam. ``GdnMetadataBuilder._finalize``
refuses a build with recurrent zero-clears pending, and until this line
the contract was honoured by CALL ORDERING: the capture families and
``precapture_compile_warmup`` each drain at their own entry, and the
dozen ``_run_metadata_builders`` sites below them inherited it. That held
for eleven of them and not for the twelfth —
``postresize_compile_warmup_embed_override`` is invoked straight from
``build.py`` rather than through ``precapture_compile_warmup``, so the
build it reaches via ``_compile_warmup_embed_override`` had no drain
upstream of it at all. Ordering across four modules is not a contract
anybody can check by reading one of them; draining HERE makes it one that
holds by construction for every present and future caller. Idempotent and
a no-op on a pool with nothing queued.

Capture every eligible layer at one ``num_tokens`` bucket.

Builds a synthetic single-request (``B=1``) prefill batch of
length ``num_tokens`` whose per-step batch tensors reference
prefix slices of :attr:`Engine.piecewise_buffers`. Runs each per-
kind metadata builder once so the ``*_meta`` dataclasses also
reference persistent slices, then drives ``model.forward`` with
``_piecewise_capture_on_miss=True``. Each eligible layer is hit
by the dispatcher on its first forward pass at this
``num_tokens`` and captured by
:func:`arbi_serve.runtime.capture.dispatch.capture_layer`.

Persistent-buffer contract. Capture binds the per-layer graphs
against the ``data_ptr``s of the
:class:`arbi_serve.runtime.capture.dispatch.PiecewiseBuffers` slices —
not against fresh per-call tensors. The live hot path's
:func:`arbi_serve.runtime.model_runner._build_batch` writes per-
step content into the same persistent storage before
``model.forward`` runs, so replays read fresh data. Without
this redirect the captured kernels would forever read the
synthetic-zeros tensors and KV-cache scatter / paged-attention
gather would walk past their valid bounds.

Returns the count of layers captured this call. Failures are
raised — the caller's ``precapture_layer_graphs`` wraps this in
a try/except and skips the bucket.

Every rank serves the same boot-time attention picks.

The verify split table (``tkv.runtime.attend.verify_split_sweep``) is a TIMING
sweep. Each rank runs its own, alone, against its own clock, and publishes the
winner into its own process. Nothing compares them, and two things make that
unsafe rather than merely untidy:

* the split count is a kernel constexpr AND a grid dimension, frozen into the
  MTP-verify CUDA graph each rank captures, so ranks that picked differently
  replay different kernels on every decode step for the life of the process;
* the prefill prewarm enumerates its verify variants FROM the published table
  (``mtp_verify_num_splits(..., shape=shape)``), so a divergent pick also
  gives the ranks different COMPILE sets — a per-rank-variable boot latency of
  exactly the kind ``compile_warmup_piecewise_prefill`` exists to keep out of
  a rank-lockstep collective sequence.

The records on disk cannot catch either of them. ``VerifyShape`` carries
``device_index``, so two ranks fingerprint to two different files, each writes
its own, and no lookup ever compares them — the divergence is invisible by
construction rather than merely unchecked.

turbo-attn #1054 removed the int8 hybrid crossover, the second sweep this
reconciled. Its import shared the ``try`` below with the verify-split
enumerators, so leaving it in place would have made a missing module take the
LIVE half down with it and degrade the whole consensus to a warning at TP>1 —
a check that could not run reporting success.

So consensus is imposed HERE, where the process group lives: turbo-attn has no
business knowing about ranks, and a sweep that reduced over one would be a
distributed dependency in a kernel library. Rank 0's picks win and every rank
re-publishes them under its OWN shape key, which differs from rank 0's only in
the device index.

Rank 0 is chosen over a reduction because the quantity is a pick, not a
measurement to average: any rank's answer is a legitimate winner of the same
contest, and what matters is that one of them is served by all. A reduction
would also have to agree on a rule for cells only some ranks measured.

Every rank still runs its own sweep. That is deliberate: this adds a
RENDEZVOUS, never a skip, so both ranks always reach the broadcast and none
can wait on a peer that took a different path. The cost is the sweep time
already being paid, and the boot keeps the symmetric shape it had.

``shape`` with its device index dropped.

The join key between ranks. Two ranks of one TP group hold the same
attention geometry on different cards, so their shapes agree on every
field but this one; dropping it is what lets rank 1 find the pick rank 0
published for the same launch.

``(local_shape, pick)`` for every local shape rank 0 also measured.

Pure, so the reconciliation can be tested without a process group or a
GPU. A local shape rank 0 did NOT publish is left alone rather than
cleared: it is a geometry only this rank drove, which pure tensor
parallelism does not produce, and dropping a pick it is serving would be a
worse answer than keeping an unreconciled one. The caller reports those.

Make every rank serve rank 0's verify-split picks.

Best effort and SYMMETRIC in its failure: every rank runs the same code
over the same geometries, so a build without the turbo-attn enumerators,
or a group that is not initialised, returns on all ranks together and
leaves the boot exactly as it was. It must never be the thing that hangs a
boot — but note that returning early on ONE rank only would, which is why
every guard below reads a value that is identical on all of them.

Wall-clock + compile-count ledger for the compile-warmup phase.

The phase reports ONE duration for a dozen distinct pieces of work
(per-shape compiles, kernel prewarms, cache persistence), so a slow boot
says nothing about WHICH piece is slow. :meth:`mark` closes the segment
that began at the previous mark and attributes it the compiles the JIT
detector recorded meanwhile; :meth:`log` emits the itemized line
unconditionally, so a warm boot's residual cost is attributed rather than
implied.

Telemetry only: every method swallows its own failures.

Warm the DFlash draft forward at every shape the capture sweep records.

Duck-typed on the drafter: a boot with no DFlash drafter, or one whose
draft-graph capture is not armed, warms nothing and returns.

Two things follow from running the drafter's warmup HERE rather than
letting the Phase-5 sweep drive it:

  * the compile happens before ``save_compile_cache_artifacts`` writes
    this boot's blob, so the drafter's Inductor artifacts are IN it and a
    later boot on the same key reloads them;
  * the phase ledger can charge it, which it cannot do for work that
    happens after the ledger has closed.

Best-effort. A failure here leaves exactly the previous behaviour — the
sweep warms and compiles each shape itself — so it is logged, not raised.

Cumulative seconds Dynamo attributes to each compile phase, or ``{}``.

``{}`` when torch does not expose the accounting, which reads as "no
split available" rather than as a zero-cost compile — the two are not the
same and only one of them is good news.

Run the forwards that must happen outside any cudagraph capture window.

Named for compilation, but on a warm boot this is EXECUTION: the Inductor /
Triton caches remove the compile, not the forward. What remains is the
synthetic forwards themselves, Dynamo's guard checks, and loading and
linking the compiled modules (the ``cuModule`` loads that surface as
``driver.modules_loaded``). A cache hit shrinks this phase; it cannot
remove it.

When the model is
decorated with :func:`support_torch_compile` and the engine is
booted with both compile + cudagraphs enabled, the very first
forward inside ``torch.cuda.graph(...)`` would trigger Dynamo +
Inductor lazy compilation (via the trampoline's ``_ensure_compiled``).
Inductor's tracing performs unpinned CPU→CUDA tensor copies, which
CUDA forbids during stream capture — the capture aborts and the
allocator's ``captures_underway`` counter ends up in an invalid
state, taking the boot down with SIGABRT.

Run a synthetic forward through ``compile_capture_ctx`` before the
capture sweep so Inductor compiles in a regular allocator window.
The trampoline caches the compiled callable; subsequent
capture-region forwards just call the cache, no compile inside
cudagraph.

Per-capture-function warmups (in ``capture_decode`` /
``capture_prefill`` / ``capture_drafter_chain``) also wrap their
own warmup forwards in ``compile_capture_ctx``, which provides
per-shape compile coverage inline. This top-level function is
defense-in-depth: it forces an early compile for the canonical
``(B=1, S=1)`` decode shape, so even capture call sites that
skip the per-function warmup (or that hit a recurrent code path
where the per-function warmup runs the eager fallback) find
``state["compiled"] is True`` on entry.

No-op when:

  - neither the model nor its ``.backbone`` is compile-decorated
    (:func:`~arbi_serve.compile.capture_bridge.compile_eligible_root`
    returns ``None`` — see its docstring for why a composite
    multi-modal wrapper is resolved through ``.backbone``), or
  - the resolved compile root has no ``force_compile`` method (older
    trampoline variant — the per-capture warmup still covers it), or
  - ``eng.piecewise_buffers`` is ``None`` (the boot orchestrator
    builds it — transiently under ``cuda_graphs=False`` — before
    calling this function; a caller that skips that step leaves
    nothing here to warm against).

Idempotent: safe to call multiple times (the trampoline state's
``compiled`` flag short-circuits subsequent invocations).

Called from the boot orchestrator (post-``build_active``,
pre-``_precapture_*``).

Deterministic attention-window tuple, global (``None``) first.

The window flips ``is_local`` in the CuTeDSL prefill kernel-cache key,
so a window that is not listed here stays COLD. Order is fixed only so
the boot log is diffable across runs.

``(field, tag, counter)`` for every int8 prefill route this boot armed.

BOOT window: the prefill prewarm runs once while the engine builds, and
each route is selected inside tkv's dispatch, which carries no counter on
this side. A fire therefore says "boot compiled the route's variants",
never "the served prefill ran them".

An empty list with every flag off is the inert state, and it is
deliberate: the bf16 route is what such a boot selects, so there is no
claim to witness — and a counter that fired (or refused) on every boot
would say nothing about the boots that asked for a route.

Record, per armed int8 route, whether the prewarm drove it.

Each route is an axis of the prefill kernel-cache key, so a boot that
selects it and warms only the other variants leaves the SERVED kernel
cold: the first live prefill compiles it and ``jit_compile_serving``
names that compile. The refusal is what turns that into a boot-time
signal instead of a serving-time surprise.

Whether the installed turbo-attn's prewarm takes the verify-band args.

The pinned engine image and the mounted turbo-attn checkout can
disagree, and ``_run`` reports a TypeError as "the whole prewarm
failed" — which would trade a cold verify band for a cold
EVERYTHING. Probe the signature instead, and say so when the build
is too old.

The boot-time sweep that measures the Turbo prefill verify split per
batch for this deployment, or ``None`` with the reason logged.

The split count is a kernel constexpr and a grid dimension frozen at
capture, so turbo-attn measures it here — eagerly, per attention
geometry, against this deployment's batch range and served context —
and the prewarm that follows warms the count the table will serve.
``mtp_kwargs`` is the verify band the prewarm was handed: no band, no
verify launch, nothing to measure. A build without the sweep serves
the closed form and says so once; it is not a boot failure, because
the closed form is what every boot served before the table existed.

Pre-compile every CuTeDSL prefill kernel this deployment can reach.

WHY THIS EXISTS. The tkv Turbo prefill dispatch caches one compiled
CuTeDSL kernel per shape class and calls ``cute.compile`` on a miss —
seconds of MLIR trace + lowering, even with the on-disk artifact cache
warm (that cache elides the cubin build, NOT the trace). A miss taken
after the engine declares itself serving-ready blocks the whole
scheduler step, so at concurrency C it is paid by C in-flight requests
and surfaces as an unexplained TTFT spike.

NO OTHER BOOT WARMUP REACHES THESE KERNELS, and the reason is
structural rather than a missing shape: every synthetic prefill at boot
runs against an EMPTY KV pool, so ``TKVCore.should_bypass`` / the
first-chunk bypass takes the raw-bf16 route on all of them. The tkv
CODEC prefill kernels — the unified ``TkvLoader`` and the
fresh-diagonal ``HybridLoader``, on both the ``B == 1`` padded-gather
route and the varlen paged route — are therefore never driven at boot,
no matter how many synthetic prefills ``_compile_warmup_prefill`` adds.

MEASURED (post-ready compile traces from server logs, RTX 4090). The
three bf16 arbi cells logged ZERO post-ready ``cute`` compiles — the bypass
kernels are already warm from the synthetic forwards. The three tkv
cells logged exactly one ``TurboPrefillForwardHybrid`` each, plus one
``TurboPrefillForward`` at c16 and c64. Two kernels, on the codec path
only: a bounded warmup-COVERAGE gap, not a "prefill shape classes are
unbounded" problem. That
story is disproven — the dispatcher's kernel-cache key carries neither
``S_q`` nor ``S_kv``, and its batch element is the literal
``0 if is_varlen else B``; sequence length reaches the key through
exactly one boolean, the asym-long tile choice
(``S_kv >= 4*S_q and S_kv >= 16384``).

So the closure is enumerable, and ``tkv.runtime.attend.prewarm``
enumerates it: route (padded-B1 / varlen) x loader (unified /
fresh-diagonal) x asym-long (off / on) per attention window. Each
variant is DRIVEN through the production entry point rather than
re-derived, so this cannot drift out of sync with the dispatcher: if
the key grows an axis, these calls produce the new key too — including
the int8 route ``TKV_PREFILL_INT8`` selects, which tkv reads from the
environment on the same dispatch this drives.
:func:`_note_int8_route_warmup` then checks the report for each armed
int8 route by name, because "the prewarm ran" is not evidence that the
variant this boot will serve is among the ones it compiled.

GEOMETRY COMES OFF THE CONSTRUCTED ``eng.attn_ops``, never off
``model.layer_specs``. The ops carry the PER-RANK (TP-sharded) head
counts and each layer's own smart-mix ``(k_bits, v_bits)`` — which is
what the kernel cache is keyed on. ``layer_specs`` carries GLOBAL head
counts, which at TP>1 name a kernel the runtime never requests (the
same trap ``warm_tkv_kernel_modules_for_engine`` documents). Distinct
``(head_dim, n_q, n_kv, k_bits, v_bits)`` groups are warmed separately
— Gemma-4 mixes head_dim-256 sliding layers with head_dim-512 full
layers, and smart-mix gives different layers different widths — and
each group is warmed over the distinct WINDOWS its own layers use.

``include_codec`` / ``include_bypass`` follow what the deployment
actually serves: codec routes only when a :class:`TkvAttnOp` is wired,
the raw-bf16 ``BypassLoader`` route only when a
:class:`TkvBypassAttnOp` is. A pure-bf16 boot does not compile codec
kernels it can never dispatch, and vice versa.

TP-safe, but not by geometry alone. The cute prefill kernels issue no
collectives, so nothing here can desync mid-call. The ENUMERATION, though,
is no longer a pure function of the per-rank geometry: the MTP-verify
variants are one per distinct split count, and the split count comes from
the boot-time table (``mtp_verify_num_splits(..., shape=shape)``), which is
a per-rank TIMING sweep. Two ranks that timed differently would compile
different variant sets and freeze different constexprs into their verify
graphs. What makes the enumeration identical again is
:func:`~arbi_serve.engine.capture_admin.boot_pick_consensus.agree_boot_picks`,
which the tuner runs before this prewarm reads the table back; the geometry
half — head counts, widths, windows — is per-rank and identical as before.

Best-effort. A turbo-attn without the prewarm entry point, or a single
variant that raises, must never break a boot that would otherwise
serve. A skipped variant is logged at WARNING, not debug: it is exactly
a cold kernel some live request will pay for, and ``jit_detector`` will
then name it under the hard ``jit_compile_serving`` gate.

Mark the recurrent metadata's num-seqs dim dynamic on ``batch``.

A recurrent (GDN / Mamba / ShortConv) block reads its row count off
``meta.state_indices`` and branches on ``n_tokens == n_seqs``, so a
statically-shaped ``state_indices`` specializes the compiled layer per
CONCURRENT-ROW COUNT — leaving every slate width the warm did not
itself run to re-trace on a live request. The metas are built by
:func:`_run_metadata_builders` from storage the batch tensors' own
marks do not reach, so they carry the mark separately.

``maybe_mark_dynamic`` is advisory (a no-op under nesting) and safe —
the dim genuinely varies at runtime.

Pre-compile the multi-row prefill block graphs at boot.

Runs synthetic prefill forwards whose flat-token dim is > 1 and
marked dynamic, so Dynamo promotes the token dim to a symbolic shape
in one shot (its automatic-dynamic otherwise needs to see two
distinct non-1 values before generalizing). After this the live
multi-row prefill path replays a warm compiled graph for any token
count, so a concurrent burst's first-token latency is the forward
cost, not an Inductor re-trace.

Best-effort + idempotent. Skips silently if the model is not the
per-layer (block) compile variant or the synthetic forward raises.

Pre-compile the single-row (B=1) prefill block graphs when eager.

``precapture_prefill_graphs`` normally compiles AND captures one
whole-forward CUDA graph per configured ``num_tokens`` bucket at
``(B=1, S=num_tokens)`` — every live single-row prefill chunk count
is covered because it either matches a bucket or pads up to one. That
sweep refuses outright when ``cuda_graphs=False`` (its first check).
``_compile_warmup_prefill`` above deliberately excludes the
single-row case too, relying on that (here, absent) capture sweep —
so on a ``cuda_graphs=False`` boot nothing compiles the single-row
prefill shape ahead of time. Every distinct token count a live
single-row prefill reaches for the first time — a duplex tick's
per-frame STT chunk, typically a handful of tokens — is then a live,
uncovered Dynamo recompile fired directly on that request.

Mirrors ``_compile_warmup_prefill``'s dynamic-shape technique:
Dynamo's automatic-dynamic promotes the flat-token dim to a symbolic
size once it observes two distinct non-1 concrete values under the
same guard, so a HANDFUL of representative token counts is enough —
a wide dense sweep is actively counterproductive here, because a
per-layer attention block ALSO guards ``attn_ops[self.layer_idx]``
(a plain Python list index, not a tensor dim — every distinct
attention-layer index needs its own compile regardless of shape
generalization). Each extra warmed shape multiplies the attention
code object's compiled-variant count by the number of attention
layers; a too-wide sweep burns through Dynamo's per-code-object
``recompile_limit`` before every layer index is even covered once
(measured: 66 shapes x 4 attention layers = 264, over the 256
ceiling ``arbi_serve/compile/dynamo_config.py`` raises it to — boot
warmup silently under-covers and live traffic pays for the tail).
Picks span the small end a duplex tick's per-frame STT chunk
reaches (2, 8) and the turn-based chunked-prefill range (128, 1024,
cap).

Runs whenever the whole-forward prefill capture sweep will NOT cover
the single-row shape. That is TWO config states, not one:

  * ``cuda_graphs=False`` — the sweep refuses outright (its first check);
  * ``prefill_capture='eager'`` — ``build.build`` skips both the
    whole-forward and the piecewise sweeps and records the dormancy on
    ``prefill_capture_sweep`` / ``piecewise_sweep``.

Gating only on ``cuda_graphs`` covered the first and MISSED the second,
which is the shipped production combination (cudagraphs on for decode +
verify, prefill eager). On that boot NOTHING compiled a ``B=1`` prefill
ahead of time: ``_compile_warmup_prefill`` warms ``B ∈ {2, 3}`` and the
compiled layer guards ``2 <= state_indices.size()[0]``, so a live
single-request prompt — the overwhelmingly common shape — fell outside
every cached entry and JIT-compiled the GDN layer inside the request
(``jit_compile_serving``, once per boot, on the first live request).

Pre-compile the external-embedding decode and prefill variants.

Models that accept ``pending_embed_override`` take a distinct Dynamo
specialization because their forward receives ``inputs_embeds`` instead
of performing its token-embedding lookup.  Duplex speech uses that branch
for both its fused seed prefill and every frame-lockstep decode step.  A
newly admitted request also has a one-column block table, which Dynamo
specializes independently from wider tables.  Both phases and table
widths are compiled before readiness.

Pre-warm the piecewise prefill kernel surface at boot, TP-lockstep-safe.

The enabler for TP>1 piecewise prefill capture. The
``precapture_layer_graphs`` capture sweep drives each eligible layer
through a rank-lockstep sequence of NCCL collectives (the GDN / MLP
``RowParallelLinear.all_reduce`` inside ``capture_layer_split_gdn``'s own
warmup + capture passes). If a per-layer FLA ``chunk_gated_delta_rule`` /
Triton kernel fires its first-call autotune (or an Inductor
``benchmark_combo_kernel`` when ``inductor_combo_kernels`` is on) inside
that sweep, the autotune's per-rank-variable latency stalls one rank out
of a collective while its peer spins inside the next one, so both GPUs
busy-wait forever.

This closes it: for every piecewise capture bucket (the same ladder
:func:`precapture_layer_graphs` builds, via the shared
:func:`_piecewise_capture_bucket_ladder`), drive one eager synthetic
``B=1`` prefill forward under ``enable_piecewise_cuda_graph()`` +
``enable_piecewise_capture_sweep()`` — the identical context + kernel
surface the capture sweep's internal warmup pass hits — on every rank.
The forward runs the GDN custom-op prefill path (``arbi_serve::gdn_-
attention``) and the row-parallel ``all_reduce`` at each rung's flat-
token width, so the FLA / Triton autotune caches and the per-width
lm_head all-gather ``_staging`` buffers are all populated here, outside
any capture window. When the capture sweep later runs, every per-layer
kernel is a cache hit on both ranks, so the collectives rendezvous in
lockstep and no rank stalls in autotune.

Runs with ``_piecewise_pool`` still ``None`` (the capture sweep has not
populated it yet), so the layer dispatcher takes its eager
``layer(hidden, *block_args)`` fast path — nothing is captured or
recorded here; this is purely a compile/autotune warm.

TP1: skipped entirely (no collectives, so no desync is possible; the
capture sweep's own per-bucket internal warmup already covers the
single-rank case). This keeps TP1 boot byte-identical.

Best-effort per bucket, but a warm failure is deterministic across ranks
(same synthetic shape, same code) so it fails symmetrically — it does not
reintroduce a desync. No-op when the model has no split-piecewise layers.

Pre-warm the live per-width decode path at boot, TP-lockstep-safe.

Hybrid models (GDN / Mamba / ShortConv) compile their per-token
recurrent decode Triton kernel on the first invocation at a given
batch width, and the tensor-parallel lm_head all-gather allocates
its per-width ``_staging`` buffer on first use. If either happens on
the first live decode of a width, the request pays a one-off JIT
stall (TTFT spike) — or, at tight ``gpu_memory_utilization`` where
the KV pool has grown to fill VRAM, the staging ``torch.empty``
OOMs mid-serve.

This drives the same ``run_model_forward`` live path
:func:`_compile_warmup_prefill` uses, once per scheduled decode width in
:func:`_decode_b_ladder`. Every rank builds the byte-identical synthetic
batch, so the row-parallel collectives rendezvous in lockstep on every
rank (SPMD parity) — no rank stalls out of a collective on an uncached
compile, and there is no cudagraph capture region open, so a mid-forward
Triton/Inductor compile cannot abort a capture. That makes this
deadlock-safe at TP>1, unlike a direct ``eng.model.forward`` call that
bypasses the SPMD-parity machinery.

Placed in ``precapture_compile_warmup`` (before the capture sweep and the
post-capture KV grow): the recurrent kernels + Inductor decode graph +
the per-width lm_head ``_staging`` buffers are all warmed/allocated while
VRAM is still free, and the persistent staging buffers survive the grow.

Best-effort + idempotent. No-op when the model has no recurrent pool
(nothing compiles per width) or the synthetic forward raises.

Pre-compile the MTP verify-pass hidden-gather forward.

The MTP verify pass calls ``EagerModelRunner.forward(batch,
return_hidden_state=True, compute_logits=False)`` to recover the per-token
hidden state. On a cold / K=0-collapse verify step the batch is at
``max_query_len == 1`` (plain-decode width), so the captured-graph lookup
hits the S=1 plain-decode graph — which is captured without ``hidden_out``
(only S>1 verify captures retain per-token hidden) — and the runner falls
through to the live compiled ``_run_model_forward`` (model_runner.forward).

Every other compiled-forward warm at boot (:func:`precapture_compile_warmup`
B=1/S=1, :func:`_compile_warmup_prefill`, :func:`_compile_warmup_decode`)
drives the forward without ``return_hidden_state``, so the hidden-gather
Dynamo specialization is never compiled at boot — it therefore compiles on
the first live MTP request instead.

This closes the gap by driving the exact fall-through path —
``run_model_forward(runner, batch, return_hidden_state=True,
compute_logits=False)`` — over the plain-decode width ladder, so the frame
compiles pre-ready and the first live request replays a warm graph. The
identical synthetic batch on every rank keeps the row-parallel collectives
in SPMD lockstep (deadlock-safe at TP>1), and it runs before the capture
sweep so no captured decode graph exists to race and the compile lands
outside any capture window. The num-seqs dim is marked dynamic so Dynamo
generalizes across widths (covering c1 and the batch>1 shapes) rather than
re-specializing per concrete B.

Best-effort + idempotent. No-op when MTP is disabled (a non-spec engine
never takes the verify hidden-gather path) or the synthetic forward raises.

Boot-refuse ``ARBI_SPMD_VERIFY_OVERLAP=1`` on a cudagraph boot.

Every captured graph (decode / MTP-verify / drafter / prefill /
piecewise) records into the one shared ``capture.cudagraphs`` mem pool
and overlays per-capture transient blocks — safe only while replays are
strictly serialized. The overlap path enqueues the verify forward and
the drafter chain back-to-back without the intervening host join, which
permits two captured graphs to be in flight concurrently; concurrent
replays of overlaid graphs read/write the same pool addresses (silent
corruption, not a crash). Refuse at boot, loudly, rather than serve a
config whose failure mode is wrong tokens.

The one ``(max B, vocab)`` logits buffer the whole decode ladder shares.

Returns ``None`` (every shape then allocates its own, as before) when
there is nothing to share — no GPU, no ``capture.io_buffers`` pool, or
an empty sweep.

Sizing is the sweep's OWN ``max(B)``, not ``cfg.batch.max_batch``: the
ladder need not reach the configured ceiling, and a buffer sized past
what any captured shape can use is exactly the kind of invisible
over-reserve that comes out of KV. The saving is ``sum(B) - max(B)``
full-vocab rows: on a Qwen3.5-0.8B TP1 bf16 mb64 boot the 15-shape

Idempotent and NEVER shrinking-or-moving in place: a re-sweep (backend
hold this buffer's ``data_ptr`` baked into their kernel args, so an
existing buffer that still fits is reused rather than reallocated. A
per-model attribute (``ENGINE_MODEL_ATTRS``) so a pool member's buffer
is snapshotted and repointed with the rest of its state.

Capture decode graphs for every shape in ``cfg.cudagraph_shapes``.

When MTP is enabled, the sweep is auto-extended via
:func:`_decode_capture_shapes` to cover every ``(B, ``(B, K + 1)`` rows of ``cfg.cudagraph_shapes`` to get the verify
forward replayed. Verify-shape captures route their persistent
inputs through :class:`VerifyBuffers`, so replay skips the
inner ``copy_()``s — :meth:`VerifyBuffers.load_step` already
populated the same storage on the live path.

Failures on any single shape log + skip that shape only — the rest
of the sweep proceeds, and ``_step()`` falls through to the live
path for missing shapes. Under TP>1 a per-rank skip is not safe
(the peer rank is already inside the shape's capture collective),
so failures raise instead — see
:func:`_raise_on_tp_capture_failure`.

Sleep-mode interaction: in the stable-VA mode, every captured
graph's persistent buffers are registered with
:class:`SleepableTensorPool` after capture so their data_ptrs
survive release / resume. Verify-shape graphs share buffers with
:class:`VerifyBuffers` (already registered separately by
:func:`arbi_serve.engine.sleep.build_sleep_pool`) — registering
them again here is harmless (the pool is keyed on data_ptr and
treats duplicates idempotently) and unnecessary. The drop-and-
realloc mode skips this registration; the captured-graph pool is
cleared on release and re-captured on resume.

Boot-fail (by default) when any attempted decode-capture bucket failed.

``eng._capture_failures`` holds ``(shape_desc, exception)`` for every
decode/verify bucket the config implied that was attempted and raised
during the sweep (recorded by :func:`_note_capture_failure`), plus a
failed ``precapture_compile_warmup`` synthetic forward. By-design
uncaptured shapes (verify pre-flight refusals, S over the served
block_m ceiling, whole-model ``_can_capture_decode`` refusal) are never
in the ledger — this gate fails only on genuine failures.

Why hard-fail: a failed bucket silently serves eager decode at that
shape — correct output, degraded TPOT, invisible to the client.
cuda_graphs is the serving config; a config the engine cannot deliver
is a boot error, mirroring the ``backend_version_guard`` /
``_assert_required_verify_captured`` precedent.

``ARBI_ALLOW_DEGRADED_CAPTURE=1`` is the explicit escape hatch for a
deliberate degraded boot (tight card, bring-up): it downgrades the
raise to a warning and registers the degradation on
``eng._serving_degraded`` so the ``serving_degraded`` gauge / the
optional ``x-arbi-degraded`` header keep it visible.

Return the PAGED_KV metadata builder the drafter chain must decorate
its per-step metas with, or ``None`` when the backend needs none.

The MTP layer is a full-attention layer (PAGED_KV state kind); its
metadata builder is shared across every attention layer. The drafter
builds bare per-step metas that skip ``build``/``_finalize``, so the
per-step page buffers the decode kernel reads default to ``None`` and
the loader crashes on ``None.contiguous()``. Any builder exposing
``decorate_external_meta`` (TKV's ``_tq_*`` mirrors and tkv-bypass's
``_bt_*`` CSR triplet) must decorate them; tkv-bypass / GDN
builders don't expose it and need no decoration → ``None``.

Prefers the drafter-private TKV builder so the captured drafter
recomputes its page metadata into a ``TQBufferPool`` distinct from the
main-model verify builder (never the verify graph's baked scratch);
falls back to the shared PAGED_KV builder otherwise. Mirrors the
live-path resolution in :func:`decorate_draft_metas`.

Capture the external draft-model chain at every ``(B, K)`` bucket.

The external drafter (:class:`arbi_serve.spec_decode.external_drafter.
ExternalModelDrafter`) is a separate small model running its own
autoregressive K-step decode chain against its own paged-KV pool. This
sweep mirrors :func:`precapture_drafter_chain` but captures the
drafter's own ``model.forward`` chain (via
:func:`arbi_serve.runtime.capture.drafter.capture_external_drafter_chain`)
keyed on the same ``(B, K)`` ladder (:func:`_drafter_chain_buckets`) the
scheduler issues. Per-bucket failures log + skip; the drafter then runs
the eager chain for that shape (TP=1 — slow, never a deadlock).

No-op when the drafter has no captured-chain pool attached
(``cfg.cuda_graphs`` off — :func:`build_external_drafter` only attaches
one under cuda_graphs).

Capture the MTP drafter chain at every ``(B, K)`` bucket.

Each bucket runs one capture against the per-arch bundled
:class:`MtpHead` (today :class:`Qwen3_5MtpHead`). Failures on a
single bucket log + skip that bucket — the rest of the sweep
proceeds, and :meth:`MtpDriver.draft` falls through to the live
chain on misses.

Sleep interaction: every captured chain's persistent buffers are
registered with :class:`SleepableTensorPool` after capture so their
data_ptrs survive release / resume (mirroring
:func:`precapture_decode_graphs`).

Rung -> why it was not captured, for the rows THIS sweep appended.

The rows
that drop a capture bucket already carry ``truncates_capture_pool`` for the
graph-pool persist guard; this asks a different question of the same rows —
which rung, and was it the one the operator named — so no second flag is
needed and there is no parallel record to fall out of step.

``ledger_mark`` scopes the read to this sweep. An engine with no
``boot_state`` records nothing and therefore loses nothing here, the same
way :func:`note_substitution` no-ops on one.

The ladder rung that IS the operator's ``--chunk-prefill``, if any.

Returns ``None`` — the answer for every boot that does not name a width —
when ``chunk_prefill`` is the engine's own choice, and when the operator's
width is not a rung of this ladder at all. The second case is not a
substitution: ``max_batched_tokens`` or ``ARBI_PIECEWISE_TOP`` caps the
ladder below the chunk by CONFIGURATION, so the sweep was never going to
capture that width and nothing was swapped for anything.

Read through :func:`getattr` so an engine double built before the
provenance field existed answers "engine-chosen" rather than raising.

The refusal for a rung the operator named and the boot could not capture.

Names the three things the log line it replaces did not: the width that
was asked for, what runs at that width instead, and the knobs that move
the outcome. The last item is the one that makes this a refusal rather
than a louder failure — an operator who is told only that a capture failed
has no next action.

Boot-time piecewise per-layer capture sweep.

For each ``num_tokens`` bucket in
:data:`arbi_serve.runtime.capture.dispatch.DEFAULT_PIECEWISE_BUCKETS`
(capped at ``cfg.batch.max_batched_tokens``), build a synthetic
prefill batch, run ``model.forward`` with the dispatcher in
capture-on-miss mode, and register each captured per-layer graph
in :attr:`Engine.layer_captured_graphs`.

Per-bucket failures log + skip; the rest of the sweep proceeds.
Layers whose class declares ``_cudagraph_eligible = False`` are
skipped at the dispatcher and stay eager forever.

Raises :class:`~arbi_serve.engine.boot_degradation.
RequestedConfigurationNotInstalled` when the rung that was skipped is the
one ``--chunk-prefill`` explicitly named. Skipping a rung the ENGINE chose
is the extension working as designed; skipping the one the operator chose
installs a prefill path nobody asked for behind an unchanged ready banner.

After this returns, ``eng.model._piecewise_pool`` points at the
populated pool — the steady-state hot path's
:func:`piecewise.model_dispatch` calls will hit replay for any
bucketed ``num_tokens`` and fall through to eager for non-bucketed
ones (no lazy mid-step capture).

True when ``model`` routes any EXL3 weight through grouped MoE experts.

The dense EXL3 linears are fine under the fused mixed forward — the
row-split declaration keeps their decode block at a fixed M. The
grouped experts are not, because

The EXL3 backend is an optional extra (``arbi-serve[exl3]``), so a
bf16 / AWQ / FP8 deployment may not have it installed; an ImportError
answers the question — nothing is bound to a leg that does not exist.

The ``(prefill_tokens, context_len)`` graphs the sweep will capture.

Pure, so the coverage claim below is testable without a GPU.

Two capture FAMILIES per bucket, keyed by ``prefill_context``:

  ``context_len == 0``
    The prefill row is a FIRST chunk. TKV's per-row route gate classes
    it ROW_BYPASS and the graph bakes the bypass attend over the
    chunk's own contiguous K/V.
  ``context_len > 0``
    The prefill row is a CONTINUATION chunk. The gate classes it
    paged, whose per-row K extent is a device read of ``seq_lens``, so
    ONE graph serves every context length the replay refreshes it to.

``ARBI_MIXED_CAPTURE`` itself, with no second
opt-in, because they are not "coverage" and "extra coverage" here —
they are the whole feature and nothing. A prefill row is co-admitted
with decode rows on the chunks AFTER the first far more often than on
the first: the first chunk of a fresh request enters a slate that
holds no decode row for it yet, while chunks 2..n ride alongside
whatever is already decoding. Measured on the served 27B at 16k
prompt / chunk 2048: 25 mixed steps, 23 of them a continuation chunk,
and a first-chunk-only ladder fired on 0 of them.

FULL ``max_context`` coverage so
it bakes the widest block-table view — same reasoning as the prefill
sweep's continuation family.

ORDER, because ``put`` can evict under a tight graph-pool budget and
the rungs are not equally worth having: the scheduler bounds every
prefill chunk at ``chunk_prefill``, so that rung serves the modal full
chunk, and within a rung the continuation family carries the traffic.
So chunk-rung continuation first, then chunk-rung first-chunk, then
the remaining rungs ascending.

The prefill-token buckets :func:`precapture_mixed_graphs` will sweep.

Pure and config-only, so the sweep, the graph-pool cache key and the
scratch-pool pre-size all read ONE answer instead of three copies of
the same derivation. Empty when nothing arms the sweep.

How many MIXED graphs :func:`precapture_mixed_graphs` would capture.

``0`` when nothing arms the sweep.

This exists for the graph-pool budget CACHE KEY. The mixed ladder is
two graphs per prefill bucket and it lands in ``capture.cudagraphs``
alongside the decode / prefill ladders, so it moves the measured pool
— and the persist is monotonic, so a budget measured for a narrower
mixed ladder and reused for a wider one under-reserves the pool while
a budget measured for the wider one is pinned onto every later boot.
Same invariant, and same failure mode, as ``decode_shape_count``.

Widest FLAT token count any captured MIXED graph will run at.

The canonical mixed layout is ``d_cap`` one-token decode rows plus one
prefill row of ``bucket_n`` tokens, so the graph's flat width is
``(max_batch - 1) + max(bucket)``. ``0`` when nothing arms the sweep.

That width is REACHABLE ONLY BY THE CAPTURE, never by a live step: the
scheduler bounds a step at ``max_batched_tokens``, and the top bucket
is already clamped to it, so a live mixed step is at most
``max_batched_tokens`` flat tokens while the captured graph always runs
the padded ``d_cap + bucket_n``. Any boot-time scratch pre-sized from
``max_batched_tokens`` alone is therefore ``max_batch - 1`` tokens too
narrow for the capture — measured on the served 27B at max_batch 4 /
chunk 2048, where ``TQBufferPool.ensure`` refused ``num_tokens=2051``
against a cap of 2048 and BOTH families failed to capture. A capture
that cannot allocate is the same zero as a capture that is never
reached, and this is the number that closes it.

Capture whole-forward MIXED decode+prefill graphs
(``ARBI_MIXED_CAPTURE``; default OFF).

Two graphs per prefill bucket — one per prefill FAMILY (first chunk
/ continuation chunk; see the capture-plan note below) — at the
canonical composition ``(d_cap = max_batch - 1)`` decode-row slots +
one prefill row of ``bucket_n`` tokens (``B = max_batch`` rows). A
single decode-slot rung suffices: live steps pad ``d_live <= d_cap``
decode rows with
benign scratch rows (the decode-pad contract) — decode rows are one
token each, so padding even from ``d_live = 2`` to ``d_cap = 63``
is noise next to the prefill chunk's cost — and every mixed slate
fits (``d_live <= max_batch - 1`` whenever a prefill row is
co-admitted).

The prefill-token dim reuses ``cfg.prefill_cudagraph_buckets``
(auto-extended to cover ``chunk_prefill``, capped at
``min(max_batched_tokens, max_context)`` — same ladder discipline
as :func:`precapture_prefill_graphs`) so the operator tunes one
bucket ladder for both prefill flavors.

Gates mirror the prefill sweep: cuda_graphs on, TP == 1 (per-rank
capture asymmetry desyncs the lockstep collectives; the runtime
lookup path is bridge-gated), and the
:func:`_can_capture_prefill` model pre-flight (the mixed graph runs
the same varlen kernel set). Per-bucket failures log + skip — the
live path serves the missing shapes via split_mixed / eager.

Capture whole-forward prefill graphs for every configured bucket.

Mirrors :func:`precapture_decode_graphs` for the prefill shape.
One captured graph per ``num_tokens`` bucket — replays the full
prefill forward at ``(B=1, S=num_tokens)`` instead of the
piecewise per-layer pool's hundreds of layer graphs.

Walks ``cfg.prefill_cudagraph_buckets`` (a tuple of token-count
buckets), capped at ``cfg.batch.max_batched_tokens`` (a bucket
larger than the engine's per-step token budget can never be
replayed at runtime). Each successful capture registers in
``Engine.captured_graphs`` keyed at
``(B=1, S=num_tokens, lora_bucket=0, is_prefill=True)`` — the
same pool the decode + MTP-verify captures live in. The
``is_prefill`` key dim keeps the prefill captures from colliding
with batched-decode entries at the same ``B*S`` flat-token count.

Pre-flight is :func:`_can_capture_prefill` (refuses on hybrid /
SWA / MLA / recurrent models). On refusal we log + return
cleanly; the engine continues to boot with prefill running
eager (or piecewise, if that's enabled). A per-bucket capture
failure (typically OOM on a tight card) logs + skips that
bucket only — the rest of the sweep proceeds and the live path
handles the missing buckets.

Sleep-mode interaction. Prefill captures own their own
persistent buffers (no shared :class:`VerifyBuffers`); every one
is registered with :class:`SleepableTensorPool` so its data_ptr
survives release / resume. Release keeps the captured
``cudaGraph_t`` and only de-instantiates its exec, so resume
re-instantiates rather than re-capturing — this sweep is a boot
path, never a wake path.

Cold-boot two-pass cudagraph-pool measurement (GMU-0.99 no-OOM).

THE PROBLEM. At a high ``gpu_memory_utilization`` (0.95 / 0.99) the
deferred KV pool resize (build.py Phase 4c) sizes the KV pool against the
*cold-boot ESTIMATE* of the cudagraph private pool
(``eng.boot_state.graph_pool_budgeted_bytes``). That estimate under-counts the TRUE
capture pool, so KV is sized too big and the OOMs the card — and the cuMem map-time cap can NOT contain it, because
CUDA-graph capture allocations bypass the pluggable allocator entirely
(``my_malloc`` is never called during capture).

THE FIX — measure-then-fit. The cudagraph pool size is ~INDEPENDENT of the
KV ``num_pages``: a captured decode/prefill graph bakes a block-table whose
width is ``ceil(max_context / page_size) = max_blk`` (a fixed shape), NOT
the pool's page count. So we can MEASURE the true full graph pool against a
SMALL measurement KV pool (lots of free VRAM → nothing truncates), then
size the real KV pool around the measured value, then capture for real
against the right-sized pool → fits by construction.

This module owns the measurement pass, run IN THE SERVING PROCESS (no child,
no second weight load). The capture sweep records into the cuMem-backed
``capture.cudagraphs`` pool; PyTorch parks that MemPool until process exit
(``torch.cuda.empty_cache`` cannot return it — pytorch#145168), but the
physical pages behind it were mapped through the cuMem pluggable allocator,
so cuMem can UNMAP them directly (the raw ``cuMemUnmap`` + ``cuMemRelease``
the discard-sleep path drives). The sequence, all in one process:

  1. weights are already loaded (build.py SMALL measurement KV pool (ample free VRAM),
  3. capture the full sweep into ``capture.cudagraphs`` → measure the true
     peak (free-VRAM consumed delta),
  4. RELEASE the capture pool's cuMem physical and recreate the capture
     pools fresh (a mere discard-sleep would fault the Phase-5 re-capture:
     PyTorch reuses the pool's cached segments, whose physical is now gone —
     so the pool object is dropped and rebuilt so KV pool,
  6. build.py Phase 4c sizes the real KV pool around the measured value and


Correct VRAM sizing ALWAYS runs when its real preconditions hold — it is NOT
gated on ``cfg.vram_mode`` (bench and serving size KV identically; the only
thing ``vram_mode`` governs is whether an incomplete/undersized engine FAILS
LOUD or WARNS, see
:func:`arbi_serve.engine.build.verify_post_capture_headroom`). It is invoked
from :func:`arbi_serve.engine.build.build` only when ALL of:

  * ``cfg.cuda_graphs`` is on,
  * the cuMem driver is available (``cumem_pools`` + driver), and
  * this is a COLD boot — no valid persisted graph-pool measurement for
    this exact config (a warm boot already sizes KV right from the cache,
    so it single-captures).

The pass is boot-only; it adds zero per-step / per-token hot-path cost.

True iff this boot has NO cached graph-pool measurement for the config.

On a cold boot Phase 4c must reserve the graph pool from FREE VRAM (the
profiler ceiling was computed from an estimate). On a warm boot the
ceiling already excludes the cached graph pool, so Phase 4c keeps its
graph pool) and a known cache key.

True iff the cold-boot in-process graph-pool MEASUREMENT should run.

cuda_graphs + cuMem driver + COLD boot (no cached measurement), unless
force-disabled via ``ARBI_SERVE_GRAPH_POOL_MEASURE=0``. A warm boot already
budgeted the measured value, so we single-capture. NOT gated on
``vram_mode``: correct KV sizing must run identically for bench and serving.

The cuMem driver is a hard precondition: the in-process measurement releases
the capture pool's physical through the cuMem discard path, so a boot without
cuMem never measures — it falls back to the conservative cold reserve floor
(see build_phases_kv ``_setup_graph_pool_reserve``).

Run the exact MTP for the measurement).

Mirrors the build.py MTP captures are intentionally excluded — the
cold two-pass targets the no-MTP path (the MTP cold estimate is already
biased up; see ``cold_boot_graph_pool_upper_bound_bytes``).

Capture the full sweep against a SMALL measurement KV pool and return the
measured persistent capture-pool bytes (``None`` when it could not measure).

Builds a small measurement KV pool (so the card stays mostly free → nothing
truncates), captures the full sweep into ``capture.cudagraphs``, takes the
consumed-VRAM delta (the true value KV must reserve — it includes the
private-pool reserve + fragmentation that ``total_cudagraph_bytes()`` misses),
subtracts the separately-budgeted ``cudaGraphInstantiate`` reserve, and
persists the pool-only value to the shared budget cache.

Returns the persistent pool bytes, or ``None`` when the pass cannot run
(non-paged-KV / unknown per-page) or the measurement truncated (a capture
OOM even against the small pool). Leaves the captured graphs + the
measurement KV pool in place — the caller RELEASES them (in-process) or the
process exits (child-free path retired).

Drop + rebuild the cuMem capture pools so FRESH pools.

After :meth:`Engine.clear_all_cudagraphs` has reset every measurement graph,
the capture pools are graph-free but their cuMem physical is still mapped
(a captured pool's blocks are not returned by ``empty_cache`` — pytorch
#145168). A bare discard-sleep would unmap that physical while PyTorch still
holds the pool's cached segments, and the Phase-5 re-capture would reuse a
cached segment at a now-unmapped VA → ``cudaErrorIllegalAddress`` (verified).
So each capture pool is discard-slept AND its wrapper dropped so PyTorch's
segment metadata goes with it, then a fresh empty pool is registered under
the same name for

A pool that still holds LIVE allocations after ``clear_all_cudagraphs`` (its
persistent buffers are pinned elsewhere — e.g. the sleep registry) is left
mapped: its blocks are correct for Mirrors :meth:`NamedPoolRegistry.free_all`'s
per-pool discipline (discard-sleep only a graph-free pool; the
``capture.cudagraphs`` MemPool stays pinned in ``_CUDAGRAPH_POOL_PINS`` so it
parks with physical already released — the VA is reclaimed at process exit).

Release the measurement capture pool + restore the pre-measurement KV pool.

Undoes every side effect :func:`_measure_capture_pool` left, so build.py's
Phase 4c / Resets the captured graphs, recreates
the cuMem capture pools fresh (:func:`_recreate_capture_pools`), and rebuilds
the pre-measurement KV pool at ``restore_pages``. Logs the free-VRAM delta so
a leak is visible.

Measure the true capture-pool peak IN THE SERVING PROCESS, leak-free.

Runs the full capture sweep against a small measurement KV pool to measure
the persistent capture pool, then RELEASES the pool's cuMem physical (the
make-or-break step: the CUDAGraph MemPool PyTorch parks until exit is unmapped
directly through cuMem) and restores the pre-measurement engine state so
build.py sizes the real KV pool around the measured value and captures ONCE.
No child, no second weight load, no cold-estimate fallback when it succeeds.

Returns the measured persistent pool bytes (also persisted to the budget
cache, so the next boot is warm), or ``None`` when the measurement could not
run / truncated — the caller then falls back to the conservative cold reserve.

WHICH model, and WHICH drafter — the two identities the memory card's two
biggest rows have to state to be readable.

"Model weights — 13.56 GiB" names a quantity and no thing. The reader already
knows the server holds a model; what they cannot see is which checkpoint it is,
what it was quantised to, and what inside it is not quantised at all — and
those are exactly the facts that decide whether the number is what it should
be. The same is true one row down: "the drafter" is a row that is 0 on a boot
whose draft head ships inside the main checkpoint, and a 0 with no sentence
beside it reads as a broken feature rather than as a different topology.

READ FROM STATE, NEVER FROM A NAME. A served name that happens to spell out
``-exl3-4.0bpw`` spells it out for this checkpoint and not the next one, so
nothing here parses one. The quant format comes from the backend's own tensor-
key detection (:func:`~arbi_serve.weight_quant.loader.detect_backend`), the
bits-per-weight from the loaded EXL3 trellis geometry where the modules are in
this process and from the checkpoint's declared ``quantization_config``
otherwise, and the weight classes from the safetensors index's own
``data_offsets`` spans. The drafter comes from ``cfg.mtp`` through the same
identity function the K-calibration cache keys on, so the display and the cache
cannot disagree about which drafter is attached.

HEADER-ONLY AND CPU-ONLY. Every checkpoint read here is a safetensors header or
a ``config.json``; nothing is materialised, nothing touches a device, and every
function returns ``None``/``{}`` rather than raising, because this is a display
path behind an admin GET.

Which weight class one tensor key belongs to.

Quantisation is asked FIRST. A quantised output head is quantised weight
and belongs with the body; asking the role first would file it under a
class the card describes as unquantised, which is the one thing these
labels exist to tell apart.

Bits per weight as the CHECKPOINT declares it, or ``None``.

The checkpoint's own ``quantization_config`` block, read with the same key
set :mod:`arbi_serve.calibration.fingerprint` reads. A checkpoint that
declares no bit width gets ``None`` — the card then names the format and
says nothing about the width, which is the honest half of the answer.

Bits per weight off the LOADED EXL3 weights, or ``None``.

EXL3 stores each linear as ``trellis[in/16, out/16, K*16]`` where ``K`` is
the bit width, so the loaded buffer's own last dimension is the measurement
— no table, no constant, and it cannot be stale the way a config file can.
Only EXL3 carries the width in its geometry; every other backend falls back
to what the checkpoint declares.

What the ``Model weights`` row is holding: the checkpoint, its quant,
and the stored bytes of each weight class inside it.

Computed once per ``(checkpoint, loaded module)`` — see
:data:`_WEIGHT_FACTS_CACHE`. The facts describe a file on disk and a
module already built; an observer polling this endpoint must not pay to
re-open and re-parse the checkpoint's shard headers to be told the same
thing again.

Args:
    model_path: the served checkpoint's directory.
    model: the loaded module tree, when this process has one. Used only for
        the EXL3 bit width, which the loaded geometry states exactly.

Returns:
    ``{name, quant_format, quant_bpw, quant_source, components}`` — with any
    term this process cannot resolve ABSENT rather than guessed.
    ``components`` are ``{key, label, detail, bytes}``, measured from the
    safetensors index's ``data_offsets`` spans, so they are the bytes the
    checkpoint STORES. What the pool row holds is those tensors after this
    rank's shard and the loader's repack, which is why the console states
    where the class figures come from rather than implying they add up to
    the bar.

Store and return one checkpoint's facts.

Every exit of :func:`model_weight_facts` goes through here, INCLUDING the
partial ones. A checkpoint whose headers could not be read this time
cannot be read next time either — it is the same file — so re-deriving
the same partial answer on every poll would leave the most expensive
version of this read (the one that fails late, after opening the shards)
as the one that is never cached.

Which drafter is attached, and whether it has weights of its own.

Returns ``{kind, family, name, path, bundled, depth}``, or ``{}`` when
speculation is off. ``bundled`` is the case the row cannot state by itself:
a native MTP head ships inside the main checkpoint, so ``model.drafter`` is
0 and those weights are inside ``model.weights``. A 0 with no sentence
beside it reads as a broken feature.

``depth`` is the drafter's K — how many tokens it proposes per step — and
is present only where the config carries it. It is the one number that
separates "speculation is off" from "speculation is on and its weights are
somewhere else", which is exactly the pair a 0-byte drafter row is read as.

Consumer-side request state (:class:`ClientRequest`) + the
:class:`OutputApplier` that feeds it.

Design (``docs/engine_core_process.md`` §3 + §8)
------------------------------------------------
The engine emits **data** — :class:`~arbi_serve.engine.proc.messages.TokenOut`
/ ``AudioOut`` / ``FinishOut`` structs coalesced into one
``OutputBatchMsg`` per step — and never touches HTTP-side objects. The
consumer half lives here: one :class:`OutputApplier` owns a registry
``{request_id -> ClientRequest}`` and applies each decoded batch to it.
HTTP routes, the realtime session, and every in-process await-the-result
consumer read ONLY :class:`ClientRequest` state; the engine
:class:`~arbi_serve.engine.request.Request` keeps its own ``output_text``
(the detok / stop-scan working buffer) and the two never share mutable
fields — the applier writes the client copy from ``text_delta``.

The applier runs on the CONSUMER side: inline mode calls
``applier.apply(batch)`` directly on the (single) loop; thread mode hands
the batch object across via ONE ``call_soon_threadsafe`` per step (see
:mod:`arbi_serve.engine.output_bus`); the future process mode (P4) feeds
it decoded ``OutputBatchMsg`` objects from the ZMQ recv thread — the
applier is identical in all three.

Ordering contract (enforced in :meth:`OutputApplier.apply`)
-----------------------------------------------------------
Per request, within a batch and across batches (batches are FIFO):

  1. ``text_delta`` lands in ``output_text`` (and audio codes in
     ``audio_codes``) BEFORE ``finish_reason`` becomes visible;
  2. the tool stream is fed from the updated ``output_text``;
  3. ``new_token_event`` fires LAST — a woken consumer always observes
     the batch's complete state (text byte before wake, finish only
     after final text). This is the structural form of the engine's
     publish-order invariant (final token's text before finish_reason).

Finish representation: the engine emits a :class:`FinishOut` for EVERY
terminal transition — token-carried finishes (stop id / max_tokens /
context / stop string) and non-token finishes (error, cancel, timeout,
embed/rerank) alike. ``TokenOut.finish_reason`` stays part of the wire
contract but is never populated by this engine's emitter, so the applier
reads finishes from exactly one place.

Lifecycle: :meth:`OutputApplier.register` at request build;
the registry entry is removed when its ``FinishOut`` applies (or via
:meth:`OutputApplier.discard` on an admission failure), so a finished
request leaks nothing. A batch item for an unknown ``request_id`` is
logged and skipped, never a crash — a late batch after a cancel is
legal (the engine's finished-row drops bound it to the in-flight
deferred steps).

A submission was rejected by pre-publish validation (P3 intake).

Raised by the engine-side ``RequestFactory``
(:mod:`arbi_serve.engine.request_factory`) for every client-input
rejection — prompt too long, out-of-vocab token ids, unknown LoRA,
over-large ``mtp_k``, grammar compile failure — and carried back to
the API side as the negative :class:`SubmitAck` resolution. Routes
translate it to ``HTTPException(status_code, message)``.

Subclasses :class:`ValueError` deliberately: every pre-P3 rejection
on this path is a ``ValueError``, so callers still catching the
older ``except ValueError -> 400`` pattern keep working even if they
have not adopted the typed form. ``status_code`` is the HTTP class
the API side must surface (400 for all current rejections; the
grammar-compile normalization maps every non-ValueError escape into
a ``ValueError`` here too, so it also surfaces as 400).

Defined in this torch-free module (not ``request_factory``) so HTTP
route modules can import it without dragging torch into the
``dump_openapi`` path.

Consumer-side view of one in-flight request.

Created alongside the engine ``Request`` (``Request.client``) and fed
exclusively by the :class:`OutputApplier`. Everything a response
consumer needs lives here:

  * ``output_text`` — the client copy of the decoded text, built from
    ``TokenOut.text_delta`` appends (+ ``trim_to`` stop-string trims).
  * ``new_token_event`` — the consumer wakeup, set by the applier
    AFTER the batch's data is applied.
  * ``finish_reason`` / ``error`` — terminal state, visible only
    after the final text (see the module ordering contract).
  * ``output_token_count`` — maintained from ``TokenOut``s
    (text-lane tokens) and overwritten by ``FinishOut``'s
    authoritative count at finish (which also covers audio/control
    tokens that emit no ``TokenOut``).
  * ``audio_codes`` + ``codes_consumed`` — the omni speech lane;
    the applier extends ``audio_codes`` from ``AudioOut``,
    the streaming consumer advances its ``codes_consumed`` cursor.
  * ``embedding`` / ``rerank_score`` — pure-prefill task results
    (``FinishOut`` payload).
  * ``tool_stream`` — the per-request
    :class:`~arbi_serve.engine.tool_stream.ChatToolStream`; the
    applier drives it from the updated ``output_text``.
  * ``streamed_chars`` — the SSE generator's wire cursor
    (consumer-written only).
  * ``prompt_token_ids`` (a copy), ``sampling`` (shared reference),
    ``timing`` (shared reference; in-process only — the process
    transport (P4) replaces it with ``FinishOut.timings``).

Applies decoded :class:`OutputBatchMsg` batches to the registry.

One instance per engine (per API side in process mode). Single
consumer-thread discipline: ``apply`` always runs on the consumer
loop (directly in inline mode, via the thread marshaler's single
``call_soon_threadsafe`` otherwise), so no locking is needed —
``register`` runs on the same loop (``asubmit`` / ``submit``).

``apply`` never raises into its caller: a bad item is logged and the
rest of the batch still applies (one poisoned request must not drop
a step's signals for every other stream).

Finish EVERY live request with an error terminal state (§7).

The process-mode engine-death fan-out (P4b): when the engine
process dies, every in-flight ``ClientRequest`` would otherwise
wait forever for a This
applies the FinishOut-equivalent through the applier's own
ordering contract — terminal state written, THEN the wakeup —
so blocking waiters return and SSE streams end with an error
frame instead of hanging. Runs on the consumer loop like
:meth:`apply`; clears the registry (each request is terminal).

Apply one step's batch to the registered requests, in order.

Applies every ``TokenOut`` (text + count), then every
``AudioOut`` (speech codes), then every ``FinishOut`` (terminal
state + registry removal), then feeds each touched request's
tool stream and fires its ``new_token_event`` exactly once —
the module docstring's ordering contract. Batch order is the
engine's commit order (FIFO), so per-request text deltas apply
in production order.

Boot-time compression-capacity advisory.

The KV pool's realized page capacity scales with the codec (tkv packs ``cfg.cache.max_context`` and ``cfg.batch.max_batch``. At boot — once the pool
is sized and ``num_pages`` is known — this logs what the *compressed* capacity
actually supports, so an operator can see when a static cap is leaving the
codec's capacity unused (or, for bf16 weights, is over-ambitious).

This is ADVISORY only: it never mutates the configured caps (auto-applying
``max_context`` / ``max_batch`` would resize the cudagraph capture buckets and
the scheduler admission window, which is a separate, validated change). The
recommendations come from the pure helpers in
:mod:`arbi_serve.engine.compression_capacity` /
:mod:`arbi_serve.scheduler.compression_admission`.

Log (and return) the compression-aware capacity recommendations.

Returns a dict ``{kv_token_slots, max_context_at_max_batch,
max_batch_at_max_context}`` (also handy for tests / metrics). Pure +
side-effecting only via the logger; safe to call on any paged-KV boot.

Compression-aware capacity helpers for max_batch and max_context.

max_context is MODEL-derived: the hard ceiling is the model's trained context
window (``max_position_embeddings``); memory can only constrain it DOWNWARD,
never past what the model supports. Compression's win is serving the model's
FULL window at lower VRAM / higher concurrency — NOT a longer-than-model
window. ``derive_effective_max_context`` enforces that: min(model_ceiling,
memory_window).

max_batch is purely memory-derived: the concurrency that fits is
``kv_token_slots / per_seq_context``, which rises with the compression ratio.
``recommend_max_batch`` (scheduler/compression_admission.py) computes it.

All terms read the realized, per-layer, per-token compressed byte cost,
so they track smart-mix automatically.

Longest per-request context that fits when running ``concurrency``
sequences concurrently against the allocated KV pool.

``num_pages`` already reflects the active backend's compressed
``bytes_per_token`` (per-layer under smart-mix), so this
recommendation rises automatically with the compression ratio.

Effective per-request context the DFlash draft-KV slabs must size for.

The DFlash drafter mirrors a request's committed context in its own
persistent draft-KV slabs, so they must be sized to the longest
context a request can actually reach — NOT the raw config gate
(``max_context``), which is the model's RoPE ceiling and at "auto"
can be 256k even though the realized paged-KV pool serves far less.

When a profiled paged-KV pool exists, this is
``derive_effective_max_context`` (``min(model_ceiling,
pool_tokens / concurrency)``) — the same memory-realistic window the
max_batch fit gate and the capacity advisory use. Without a pool
(pure-recurrent target, or a pre-profile call), it falls back to the
raw ceiling.

``concurrency`` should be the resolved ``max_batch`` (the worst-case
concurrent sequences the slabs must serve simultaneously).

The advertised per-request context window for ``max_context="auto"``.

max_context is **model-derived**: the hard ceiling is the model's trained
context window (``max_position_embeddings`` — the RoPE table size; serving
beyond it walks the table out of bounds, which
``_validate_context_and_capture_bounds`` already rejects). Memory can only
constrain it DOWNWARD — never extend it past what the model supports.

So the effective window is ``min(model_max_context, memory_window)`` where
``memory_window`` is the longest context the realized (compressed) KV pool
can serve at ``concurrency`` sequences. Compression's win is therefore
"serve the model's FULL window at lower VRAM / higher concurrency", NOT a
longer-than-the-model window.

Runtime config-variant orchestration for live, bespoke config override.

A *config variant* is a (ServerConfig, RuntimeFlags-overlay) pair derived from
the CURRENTLY-ACTIVE config by a sparse delta. Each variant is keyed by a stable
config-signature and prepared+parked ONCE via the EXISTING stable-VA pool
machinery (:mod:`arbi_serve.engine.stable_va_controller`). Re-selecting a known
variant is a park/wake swap; a new capture-affecting variant pays a one-time
build+capture.

This is the engine-side core the admin endpoints
(``POST/GET /v1/admin/config``) drive. It reuses
``ServerConfig`` + ``RuntimeFlags`` + the param registry
(:mod:`arbi_serve.config_overrides`) + the stable-VA controller; it does not
introduce a parallel config schema.

Three delta classes (correctness first; instant when safe):

  * **known variant** — the signature already names a prepared member ⇒ an
    instant ``aswitch_to`` (or a no-op when it is already active).
  * **fresh-read-only** — the delta touches ONLY the fresh-read flags
    (recurrent_prefill_chunk / prefix_grouping[_window] /
    mtp_spec_disable_batch — see :data:`_FRESH_READ_FLAGS`)
    AND maps onto the SAME GPU member as the active one ⇒ install the new flag
    overlay live (no new member, no re-capture) and record the variant key so
    re-selection is instant. ``gdn_decode_num_warps`` is NOT one of these
    despite the name: it is scope="backend" in
    :mod:`arbi_serve.config_overrides` (a Triton kernel-launch config baked
    into the captured decode cudagraph, not a per-step fresh read), so a
    delta touching it always routes to build-required below.
  * **build-required** — any capture-affecting delta, or a snapshotted-at-build
    runtime-flag delta ⇒ build+capture a NEW member as the sole VRAM resident
    (park the active first), register it parked with its overlay, then wake it.

One overlay per member, in one place. Whatever route a delta takes, the overlay
it resolves to is installed live AND recorded on the member's residency record
(:func:`_adopt_overlay`), because the record is the copy a park keeps and a wake
re-installs — and two overlays can name the SAME variant (a flag restated at, or
returned to, its boot value signs identically), so a member serving one while its
record holds the other is not a difference the signature can catch.

Snapshot the boot config + (empty) overlay as the variant baseline.

The base is the config the engine booted with; every variant is a delta
against the currently-active config, but the boot base is what ``describe_
delta`` reports against and what a key collision resolves to. Idempotent.

The (cfg, flag_overlay) of the currently-active variant.

The active overlay is the stable-VA active record's stashed overlay (the
live config-override seam); the active cfg is ``eng.cfg`` (repointed by the
member switch). Falls back to the boot base when residency is not engaged.

Full active config snapshot for ``GET /v1/admin/config``.

Returns the resolved value of every registry param (cfg-path + flag), the
active variant key + its delta vs boot, the known-param catalogue, and the
live value domains for the params whose bound only this boot knows.

The residency member that is STILL mapped and active, or ``None``.

Read after a failed override to answer the only question the sticky
fault turns on: does the engine still have a config it can serve? A
build-required override that fails is rolled back transactionally
(:func:`_rollback_member_build` re-wakes the outgoing member and reinstalls
its overlay), and a delta refused BEFORE the park (``apply_overrides``
validation, ``assert_member_build_feasible``) never touched residency at
all — in both cases the outgoing member is intact.
``None`` means the rollback could not put one back, which is the case the
latch exists for.

Residency is necessary but NOT sufficient, so the resident member is also
asked whether it can still FORWARD (:func:`_unservable_reason`). A build
that mutated shared state and aborted before restoring it leaves a member
that is mapped, active and unable to serve; reporting that one as "still
serving" is the lie this function must not tell.

Every module graph the engine currently FORWARDS THROUGH.

``eng.model`` is not all of it. The speculative drafter is a SEPARATE
graph with its own quant linears and its own uids, and it is on the step
path for every speculated token — so a walk that stops at ``eng.model``
both fails to republish the drafter's kernel descriptors and reports a
member servable while the drafter cannot forward. GPU-measured: a rollback
whose republish covered the served model alone rebuilt 401 descriptors and
the next request still died, in the drafter, on a uid the walk never
reached.

Deliberately enumerated, never discovered by scanning attributes: the
process also holds the module graphs of PARKED members, and touching one
of those reads buffers at VA whose physical the park unmapped. The set is
exactly what is live on the engine.

Why the ACTIVE member cannot forward, or ``None`` if it can.

Checks the process-global EXL3 dispatch table against the live model: a
linear whose uid is not published there raises out of the custom op on the
first forward, so residency bookkeeping alone cannot answer "is this member
serving". The gap is what an aborted build leaves when a compaction RELEASE
hook ran without its rebind — and on a donor-share build the released
modules are shared by object identity with the member that was serving.

Cheap enough to run on the failure path only (one ``modules()`` walk), and
it is the failure path that decides between "CONTINUES to serve" and the
sticky fault.

Remember the config that is ACTUALLY serving, for a later restore.

Recorded only after an override succeeds, so it names a config the engine
has demonstrably built and switched to — never one that was merely asked
for.

Resume serving the last config that actually served.

The deliberate counterpart to the sticky refusal in
:func:`aapply_overrides`. That refusal latches only when a failed override
left NO member able to serve, so this is the way back from that state: the
residency rollback has since put a member back (or an operator is choosing
which config to return to), and what this call restores is permission to
serve, plus that member's flag overlay. An override that failed while the
previous member kept serving never latches and needs no restore.

Refuses loudly when no config has ever successfully served: there is
nothing to return to, and inventing one would be the silent substitution
this whole path exists to prevent.

Apply a sparse config delta against the currently-active config.

Computes the resulting (cfg, overlay) variant, keys it, and:

  * known key ⇒ instant ``aswitch_to`` (or no-op if already active),
  * fresh-read-only delta on the same GPU member ⇒ live overlay install,
  * else ⇒ build+capture a new member, park-others, switch.

Returns the now-active full config (``current_config``) plus timing/route
metadata. Requires stable-VA residency engaged.

A FAILURE does not by itself suspend serving. The build-required route
rolls back transactionally and a delta refused before the park never
touched residency, so the outgoing member is normally still mapped and
stepping; the call raises and the log names the config that is actually
serving, and nothing is latched. The sticky fault
(:func:`~arbi_serve.engine.infra_health.note_swap_fault`) is reserved for
the case the rollback could not put a serving member back — where the
engine genuinely has nothing to serve.

``rollback_on_failure`` is an affirmative, default-off authorization to
resume serving the previous config in THAT case — for a caller (a bench
harness) that would rather keep serving than inspect the failure. It means
"on failure, resume serving the previous config"; it does not mean "undo
the residency change", which happens either way. It runs the same
:func:`arestore_last_valid` an operator would call by hand, and the failure
is still raised — with the restore named in the message, so a caller can
never read a resumed old config as the new one.

``drop_previous`` is an affirmative, default-off opt-in that, on a
BUILD-REQUIRED swap only, fully reclaims the member this call supersedes
(``adrop_variant`` — the same primitive an ephemeral A/B teardown uses)
instead of leaving it parked. The default (``False``) preserves today's
A/B-friendly behaviour: a parked member costs a real, standing VRAM tax
(captured-graph exec memory + torch-MemPool held-free segments neither
``empty_cache()`` nor the cuMem sleep path can return) but stays instantly
re-selectable. Set ``drop_previous=True`` for a caller that walks through
many genuinely distinct configs it will never revisit (a floor-and-grow
bit-allocation search is the motivating case) — without it, that standing
tax is paid again on EVERY swap with no bound, a real monotonic VRAM leak
that eventually OOMs the card. Silently a no-op on the boot/declared-pool
member (never dropped — every same-model config variant donor-shares its
weights) and on the non-build-required routes (nothing new was parked).

How many superseded runtime variants may stay parked (``>= 0``).

Read fresh per swap from ``ARBI_SWAP_VARIANT_CACHE`` so an operator can
change the retention policy without a restart. The bound is on RETAINED
(parked, ephemeral) variants only — the active member and every declared
residency member (the boot model, a ``--pool-member`` model) sit outside
it and are never evicted by it.

How many parked variants may survive a build-required swap's eviction.

The bound counts the set as it will stand AFTER the swap, and the swap is
about to park the member it supersedes — so when that member is itself an
ephemeral variant, one slot belongs to it and the eviction must free it
now. Deriving that here rather than at the call site is what keeps the
arithmetic testable without a build.

How many parked variants may survive an instant switch's eviction.

Same arithmetic as :func:`variant_cache_keep_for_build` with the switch's
two differences: the member being woken LEAVES the parked set (so it is
not one of the retained), and the member being parked joins it. Both
routes have to enforce the bound or it is not one — an instant switch
parks the outgoing member exactly as a build does, so enforcing only on
the build route lets a caller alternating between prepared variants grow
the resident set without bound anyway.

Registry records that are runtime config-variants (evictable), in
registration order.

A record is ephemeral iff it is BOTH a residency record and a key
``config_override`` created — the same conjunction
:func:`droppable_variant_refusal` refuses on. A live-overlay variant owns
no record (it shares the active member's) and so is absent here; a declared
pool member owns a record but no variant entry and is absent too.

Stamp ``key`` as the most recently used variant (drives LRU eviction).

A monotonic sequence, not a clock: two swaps inside the same coarse timer
tick must still order, and a wall clock that steps backwards must not
reorder the retained set.

Evict least-recently-used parked variants until at most ``keep`` remain.

THE BOUND IS THE WHOLE POINT. Every superseded member left parked costs
resources that neither ``empty_cache`` nor the cuMem sleep path returns —
driver-side cudaGraphExec memory, a private torch MemPool's held-free
segments (pytorch#145168), the pinned host RAM its pools were offloaded
into, and a VA arena. Those are per-member costs, so retaining every
superseded member multiplies them by the swap count: a walk of N configs
is O(N) in device memory, host RAM and VA, and ends at whichever runs out
first. Bounding the retained set makes all four O(1) in N. What the
per-member cost IS on a given deployment is what that boot reports — the
build's "VRAM at profile" line and the park/wake census — not a constant
this docstring can hold.

Evicts oldest-used first so the cache keeps what a caller is actually
alternating between. Never touches the ACTIVE member or a declared
residency member — :func:`droppable_variant_refusal` is the single
predicate for that, and it is consulted per key rather than re-derived.
``exclude`` additionally spares the member a switch is about to WAKE:
it is not active yet, so no other predicate protects it, and destroying
the target of the switch that called this is the one outcome that cannot
be right.

Never raises: an eviction that cannot run leaves the member parked, which
is exactly the pre-existing behaviour, and the swap that asked for the room
carries on with whatever is free. Returns the keys actually torn down.

Record ``overlay`` as member ``key``'s own.

The record is the only copy of a member's overlay that outlives its
residency: a park keeps it and the wake re-installs it
(``stable_va_controller._install_flag_overlay``). An overlay installed
live and not recorded here is therefore served only until the member is
parked, and the wake then replaces it with whatever the record still
held — an override that stops applying at a swap nobody connected to it.

Install ``overlay`` as the active-member flag overlay (if changed) and
re-snapshot the live runner's cached flag attrs.

``set_active_flag_overlay`` makes ``runtime_flags()`` return the overlaid
values immediately (fresh-read flags pick it up automatically). The runner
additionally caches a handful of ``capture_affecting=False`` flags into
instance attrs at build (the hot path reads a plain bool, not a dict
lookup), so a live overlay flip of one of those must also poke the runner
to re-read them — otherwise the flip is invisible until the next rebuild.

Build+capture a NEW member as the sole VRAM resident, register it parked.

Runs under the engine critical section (scheduler drained + paused) so no
forward runs against half-mapped physical during the park-active→build→park
window. Mirrors :func:`prepare_model_pool`'s per-member step but for ``aswitch_model(key)`` then performs the user-visible swap.

TP>1 lockstep + commit/abort barrier. Every rank reaches this function
deterministically from the same broadcast delta and runs the identical
park-build-capture sequence; the build's NCCL collectives rendezvous by
construction. After each rank attempts its build, a cross-rank
confirmation barrier (:func:`member_build_barrier.confirm_build_across_ranks`)
decides the cluster outcome: a build that FAILED on ANY rank rolls back
on EVERY rank (the locally-succeeded ranks drop their freshly-built
orphan and re-wake the previous member too), so the ranks never desync
into one serving the new member while another serves the old. Single
process / TP=1 degenerates to "local outcome is the verdict".

Tear down every ephemeral parked variant except ``keep``. Keys freed.

The rollback's funding source, and the same primitive the bounded cache
evicts through, so there is one teardown path and not two that drift.
Declared residency members are refused by
:func:`droppable_variant_refusal` and are never touched — the baseline the
last-resort restore needs has to survive this.

Never raises: it runs on the path that exists because something already
failed.

Last resort after a failed re-wake: bring the DECLARED baseline back up.

A wake that fails part-way leaves its member half-mapped and un-retryable,
so the member the rollback was aiming at is gone as a serving target. The
boot model is not: it is a declared residency member, it was parked by the
ordinary path, and every same-model variant donor-shares its weights — so
it is the one record that is both intact and guaranteed present. Discard
the wedged member's physical + record, then wake the baseline.

Serving the BOOT config is not serving what the caller asked for, and the
caller still gets a refusal saying so. It is, however, a running server
instead of a process holding a card with nothing resident, which is the
only outcome the alternative offers. Returns the key now active, or
``None`` when there is no baseline to fall back to (the failed member WAS
the baseline) — the caller then reports the unrecoverable state.

Drop a failed/aborted member build and re-wake ``prev_active``.

This is a transactional rollback: leave the engine serving
``prev_active``. Used on a local build failure and on a cross-rank abort
(a peer rank failed), so the rollback is identical on every rank.

The failure may have struck during the build (partial/no new physical),
after the build succeeded but registration was refused (the new member's
full physical is live-mapped and occupies the VRAM ``prev_active``'s
remap needs), or after a successful local build and register but the
cross-rank barrier aborted (``registered=True`` — the orphan is a full
registry record we must also drop). Either way we discard whatever the
build mapped before re-waking ``prev_active``; otherwise
``resume_engine_state_sync``'s ``map_range`` for prev's KV slab has no
room and raises, stranding the boot model parked with active=None
(engine 500s on every request).

``prev_cfg`` is the config ``prev_active`` was built and captured with,
snapshotted by the caller before the builder repointed ``eng.cfg``. It is
what the engine is repointed to, because it is what the re-woken member's
captured graphs and pools actually implement. Falling back to the BOOT
config is only correct when the previous member IS the boot member: for a
rollback landing on an earlier VARIANT it reports one config while serving
another's graphs.

Re-publish the re-woken member's quant kernel descriptors.

The last thing a rollback owes the member it is putting back is the
ability to forward. A member build that aborts anywhere after the load's
compaction RELEASE pass leaves that pass unmatched: each quant linear's
derived kernel descriptor is torn down and its uid dropped from the
process-global dispatch table, and the matching
``rebind_after_compaction`` is only reached at the END of a build that
completes. On a DONOR-SHARE member the released modules are the donor's
OWN objects (``share_donor_head_quant_modules`` installs the donor's head
modules on the member's graph by identity), so the uids that go missing
belong to the member this rollback is restoring — and a residency wake
restores mappings, not descriptors.

Firing the rebind hook here re-derives every descriptor over the buffers
the woken member now holds and re-registers each uid, which is what makes
"the previous config CONTINUES to serve" a true statement rather than a
claim about bookkeeping. Idempotent and cheap on a member that never lost
one: the hook rebuilds from the live buffers either way, and a member
holding no quant linears fires zero hooks.

A failure here is logged CRITICAL rather than raised: raising would
replace the build error the operator has to read with a teardown error,
and the consequence is detected anyway — the caller's fault handler probes
the restored member for missing descriptors (:func:`_unservable_reason`)
and latches the sticky fault, so ``/health/ready`` goes 503 either way.

The parked snapshot, shaped like the engine surface the graph walk reads.

:func:`arbi_serve.engine.sleep.iter_owned_cuda_graphs` reads two things off
an engine: the cudagraph POOLS it walks structurally, and every attribute in
``_GRAPH_OWNER_ATTRS`` — the owners that keep graphs in their own containers
and can only be reached by declaring them (``owned_cuda_graphs``). A shim
built from a hand-written subset of those names silently answers ``None`` for
the rest, and a declared owner the walk cannot reach keeps its captured
``cudaGraph_t`` for the process's life. Reading the roster is what keeps the
two ends of that contract on one seam.

``mtp_driver`` is a forwarding alias for ``drafter`` on the engine
(``Engine.mtp_driver``) and therefore has no key of its own in the per-model
snapshot; the shim reproduces the alias so the walk sees the same surface it
sees live. The walk dedupes by identity, so the pair costs nothing.

Reset the captured cudagraphs held in a PARKED record's OWN snapshot.

A parked member's graphs live in ``rec.extra["state"]``
(``attrs["cudagraph_pools"]`` + ``attrs["captured_graphs_multi"]`` + the
declared owners of :func:`_snapshot_graph_owner_shim`), NOT on the live
engine — so the engine-scoped ``eng.clear_all_cudagraphs`` would reset the
ACTIVE member's graphs instead of this variant's. Point the shared
owned-graph walker at the snapshot's own surface and ``reset()`` each, which
returns the retained ``cudaGraph_t`` driver memory a light park kept mapped
(GC alone does not reliably drive ``~CUDAGraph.reset()``).

A park destroys the EXEC and keeps the topology, by design — the wake
re-instantiates. A drop is the other lifecycle: nothing will wake this
member, so the topology is driver-side bytes that no allocator we own can
see, and only a ``reset()`` returns them.

Returns the number of graphs still reporting an instantiated exec after the
sweep — the residue receipt reads it.

Physical still cuMem-MAPPED under one member's tag namespace.

Reads the allocator's own per-tag ledger (the same rows the boot log
prints), not a running counter, so a teardown that returned nothing and one
that was never asked read differently. Zero when the driver is unavailable
(the CPU lane), which is also the value a released namespace has.

One receipt: did the dropped member return everything it owned?

A drop is the only lifecycle edge with no counterpart — nothing wakes this
member, so anything it still holds is held for the process's life, and the
next member's KV grow sizes itself from a card that much smaller. The
symptom (a served context that shrinks a little on every A/B flip) is many
rebuilds downstream of the cause, so the cause states itself here.

Three residues, each the failure of one teardown step: physical still mapped
under the member's tag namespace (the unmap), captured graphs still holding
an instantiated exec (the graph reset), and a bookkeeping entry that still
names the key (the scrub). Fires
:data:`_DROPPED_MEMBER_VRAM_NOT_RETURNED`, which must not fire.

Grow the ACTIVE member's growable KV back to what it held, against the
free VRAM a serving step does not need.

``max_pages`` is the restoration cap: the reclaim RESTORES, it does not
maximise. Growing past it would size the pool from a drained moment's free
VRAM, which includes the reserve the next serving step is about to take.
Omit it and the cap is resolved from the ACTIVE residency member's own
history (``StableVaResidencyController.active_kv_restoration_cap`` — its
high-water page count, not what it last happened to hold); a caller that
knows the member it is restoring, like the wake seam, passes it directly
because the record it is waking need not be the active one yet. ``0`` /
absent on both routes means no cap is known and the free bound governs
alone.

TWO owners, one reason. A variant DROP frees the variant's physical, and a
budget-capped WAKE re-provisions a discard-parked slab
(:meth:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool.wake_growable_kv`)
against free VRAM read before the wake had finished. Both leave the active
member holding fewer pages than the card can back, and nothing else grows a
woken pool: the cap only ever lowers the count, and the next park
re-provisions from the already-lowered size, so the shortfall compounds
across swaps into a permanent KV-capacity tax while ``free`` VRAM sits idle.

Closes by re-deriving the served ``max_context`` from the pool it leaves
(:func:`_resync_served_context`) — the pool's size is what that bound is a
statement about, and this is the only place a wake or a drop changes it.

Uses the REBUILD-path grow (:meth:`grow_kv_to_fit`), NOT the post-capture
grow: a variant build's profiling clobbers the engine-global
``boot_state.kv_serving_ceiling_pages`` down to the co-resident (reduced)
budget, and switching back to the baseline never restores it — so
``grow_kv_after_capture`` (bounded by that clobbered ceiling) would leave the
freed VRAM idle. ``grow_kv_to_fit`` instead fills free VRAM up to the slab's
OWN reserved-VA ceiling (``num_pages``, fixed at boot and preserved across
the park/wake), restoring full KV capacity. Mirrors the ``rebuild_pool``
reclaim (backend swap / sleep-wake), plus an all-rank-MIN collapse of the
free-VRAM read so every TP rank maps the identical page count (a per-rank
divergence would fault under load). Must run under the engine critical
section (drained). Best-effort: a grow that cannot restore full capacity
(real scarcity) leaves the pool at what fit; the returned page count lets the
caller report the realized capacity.

Re-derive the ADMISSION-time served context from the pool as it now is.

``max_context=auto`` is a statement about the REALIZED pool
(:func:`~arbi_serve.engine.build_memory_sizing.finalize_served_max_context`),
and until now only a BUILD made it. A wake re-provisions the woken member's
slab and this reclaim regrows it, so by here the pool is a different size
than any build ever measured and the bound no longer describes it. Both
directions are wrong and one is unsafe: too high and admission accepts a
request the pool cannot hold, too low and the member serves a fraction of
what it mapped.

It is also the only thing that repairs the bound after ANOTHER member wrote
it. A capture-affecting delta over RuntimeFlags alone produces no
ServerConfig delta, so the variant is built against the SAME ``cfg`` object
the boot member holds, and its build's auto-narrow rewrites the field both
members read. Re-deriving here from the pool that is actually resident
makes the field describe the member that is actually serving.

Bounded, as at boot, by what the WOKEN member's own sizing-time buffers can
address — ``verify_buffers.max_pages`` is per-member and already repointed
— so this can never admit past the block tables. Best-effort: a pool whose
context cannot be re-derived keeps the bound it had.

Why ``key`` may not be dropped, or ``None`` when it may.

ONE predicate for both callers, because they need the same answer and
disagree only on what to do with it: :func:`adrop_variant` raises it at an
operator who named the member, while a switch that reclaims what it leaves
behind treats it as "keep this one" and carries on.

A declared residency POOL MEMBER — the boot model, a ``--pool-member`` model
— is a registry record but not a runtime config-variant (it is absent from
``_cfgvar_variants``). Dropping one tears a model out of the multi-model
pool, and every same-model variant donor-shares the boot member's weights,
so it is also the record a switch BACK is heading for. Only variants built
through ``config_override`` are ephemeral.

Release a PARKED runtime variant's VRAM, VA and record. Bytes freed.

The teardown half of :func:`adrop_variant`, and the half a switch runs in
the one window where it helps: BETWEEN the park and the wake, while the
outgoing member sits parked and the incoming one has yet to ask for its
memory. Called there, it is what makes a card sized for ONE member
able to complete the swap at all.

Caller's contract: hold the engine critical section (drained), have ``key``
already PARKED and not active, and have checked
:func:`droppable_variant_refusal`. Does no KV regrow — the wake's own
reclaim seam owns that, and between a park and a wake there is no active
member whose pool could be grown.

``successor`` names the member that inherits the dropped one's standing as
the last config known to serve. A switch must pass its TARGET: between a
park and a wake the registry has no active record to fall back to, and a
``config_restore`` pointed at ``None`` restores nothing.

Every engine-live-member-scoped primitive is avoided, exactly as in
:func:`adrop_variant`: only the variant snapshot's graphs are reset and only
its ``"<namespace>/"``-tagged physical is unmapped, so the member the engine
is about to wake keeps everything of its own — including the weights this
variant donor-shares from it.

Forget ``key`` everywhere the variant registry names it.

Leaves the residency looking exactly as it did before the variant existed,
so a later override with the same delta re-prepares rather than resolving to
a record that is gone. Covers a live-overlay key (no record of its own) too.

Ephemerally drop a PARKED (non-active) config variant, reclaiming ALL its VRAM.

The inverse of a build. After an A/B a build-required variant sits
co-resident/parked (a standing +VRAM tax on KV/memory under light-park);
this returns the residency set to EXACTLY the baseline member — no permanent
tax. Runs under the engine critical section (scheduler drained) and refuses
to drop the ACTIVE member.

Composes the same teardown primitives as :func:`_rollback_member_build`, but
for a variant that WAS made active and is now parked while a DIFFERENT member
is live. Every engine-live-member-scoped primitive is therefore avoided: this
NEVER calls ``eng.clear_all_cudagraphs`` (would reset the active baseline's
graphs), ``release_engine_state_sync`` / ``_park_mechanics`` (act on the live
member), or ``sleep_namespace(base)``. It resets only the variant snapshot's
graphs and unmaps only the ``"<namespace>/"``-tagged physical.

A live-overlay variant (``gpu_member`` is the active member, no record of its
own) has no member to drop — its bookkeeping is scrubbed and the call is a
no-op on VRAM.

Drop EVERY non-active config variant, returning residency to the single
active baseline member.

The ephemeral-A/B reset: after a run, reclaim all variant VRAM/KV/arena so
there is NO permanent tax. Iterates the residency records (the authoritative
set of GPU members), not ``_cfgvar_variants`` (which also holds live-overlay
keys that share the active member and own no record); each non-active record
is torn down via :func:`adrop_variant`. Any lingering live-overlay
bookkeeping is scrubbed at the end so the pool listing is left clean too.

CPU-side prep worker for prefill batches.

The :func:`_build_batch` path runs synchronously in front of the GPU
launch on every step. At admission time we kick off a small piece of
per-request CPU prep — convert ``prompt_token_ids`` (a Python list)
into an int32 numpy array — so step-time slicing into the pinned-host
buffer skips Python's list→numpy coercion.

For a 2048-token prompt at chunk_prefill=512 the prefill is split into
4 steps of 512 tokens each. The serial path was:

    flat_ids = [];  flat_ids.extend(prompt_token_ids[s:s+512])  # 512 PyObject borrows
    pb.h_input_ids.numpy()[:N] = flat_ids                       # list → int32 numpy

With prep done at admission:

    pb.h_input_ids.numpy()[:N] = req._prep_prompt_ids_np[s:s+N] # C-level memcpy

The list-extend + numpy coerce runs per chunk on the build path;
pre-building an int32 numpy view at admission moves that work off the
TTFT critical path (the prep runs while the scheduler thread is doing
other slate work).

Decode rows are 1 token each so the CPU prep there is sub-microsecond
and not worth optimizing through this seam — those still go through
the original Python list path. Chunk-prefill is naturally supported:
the prep numpy array covers the whole prompt; per-step ``_build_batch``
slices into it at ``[prompt_consumed : prompt_consumed + n]`` for that
step's chunk. The position arange is also full-prompt-length so the
same slice-by-offset pattern works.

Worker model:

This is GIL-bound CPU work (numpy from a Python list). A single
shared :class:`~concurrent.futures.ThreadPoolExecutor` (size 2) lets
multiple admissions overlap with each other and with the engine
step thread; for a single admission the prep runs while the scheduler
is still doing its FCFS / preempt / quota dance, so by the time the
slate reaches ``_build_batch`` the numpy array is usually ready.

Best-effort: if the prep hasn't completed by the time
``_build_batch`` reads ``req._prep_prompt_ids_np``, it's still
``None`` and the legacy Python-list path runs. No correctness
delta, just no speedup that step.

Materialize the int32 prompt + position arrays.

Pure CPU / numpy. Runs on the prep executor — must not touch the
request object directly (consumers wire the result back via
:func:`prep_request`).

``prompt_ids_np`` is an integer array already holding the prompt (the
admission vocab guard's walk of the list); when given, the int32
prompt array is a cast of it and the list is not walked again.

Returns ``(prompt_ids_np, positions_np)`` — both ``int32`` and
contiguous so the per-step slice into the pinned-host buffer is
a straight memcpy.

Kick off the per-request CPU-prep on the executor.

``prompt_ids_np`` is the admission guard's integer array of the same
prompt, when one was built; it spares the prep a second walk of the
list.

Called from the admission paths (``submit`` / ``asubmit``). When
no asyncio event loop is running (sync test harness), runs the
prep inline so the test path still gets the speedup at step time.
The executor + asyncio path is the production hot path.

Idempotent: a second call against an already-prepped request is a
no-op (early return on ``req._prep_prompt_ids_np is not None``).

Drop the cached prep arrays + cancel any in-flight prep future.

Called from finish / cancel paths to release the int32 buffers
(~8 KB at prompt 2048 — small but worth dropping promptly on
long-tail finish hooks). Idempotent.

CriticalSection — drain / mutate / resume context manager.

The live-state-mutation surfaces (:mod:`engine.swap_admin`,
:mod:`engine.sleep`, :mod:`engine.cudagraph_admin`, calibration
reload) all follow the same dance:

    1. acquire ``_swap_lock``
    2. set ``_draining`` and poll ``while eng.requests``
    3. drop pool / page_table / scheduler / attn_ops
    4. mutate (build new backend, apply calibration, …)
    5. clear drain + ``_wakeup.set()``

This module collapses that into ONE async context manager scoped by
:class:`StateKind` (or the ``ALL`` sentinel for whole-engine sleep).
:class:`CriticalHandle` is the only mutation surface inside the
section; the typed handle makes mutation of unrelated engine state an
:exc:`AttributeError` by construction.

Usage::

    async with eng.critical_section(kind=StateKind.PAGED_KV) as handle:
        handle.set_backend(new_backend, spec="paged_kv:tkv-k4v4")
        handle.rebuild_pool()

The admin surfaces (``aswap_attention_backend``,
``areload_calibration``, ``areload_model``, ``release/resume``) all
reduce to a few lines of handle calls inside one ``async with``.

Drain semantics: all swaps drain all in-flight requests; the
``kind=`` parameter scopes the lock + the typed handle but not the
drain. The kind-locking is still load-bearing: it prevents two swaps
for the same StateKind from racing while letting swaps for different
kinds serialize independently. Rebuild ordering is deterministic
because each kind has its own lock; the whole-engine sleep path
acquires every kind's lock in a fixed order.

The mutation surface a :class:`CriticalSection` legitimately permits.

Constructed by :class:`CriticalSection.__aenter__`; never
instantiated directly by callers. The set of methods enumerates
every mutation that's safe inside a drained section. Anything
outside this list (e.g. mutating ``eng.requests`` directly) is
not exposed and is therefore an :exc:`AttributeError`.

Adding a new mutation surface = explicit Protocol amendment + a
matching admin endpoint, not a silent expansion.

A critical section's in-flight drain exceeded its deadline.

Raised (only when a ``drain_deadline_s`` is set) instead of polling
forever: under sustained load a config-swap/backend-flip could otherwise
wait indefinitely for in-flight requests to finish, wedging the engine and
``/health``. The swap caller catches this and refuses cleanly (HTTP 409)
with the engine left untouched — no locks held (``__aenter__``'s handler
releases them), no drain flag stuck. Unbounded callers (sleep / model
reload) pass no deadline and keep the original wait-forever semantics.

Async context manager that drains in-flight requests, exposes a
:class:`CriticalHandle`, and resumes scheduling on exit.

Scoped by ``kind``: a :class:`StateKind` for backend / calibration
swaps, or :data:`ALL_KINDS` for sleep / full-engine reloads.

Drain semantics: poll ``eng.requests`` until empty, with
``await asyncio.sleep(0.005)``.

Locking. Each kind has its own :class:`asyncio.Lock` in
``eng._kind_locks`` so swaps for different kinds serialize
independently. The :data:`ALL_KINDS` section additionally acquires
every per-kind lock in a deterministic order so concurrent kind-
scoped swaps cannot proceed mid-sleep.

Backwards-compat: the legacy ``eng._swap_lock`` is still acquired
so external callers that take it (e.g. health probes that read
``_draining`` for shutdown) keep working unchanged.

Whether a :class:`CriticalSection` body currently owns the engine.

The predicate a run loop asks before it schedules: while it is true the
engine's per-model state is being replaced, and the body may be driving
steps of its own through the very scheduler and step function this loop
would use. Defined here, beside the section that raises and lowers it, so
a loop can never test a different question than the one the section
answers.

``getattr`` because the loop is also driven by the minimal engine doubles
the CPU lane builds, which carry only what the loop body reads; for them
the honest answer is "nobody is mutating", not an ``AttributeError``.

Poll drain progress for a :class:`CriticalSection`.

Returns whether any request is in flight. The drain is engine-wide
and deliberately does not refine by ``kind`` — for two reasons:

* It cannot refine. An engine serves one active model at a
  time (swaps are sequential, under an ``ALL_KINDS``
  section), and every in-flight request traverses every layer, so
  it holds state for every ``StateKind`` the model uses. A section
  only ever targets an in-model kind (or ``ALL_KINDS``), so the set
  of "requests using ``kind``" is always the full request set —
  a kind filter would be identical to ``bool(eng.requests)``.
* Engine-wide is safer. A freshly-admitted request may not have
  allocated its per-kind ``state_handles`` yet; ``bool(eng.requests)``
  waits for it, whereas a ``kind in req.state_handles`` filter would
  let a swap race ahead of an in-flight request that has not built
  its page list.

``kind`` is accepted so the signature matches the per-kind lock
scoping the caller uses (different-kind sections don't serialize),
but the drain itself is engine-wide.

Residency exclusion: requests parked in ``eng._multi_group_queue``
are pre-admission — they are not on the GPU and hold no model state;
they sit waiting for their resident to become active. A run-loop-driven
resident swap (``maybe_advance_active_group`` → ``eng.aswitch_model``) is
triggered by such a request, so counting it here makes the swap's drain
wait on the very request it must admit → permanent self-deadlock
(``/health/ready`` 503 forever). Subtract the queued count so the drain
waits only for requests actually admitted to the scheduler / on the GPU.
A freshly-admitted scheduler request (in ``eng.requests`` but not yet
holding per-kind handles) is not in the graceful queue, so it is
still counted — preserving the safety the engine-wide drain gives.

Engine-wide count of requests admitted to the scheduler / on the GPU.

The single source of truth for "is the engine busy right now": the
size of :attr:`Engine.requests` minus the pre-admission
``_multi_group_queue`` backlog (those are parked off-GPU awaiting their
resident and hold no model state — counting them would self-deadlock a
resident swap and falsely report a freshly-booted, idle engine as busy).

Read lock-free from any thread (a plain ``len`` under the CPython GIL);
a stale read is at worst off by the handful of mutations applied during
the read — the same benign race :func:`has_inflight`'s drain poll and
the admission gate already rely on. Used by both :func:`has_inflight`
(swap drain) and the background flat-weight warm dump's idle gate
(``engine/sleep.py``) so the two never disagree on what "idle" means.

Read-only access to the engine for callers that need to
invoke a build helper that lives in :mod:`engine.build`. The
handle itself enumerates the safe mutations; callers that
need to reach beyond the enumerated set still go through
``handle.engine`` and bear responsibility for the invariant.
Every escape-hatch call site should ideally grow into a
first-class handle method or be rejected as out-of-scope.

Drop and re-allocate the per-:class:`StateKind` slab pool.

Called by sleep mode phase2 (release / resume VRAM at the same
cuMem-stable VAs), by backend swaps that change
``bytes_per_token``, and by full-engine reloads.

This teardown-and-rebuild affects all state kinds (the
pool currently holds every kind's slab in one allocation).

Captured-graph invalidation. Every captured cudagraph baked the
old pool's per-layer slab + per-op TKV-scratch ``data_ptr``s
into kernel-launch arguments at capture time. ``eng.pool`` / ``eng.attn_ops``, those backing tensors lose
their last Python reference and the caching allocator returns
their segments to the named MemPool free lists. The subsequent
:func:`torch.cuda.empty_cache` walks every allocator's free
segments and The first device sync after
empty_cache (or empty_cache's own internal sync) surfaces the
stale-pointer access as ``cudaErrorIllegalAddress``. Clearing
every captured graph before the teardown drops the last
Python reference to those captured kernels (and to the
persistent input/output buffers in Mirrors the explicit drop the calibration-reload path takes
when per-channel scales change (see
:func:`arbi_serve.engine.swap_admin.areload_calibration` —
the ``eng.clear_all_cudagraphs()`` branch).

Re-bind named calibration buffers (centroids, boundaries,
per-channel scales) without recompile.

``cal`` is the parsed calibration dict the active backend's
loader understands. Both paths do ``buffer.copy_(new_values)``
semantics — see
:func:`arbi_serve.engine.build.apply_calibration_to_tkv_ops` and
:func:`arbi_serve.engine.mla_install.apply_calibration_to_mla_ops`.
Pool teardown is not triggered; the existing tensors are
rewritten in place.

Boot-time cuBLAS / cuBLASLt workspace routing into a named pool.

PyTorch's caching allocator backs cuBLAS and cuBLASLt with a private
workspace, lazily allocated on the FIRST GEMM that touches each
``(device, stream, handle)`` and kept for the process lifetime (one
caching allocator, but because cuBLAS issues it implicitly — there is no
``torch.Tensor`` Python referrer — it escapes every
:class:`~arbi_serve.runtime.named_pool.NamedMemPool` context unless the
very first GEMM on that stream/handle happens to run inside one.

On a hybrid GDN model (Qwen3.5-0.8B) the boot-time warmup forwards run
their fused-projection matmuls on the main stream and on a fresh per-layer
side stream allocated inside each ``graph_capture`` (the piecewise layer
capture sweep). cuBLAS seeds one workspace block per ``(device, stream)``
— one for the main stream and one per per-layer side stream — all of
which otherwise live in the synthetic ``unattributed`` residual forever.
That muddies the residual's meaning: a non-zero ``unattributed`` should
mean "a real routing bug", not "library workspaces we expect".

The fix is to allocate those workspaces deliberately into a dedicated
``cublas_workspace`` :class:`NamedMemPool`, via two seams:

  * :func:`warmup_cublas_workspaces` — at boot, before the capture sweeps,
    issue a representative GEMM on the main stream inside the pool's
    ``use()``. Covers the main-stream workspace the decode/prefill
    whole-forward warmups (which run on stream 0) seed.
  * :func:`seed_stream_cublas_workspace` — called by the per-layer capture
    machinery (``runtime/capture/dispatch/capture.py``) once per capture
    side stream, before that stream's warmup forward, with a tiny GEMM
    inside the pool. Covers the per-side-stream workspaces, which are the
    dominant residual on hybrid GDN models.

In both cases the seed is a single bias-free ``torch.matmul``, which backs
the standard cuBLAS handle from the pool; every subsequent GEMM on the same
stream/handle reuses that pool-tagged workspace. Only the standard handle is
seeded — every linear in the LLM forward is bias-free, so production never
lazily allocates a cuBLASLt (``addmm`` epilogue) workspace on these streams,
nothing reuses.

Why this works:

  * The workspace is keyed by ``(device, stream, handle)`` and is
    shape-independent — a small and a large bf16 GEMM pin the same
    workspace block — so one tiny GEMM per stream/handle suffices, and
    keeping the seed GEMM tiny means only the workspace (not the real
    warmup forward's large activations) is routed into the pool.
  * ``NamedMemPool.use()`` no-ops its begin/end-allocate calls while a
    cudagraph capture is active on the stream (to avoid corrupting capture
    state), so the seed must run outside capture. Both seams run before
    their respective ``torch.cuda.graph`` window.
  * cuBLAS honours the active per-thread MemPool: a GEMM run inside
    operand tensors are freed, and ``torch._C._cuda_clearCublasWorkspaces``
    is what frees it — confirming the block is the workspace, not a
    stray operand.

Idempotent: a second call is a cheap no-op GEMM (the workspace is already
cached on the stream/handle, so no new allocation happens).

Set (or clear) the module-global ``cublas_workspace`` pool.

Boot registers the engine's pool; the engine's release registry clears
it (sets ``None``) at shutdown so the pool wrapper can reach refcount 0
and its ``torch.cuda.MemPool`` destructs cleanly.

Force the cuBLAS workspace for ``(device, stream)`` into the pool.

Runs a tiny ``matmul`` on ``stream`` inside the ``cublas_workspace``
pool's ``use()`` context, so the per-``(device, stream, handle)``
workspace cuBLAS lazily allocates lands in the pool. Because the
workspace is shape-independent, a small GEMM is enough — and keeping it
small means only the workspace (not the much larger real warmup-forward
activations) is routed into the pool, so the pool's reserved footprint
stays ~workspace-sized rather than ballooning.

Seeds only the standard cuBLAS handle (``matmul``), not the cuBLASLt
epilogue handle (``addmm`` with a bias): every linear in the LLM
decode/prefill replay path is bias-free (``F.linear`` with no bias →
the standard cuBLAS GEMM), so production never lazily allocates a
cuBLASLt workspace on these side streams. Seeding ``addmm`` here would
inflating the pool's footprint with genuinely-new memory. The operand
is freed before return; the workspace block cuBLAS pins survives (no
Python referrer) and stays tagged to the pool.

Call this once per capture side stream, before the (un-pooled) warmup
forward on that stream: the warmup's first GEMM then reuses the
already-pool-tagged workspace and allocates nothing new. Must run while
capture is not active (``NamedMemPool.use()`` no-ops under capture).

allocator rather than reclaimed per call, since ``empty_cache`` mid
capture-sweep is unsafe and the aggregate across side streams is

No-op when no pool is registered.

Seed the cuBLAS main-stream workspace into ``cublas_workspace``.

Runs one ``matmul`` (standard cuBLAS handle) on the engine device's main
stream, inside the ``cublas_workspace`` pool's ``use()`` context. Must be
called at boot before the cudagraph capture sweeps (``use()`` no-ops
under capture).

Seeds only the standard cuBLAS handle, not the cuBLASLt epilogue handle:
every linear in the LLM forward is bias-free, so the main-stream
decode/prefill warmups never lazily allocate a cuBLASLt workspace — see
:func:`seed_stream_cublas_workspace`.

After the GEMM, ``empty_cache`` returns the seed operand segments to the
driver: the operands have no live Python referrer here, whereas the
workspace block cuBLAS pins does not route through the caching allocator
as a reclaimable segment, so it survives the trim and stays tagged to the
pool. This keeps the pool's *reserved* footprint ~workspace-sized instead
of also holding the (here large) seed-operand segments. Safe because this
runs strictly before any cudagraph capture — never between ``capture_*``
calls, where an ``empty_cache`` could disturb the capture allocator.

No-op when CUDA is unavailable or the pool is not registered. Failures
are best-effort logged — a missed warmup only means the workspace falls
back into ``unattributed``, which is a metrics-attribution regression,
not a correctness bug.

Re-seed the current stream's cuBLAS workspace into ``pool``.

The counterpart to :func:`warmup_cublas_workspaces` for a caller that has
just dropped torch's process-global workspace cache
(``torch._C._cuda_clearCublasWorkspaces``). That drop frees each cached
block back to the pool it was allocated from; running one GEMM inside
``pool.use()`` re-establishes the ``(device, current stream, standard
handle)`` entry in ``pool`` rather than leaving it to be re-allocated
lazily, outside every named pool, by whichever eager GEMM runs first.

Re-seeding inside the pool is what makes the re-allocation map no new
physical. The pool already carries a workspace-sized block for this member
(:func:`warmup_cublas_workspaces` seeded it at the member's build), the
workspace is handle-fixed rather than shape-dependent, and a cuMem pool
keeps its physical mapped across a free — so the re-seed lands in the block
the drop released.

The operands and the output are allocated OUTSIDE the pool and the GEMM
writes through ``out=``, so the only allocation the pool sees is the
workspace itself; a small operand routed into the pool would take the
caching allocator's small-block size class and map a fresh segment under
the tag.

Returns True iff a re-seed ran. No-op (False) without CUDA or a pool.

Per-request detokenizer — inline (default) or async worker.

Two detok strategies share one boundary contract (decode -> update the
engine-side ``req.output_text`` buffer -> scan stop-strings -> emit a
``TokenOut`` carrying the text delta onto the typed output pipeline);
the strategy is selected per process by ``ARBI_INLINE_DETOK``
(:attr:`RuntimeFlags.inline_detok`, default ON):

* **Inline (default, ``ARBI_INLINE_DETOK=1``).** The engine commit loop
  runs ``Tokenizer.incremental_decode`` directly on the engine thread
  (the ``detok_queue is None`` path in :func:`enqueue_token`). The Rust
  ``tokenizers`` ``decode`` does not release the GIL, so a worker hop
  buys no parallelism; the direct call eliminates one
  ``concurrent.futures.Future`` + executor lock acquire/release per token
  per stream. Timing: the decode runs in the engine step's commit
  loop, which (on the async-output deferred path) runs while the NEXT
  step's forward is already in flight on the GPU — so it overlaps the
  forward and finishes before the next graph launch, never delaying it.
  The resulting ``TokenOut`` is staged on the per-step output batch
  (``run_step._emit_token``) and flushed at the GPU-sync window.

* **Async worker (``ARBI_INLINE_DETOK=0``).** vLLM's pattern (cf.
  ``vllm/v1/engine/output_processor.py``): the engine step writes
  ``(token_id, sentinel)`` to a per-request :class:`asyncio.Queue` and
  continues. A dedicated per-request task on the engine loop
  (:func:`_detok_loop`) pulls, runs ``incremental_decode`` on a shared
  :class:`ThreadPoolExecutor` worker thread, updates the engine
  ``output_text`` + scans stop-strings, then emits the ``TokenOut`` (and
  the terminal ``FinishOut`` when the item is a finish) and flushes the
  bus itself — the worker runs between steps, so no step flush would
  cover it. Detok is engine-side in every mode (proc design §4): the
  worker is created from the admission publish callback, which runs on
  the engine loop, so the queue's getter wakeup never crosses threads.
  Per-request ordering is preserved: each worker awaits one decode
  before submitting the next, so token N+1's text lands after token
  N's. Kept as a fallback / A/B comparison toggle.

Synchronous fallback: when no asyncio loop is running at publish time
(test harnesses that build :class:`Request` directly), the queue stays
``None`` and :func:`enqueue_token` runs the decode inline — the same
engine-thread path inline-detok mode uses in production.

Tool-call extraction is not driven here: the OutputApplier feeds the
consumer-side ``ClientRequest.tool_stream`` from the applied text (proc
design §4 — tool parsing consumes decoded text, which already crosses).

Lazy-init the engine's shared detokenizer executor.

Stored as ``eng._detok_executor`` — created on first use to avoid
paying for the threads in test harnesses that never run the async
detokenize loop. Call :func:`shutdown_detok_executor` to release.

Drop the engine's detokenizer executor. Idempotent.

Called from :func:`arbi_serve.engine.lifecycle.shutdown`. ``wait`` is
False by default — engine shutdown wants the run loop to return
promptly; outstanding decodes are cheap enough that the GIL holds
onto them until process exit without observable damage.

Set up the per-request queue + spawn the worker task.

Called from the engine-loop admission publish
(:func:`arbi_serve.engine.request_factory.publish_and_admit`) so the
queue + task bind to the engine loop — the loop the commit path's
``put_nowait`` runs on (detok is engine-side in every mode). When no
asyncio event loop is running (sync test path), this is a no-op —
:func:`enqueue_token` notices ``req.detok_queue is None`` and falls
back to inline decode in the engine step.

Inline-detok mode (``ARBI_INLINE_DETOK=1``, the default): also a
no-op — leave ``detok_queue=None`` so the engine commit loop decodes
inline on the engine thread (the ``detok_queue is None`` path in
:func:`enqueue_token` / :func:`arbi_serve.engine.run_step.post_token`)
instead of dispatching each token to a per-request worker on the
shared :class:`ThreadPoolExecutor`. The Rust ``incremental_decode``
does not release the GIL (see module docstring), so the
worker hop buys no parallelism; the direct call is cheaper than the
``run_in_executor`` ``Future`` round-trip, and the emitted
``TokenOut`` is batched into the per-step output flush at the
GPU-sync window. Set ``ARBI_INLINE_DETOK=0`` to restore the
per-request executor worker (A/B path).

Hand a freshly-sampled token to the request's detokenizer.

Fast path (queue present, worker mode): one ``put_nowait`` and
return — the queue and this call both live on the engine loop, so
the put is loop-local and needs no marshaling. The engine step never
blocks on detokenization.

Fallback (queue absent — the inline default and sync tests): run
:meth:`Tokenizer.incremental_decode` inline on the engine thread and
update the engine ``output_text`` buffer directly. The caller
(``post_token``) emits the resulting ``TokenOut`` after the
stop-scan.

``sync_finished`` is True when the engine step already decided
this request is terminal (stop-token-id / max-tokens / context).
The worker treats the message as the last one for this request:
it decodes, emits the final text + ``FinishOut``, and exits.

The matched stop token is appended to ``output_token_ids`` (for the
token count) but is kept out of the user-visible text by
:meth:`Tokenizer.incremental_decode`, which decodes with
``skip_special_tokens=True`` — so no special-case is needed here.

Synchronous decode used when no detokenizer worker is running.

Mirrors the legacy ``post_token`` body for the decode + stop-string
sub-step. The caller is responsible for the synchronous-by-token-id
finish checks (which still live in :func:`post_token`) and for
emitting the token's ``TokenOut`` afterwards.

This is the engine-thread inline-detok path (``ARBI_INLINE_DETOK``
default-on): decode + stop-scan run here, on the engine thread,
timed to the GPU window. Tool extraction is the OutputApplier's job
(fed from the applied client text) — nothing here touches it.

Scan ``output_text`` for any ``sampling.stop`` substrings.

Returns True iff a stop string fired (request is now FINISHED).
Side effect mirrors the legacy scan: trims the engine ``output_text``
in place, sets ``finish_reason``, calls scheduler.finished +
on_finished. Emission is the caller's job: the inline path
(``post_token``) and the worker loop both emit the (trim-aware)
``TokenOut`` + ``FinishOut`` after this returns, so the consumer
copy replicates the trim via ``TokenOut.trim_to``.

Run the resource-releasing tail of a deferred stop-string finish.

The host-ahead pipeline (ARBI_ASYNC_SCHEDULE) queues this tail from
:func:`_scan_stop_strings_inline` while an ``execute`` is in flight on
the runner thread; :func:`arbi_serve.engine.run_step._flush_deferred_finishes`
executes it once the runner returned. Byte-for-byte the tail the scan
runs inline when no pipeline overlap is active — ``scheduler.finished``
(page/slot free + radix commit), the FINISHED state flip, the LoRA ref
release, the live-map drop, and the finish metrics — plus the
idempotent ``cancel_detok_task`` the inline caller's ``post_token``
tail would otherwise have run on the state check. The trim-aware
``TokenOut`` was already emitted at scan time by the scan's caller;
the terminal ``FinishOut`` is emitted here (idempotent via
``req._finish_emitted`` — a detok worker that already emitted it on
its ``stopped`` return makes this a no-op) and rides the next step
flush. Idempotent via the FINISHED short-circuit.

Per-request detokenizer worker (engine loop; worker mode only).

Consumes ``(token_id, finished)`` tuples from ``req.detok_queue``,
runs :meth:`Tokenizer.incremental_decode` on a worker thread,
updates the engine ``req.output_text``, scans stop strings, then
emits the token's ``TokenOut`` (+ the terminal ``FinishOut`` on a
finish) and flushes the output bus — the worker runs between engine
steps, so no step flush point covers its emissions. The flush is
per-processed-item: worker mode trades the one-hop-per-step
invariant for taking the decode off the step body (it is the
non-default A/B path).

Exits when:
  - a queue item carries ``finished=True`` (sync-finished path)
  - a stop-string scan fires
  - the worker's task is cancelled (engine shutdown / request
    cancel / timeout — the canceller emits the ``FinishOut``)

Thread-side body. Runs on the executor — releases the GIL.

Takes a *snapshot copy* of ``output_token_ids`` (the engine task
keeps appending to the live list); ``detok_state`` is mutated
in-place by ``incremental_decode`` and is only ever touched by
this worker, so no lock is required.

Cancel the per-request detokenizer task. Idempotent.

Called from the engine cancel / timeout / finish paths so the
background task drops its references and exits promptly. The task
has no GPU state — cancellation is a clean asyncio operation.

Per-layer dequant relative-error report for TKV backends.

Drives the ``/v1/admin/distortion`` admin endpoint: runs every
registered backend's :class:`TkvCodec` over either the captured
per-layer ``(K, V)`` tensors (when the model exposes ``capture_kv``)
or a synthetic Gaussian sample, and reports relative L2 error per
``(layer, side)``.

Run the eval-set prompts through every TKV backend and report
per-layer dequant relative error.

Algorithm (per backend):
  1. Tokenize the prompts (truncated to ``max_tokens_per_prompt``).
  2. Run a forward pass capturing each layer's pre-codec
     ``(K, V)`` tensors.
  3. For each layer, encode + decode through that backend's
     :class:`TkvCodec` and compute relative L2 error per side.

Returns a dict shaped::

    {
        "max_tokens_per_prompt": 64,
        "num_prompts": 8,
        "backends": [
            {
                "backend": "tkv-k4v4",
                "per_layer": [{"layer": 0, "k_err": 0.092, "v_err": 0.084}, ...],
            },
            ...,
        ],
    }

For non-TKV backends, ``per_layer`` is ``[]`` (no codec means no
dequant error to report).

Synchronous blocking call — admin endpoints wrap it in
:func:`asyncio.to_thread`.

Per-layer ``(k_err, v_err)`` using a captured-K/V forward pass.

Hooks every PAGED_KV layer's QKV projection output, encodes and
decodes the K and V tensors through that layer's codec ops, and
reports relative L2.

Encode + decode → relative L2 error per side.

Uses the codec's compress / decompress public surface; falls back
to zeros if the codec API doesn't expose the round-trip on this
build. Returns ``(k_rel, v_rel)`` — relative L2 norm of the
residual against the original.

Distributed engine driver — multi-rank topology under torchrun.

Two execution paths share this driver:

  - SPMD (``spmd_tp``, the default): every rank derives the slate and
    forwards symmetrically; there is no per-forward control op or
    materialized-batch broadcast. The per-rank loop lives in
    :mod:`arbi_serve.distributed.spmd` and is driven by
    :meth:`_run_loop_spmd` (built in :meth:`_build_spmd_loop`). The
    worker bridge is not installed in this mode.
  - Driver / worker bridge (``ARBI_SPMD_TP=0``, the fallback): rank 0
    owns the scheduler and broadcasts each forward / drafter / seed op
    to forward-only worker ranks via the control ring. The rest of this
    docstring describes this path.

Single-node multi-GPU layout (driver / worker bridge path):

  - Every rank loads its TP-sharded slice of the model (parallel
    linears + per-rank KV slabs) inside :class:`Engine.build`. The
    pool sizes ``num_kv_heads`` per rank as
    ``num_kv_heads // tp_size``.
  - Rank 0 owns the scheduler, sampler, tokenizer, page table,
    request map, FastAPI surface. Ranks 1..N-1 hold none of these —
    they are forward-only workers.
  - Per slate: rank 0 runs the scheduler then enters its standard
    spec-decode strategy (legacy K=1 ``step`` and / or MTP verify
    pass). Each forward / drafter-chain / seed-fill issued by the
    rank-0 engine code paths enqueues a control op via the
    :class:`arbi_serve.distributed.worker_bridge.WorkerBridge` before
    running locally; worker ranks sit in
    :meth:`WorkerBridge.run_worker_dispatch` and execute each
    received op against their own engine in lockstep. After the
    slate's last op, rank 0 enqueues a ``"tick_done"`` op and
    workers return to their outer recv loop.
  - Control messages ride a CPU shared-memory ring buffer
    (:class:`arbi_serve.distributed.shm_message_queue.ShmControlQueue`):
    rank 0 writes, every worker rank reads the same broadcast slot.
    NCCL is reserved for tensor
    collectives only — per-layer activations cross the TP boundary
    inside ``RowParallelLinear.all_reduce`` and the MTP head's
    ``lm_head.all_gather``; no control message touches NCCL.
  - Hot-swap (per-:class:`StateKind`) is shipped as a control op;
    every rank rebuilds its per-rank pool view + per-layer attn ops
    in lockstep.

Why broadcast the materialized batch instead of the slate?

The earlier design assumed every rank could re-derive identical KV
slot indices from a deterministic page-table replay. That contract
breaks against features that already exist on rank 0:

  - :class:`RadixPageTable` aliases pages on prefix-cache hits, so
    the slot-allocation order on rank 0 depends on the radix-tree
    state — a worker rank with no request stream cannot mirror it.
  - LoRA per-row assignments need to land in the
    :class:`ScheduledBatch` shape; mirroring requires the
    ``lora_id`` per row.
  - Tenant scoping and per-request ``cache_enabled`` change the
    radix-cache lookup; mirroring requires both.

Broadcasting the post-build batch tensors instead removes the
mirroring problem entirely. Workers do not run a page table or a
scheduler; they consume the batch shape rank 0 already produced.
The protocol composes with whatever rank-0 features land later.

MTP at TP>1 (Pattern 1: rank-0-only MTP). The verify pass + drafter
chain run on rank 0's scheduler / page table / sampler, but the
MTP head's :class:`RowParallelLinear.o_proj` issues an ``all_reduce``
collective every rank in the TP group must participate in. The
worker bridge handles this: rank 0's ``MtpDriver.draft`` broadcasts
``(B, K, hidden_in, last_token_in, sampling_params)`` to workers,
which reconstruct + call ``draft`` against their local engine in
lockstep; tokens / probs are discarded on workers. The ``mtp_seed``
counter is broadcast so all ranks observe identical Gumbel-max
draws.

The constructor refuses ``world_size <= 1`` (single-GPU must use
:class:`Engine` directly) and refuses LoRA at TP>1 — it composes on
top of the dense Qwen3 path but is not supported under TP.
MTP-at-TP>1 is fully wired.

Multi-process engine driver for ``world_size > 1``.

Each rank owns one :class:`Engine` instance — same code path as
single-process — but rank 0 is the only one that admits requests,
runs the scheduler, runs the page table, and samples. Worker ranks
receive a :class:`StepPlan` per step, rebuild the batch, and run
forward only.

Control messages travel over a CPU shared-memory ring buffer
(:class:`arbi_serve.distributed.shm_message_queue.ShmControlQueue`,
built in :meth:`build`). NCCL — initialized by every rank via
:func:`init_distributed_environment` inside :class:`Engine.build` —
carries only the per-layer tensor collectives.

Forward unknown attributes to the rank-local :class:`Engine`.

The OpenAI-compatible HTTP layer reads engine attributes
(``cfg``, ``tokenizer``, ``metrics``, ``requests``, ``pool``,
``page_table``, ``scheduler``, ``model_runner``, …) directly off
``app.state.engine``. Wiring the driver as ``app.state.engine``
and forwarding every read keeps the server code path uniform —
the driver only adds explicit overrides for the methods that
carry a control-plane broadcast (e.g.
:meth:`aswap_attention_backend`).

Standard ``__getattr__`` semantics: only called when the
attribute isn't found on the driver itself, so the driver's
own methods (``submit``, ``cancel``, ``shutdown``,
``aswap_attention_backend``, ``run_forever``) win without
recursion.

Bring up the per-rank engine and the shm control queue.

Runs on every rank. After the engine is up, both ranks build the
CPU shared-memory control queue and hand-shake its handle once
over the existing NCCL TP group (the only remaining control-plane
NCCL use — it runs at boot, not per step). If shm setup fails the
exception propagates and boot aborts.

Engage per-rank stable-VA residency after the engine builds.

Runs on every rank — single-member engagement is local bookkeeping
(wrap the booted model as the active resident record; no collectives,
no new CUDA allocations), so each rank engages its own controller with
its own device index. This is what makes the hot-swap surface
(``/v1/admin/attention_backend`` / ``/v1/admin/config_override``) work
at TP>1; without it the admin paths raise "requires stable-VA residency
engaged".

With boot-declared ``--pool-member`` models this instead drives the
multi-member pool orchestration per rank (see :meth:`_engage_model_pool`),
which DOES build + capture each extra member and so runs NCCL
collectives — hence it is barrier-gated across ranks.

No-op when ``cfg.stable_va_residency`` is off (--no-cumem-pools).
Idempotent for the single-member case — :meth:`Engine.
enable_stable_va_residency` returns the existing controller, so a
later lifespan engagement (the single-process path never runs for a
prebuilt driver, but defense-in-depth) cannot double-engage.

Engage a boot-declared multi-model pool per rank, barrier-gated.

Mirrors the single-process ``prepare_model_pool`` boot orchestration
(:func:`arbi_serve.engine.stable_va_pool_builder.prepare_model_pool`)
but runs symmetrically on EVERY rank: each rank builds its own TP shard
of every member (``build_member_into_engine`` calls ``eng.build`` under
the member's namespace, which naturally shards per rank). The boot model
(``members[0]``) is already built by ``engine.build`` above, so
``build_primary=False``; the loop builds + parks each extra member and
leaves the boot model the active resident.

Each member's commit is gated through
:func:`arbi_serve.engine.member_build_barrier.confirm_build_across_ranks`
so an asymmetric per-rank build failure (one rank OOMs) aborts that
member on every rank together — never a partial commit that would
desync the cluster into a NCCL-watchdog SIGABRT. The member list +
overlays derive purely from ``cfg`` (identical on every rank), so the
ranks build the same members in the same order and stay in lockstep.

This is the ONLY pool engagement at TP>1: the rank-0 lifespan
(``_attach_prebuilt_driver`` / the process-mode ``build_engine``)
deliberately does NOT call ``_engage_stable_va_residency``, so the
orchestration runs exactly once.

Create + hand-shake the shm control queue across all ranks.

Rank 0 (writer) creates a single 1-writer / (world_size-1)-reader
broadcast queue and ships its :class:`Handle` once over the TP
group; every worker rank (1 .. world_size-1) attaches from that
handle as one of the readers. All ranks then rendezvous on
``wait_until_ready`` to defeat the zmq slow-joiner.

The queue is N-general: rank 0 enqueues each control op once and
the ring fans it out to every worker (each worker has its own
per-reader flag byte + cursor over the shared slots). TP=2 is just
the single-reader case. Scoped to a single TP group (ep_size=1 →
world_size == tp_size); EP composition adds per-group control
planes.

Signal every rank to exit the run loop.

Rank 0 calls into here from the FastAPI lifespan teardown; the
next :meth:`run_forever` tick broadcasts a ``"shutdown"`` control
op so workers also exit.

Per-rank step loop.

Rank 0 — schedule a slate, build batch, enqueue plan, run
forward+sample, commit. Worker ranks — dequeue plan, rebuild
batch, run forward, discard logits. The shm queue's FIFO order
keeps both paths in lockstep.

Attention-DP serving: per-set admission and per-rank step plans.

At ``attn_dp_size > 1`` each attention-DP set owns whole REQUESTS — their
tokens and their whole KV — and forwards only its own rows through
attention, while the FFN stays parallel over the entire TP group. This
module is the scheduler / SPMD-loop half of that: the model-side
gather/slice boundary lives in
:mod:`arbi_serve.distributed.dp_attention`.

THE ONE INVARIANT EVERYTHING RESTS ON — *rank 0 is the sole decider of
per-tick control flow, and every rank derives its control flow from the
broadcast message alone, never from its own row count.*

Concretely: rank 0 schedules every set, packs the sets' deltas into ONE
:class:`~arbi_serve.distributed.spmd.DpSlateDeltas`, and broadcasts it.
Every rank then reads the SAME object, so every rank reaches the same
verdict on the two questions that must not diverge:

  1. **Does this tick forward at all?** ``DpSlateDeltas.forwards`` — true
     when ANY set has rows. The FFN exit reduces over the whole TP group,
     so a rank that skipped the forward because its OWN set was empty
     would leave every peer blocked inside that reduce. A set with no rows
     therefore still forwards, contributing zero rows to the gather.
  2. **How many forwards does it issue?** Exactly one. The MTP verify
     tick issues a verify pass plus a legacy pass whose split is a
     per-ROW property, so at ``attn_dp_size > 1`` the sets could disagree
     on whether the second pass runs — a collective count mismatch, not a
     wrong number. :func:`attn_dp_boot_gate` refuses MTP for that reason,
     and :func:`_agree_spec_mode` fails loud if a tick ever reaches here
     with the sets disagreeing.

The existing cross-rank step-outcome reduce
(``_spmd_agree_step_failed``, over the WORLD group) still turns an
asymmetric step FAILURE into a symmetric refuse; this invariant is about
the success path, which that reduce cannot cover.

WHAT RANK 0 HOLDS. Only rank 0 runs schedulers, so it runs one per set
(:class:`~arbi_serve.scheduler.attn_dp_group.AttnDpSchedulerGroup`) and
one :class:`~arbi_serve.distributed.spmd.SpmdRankLoop` per PEER set
besides its own. A peer loop is the same object a worker rank runs, fed
the same delta, so rank 0's mirror of that set's page-table occupancy
advances exactly as the set's own does — which is what makes the peer
scheduler's admission decisions match the pool the set actually has.

HOW THE TOKENS COME BACK. Each set samples its own rows, but only rank 0
holds the :class:`~arbi_serve.engine.request.Request` objects and the
output streams. So after sampling, every rank all-gathers its set's
sampled tokens over the attention-DP group — one ``(max_rows,)`` int64
tensor per rank, padded to a width every rank reads off the same
broadcast message. Rank 0 then commits and streams every set's rows.

Refuse every configuration attention-DP serving is not wired for.

Runs on EVERY rank, from the SPMD loop build, so a refusal is
symmetric. Each clause names a mechanism that would be silently wrong
(not merely unsupported) at ``attn_dp_size > 1``:

* **Non-SPMD driver path.** The driver/worker bridge ships one
  materialized batch to every worker, which is the replication this
  axis removes.
* **Radix page table.** A prefix tree per set is only worth having
  with prefix-affinity routing (a request must land on the set whose
  tree already holds its prefix); without it the sets thrash
  independent trees and rank 0's shadow tables would have to mirror
  per-set eviction plans that no set can act on.
* **Recurrent state.** The per-request recurrent slab is allocated by
  the scheduler; a peer set's slab lives on another process and rank
  0's shadow pool deliberately has no rows to hand out.
* **MTP / spec decode.** See invariant (2) above — the verify tick's
  sub-pass count is a per-row property, so the sets can disagree on
  how many forwards the tick issues.
* **CUDA graph capture / torch.compile.** The attention-DP gather's
  width is the widest set's row count for THIS step — a runtime
  value. A captured graph would replay a stale width, and a compiled
  region would recompile or graph-break on every distinct width.
* **A rank-0 scheduler that is not the per-set group.** Without it
  every set would be admitting against rank 0's own pool.

Reads ``driver._attn_dp_size`` (derived once in ``__init__``) rather
than the parallel config, so the class default of 1 makes this inert
on a driver shell that never ran ``__init__``.

Fail loud when the sets disagree on this tick's spec mode / K.

Invariant (2) in the module docstring: the number of forwards a tick
issues must be a property of the tick, not of a set's rows. The boot
gate refuses MTP outright, so today this can only fire if that gate is
ever relaxed without the sub-pass count being made global first —
which is a deadlock, not a wrong answer, and is worth one comparison
per tick to turn into an exception.

Bind rank 0's per-set trackers and peer page-table mirrors.

Called from the SPMD loop build, after :func:`attn_dp_boot_gate`
has passed, so ``eng.scheduler`` is known to be the per-set group.

``drive_page_table=False`` on every peer loop, for the same reason
rank 0 passes it for its own set: that set's scheduler — which
rank 0 also runs — has already applied this tick's admits and
frees to the very same table. Driving them again from the loop
would double-allocate out of the shadow pool and make the mirror
of the set's occupancy pessimistic, which shows up as that set
refusing admissions it has room for. The per-step
``allocate_slots`` inside ``derive`` runs either way.

Schedule every set and pack one tick's deltas for the wire.

Returns the broadcast payload and rank 0's per-set
``[(Request, n_tokens)]`` slates (the objects only rank 0 holds,
used to stream this tick's output).

Advance rank 0's peer page-table mirrors to this tick's state.

Runs the SAME two steps, in the same order, that the peer set's
own ranks run before their forward: apply the delta (frees, then
admits), then the per-step ``allocate_slots`` inside ``derive``.
The derived tensors are discarded — rank 0 never forwards a peer's
rows — but the allocation they drive is exactly what keeps the
shadow table's free count equal to the set's real one, which is
what that set's scheduler admits against.

Land an idle tick's frees on every set rank 0 tracks.

A set can finish a request on a tick whose next schedule is empty;
without this its pages would stay held on rank 0's mirror of that
set until some later busy tick, and the set's own scheduler would
keep admitting against the stale count.

Every set's rows for a tick the ranks agreed to fail.

The step-outcome reduce fails the WHOLE tick on every rank, so
rank 0 must release the clients of every set's rows — releasing
only its own set's would leave the peers' requests hanging on a
step no rank will retry.

All-gather every set's sampled tokens over the attention-DP group.

The buffer width is the widest set's row count for THIS tick, read
off the broadcast ``dp`` — so every rank sizes it identically with
no extra collective and no host sync to agree on a shape. Rows are
padded past the set's own count and the pads are dropped on the way
out, so the result is a pure re-arrangement.

Issued by every rank: the attention-DP group holds one rank per
set (those sharing an ``attn_tp_rank``), and a set's ranks all
sampled the identical tokens, so every group's gather returns the
same per-set lists.

Advance every peer set's mirror, then commit + stream the tick.

Rank 0's OWN set has already advanced its mirror in the step body
(the same call every rank makes); the peers advance here off the
exchanged tokens. The scheduler-side commit and the per-token
stream then run ONCE over every set's rows concatenated — the
per-set group routes each row to its owning member — so the finish
/ detok / SSE path is the single-GPU path unchanged.

Rank-0-driver / worker-bridge run loops + admin broadcast surface for
:class:`~arbi_serve.engine.distributed_driver.DistributedEngineDriver`.

Holds the non-SPMD (``ARBI_SPMD_TP=0``)
run loops and the swap / config-override / profiler admin broadcasts.
Methods read ``self`` attributes set in
``DistributedEngineDriver.__init__``; the shared ones are declared as
class-level annotations below (duck-typed at runtime).

Swap the backend for one StateKind across every rank.

Rank 0 broadcasts a ``"swap"`` control op; every rank then runs
the same single-process swap path on its local engine. The
per-rank pool view is rebuilt with the per-rank KV-head count
(``num_kv_heads // tp_size``) so post-swap geometry is coherent
without an extra collective.

Workers never call this directly — :meth:`run_forever` handles
the worker-side rebuild on receiving the broadcast.

Apply a live config override on every rank.

Mirrors :meth:`aswap_attention_backend`: rank 0 broadcasts a
``"config_override"`` control op, then runs the same local
apply path (:func:`arbi_serve.engine.config_variant.aapply_overrides`).
Workers apply the identical delta in :meth:`run_forever` so the
variant route (noop / live overlay / instant switch / build) — and
any per-rank park/wake/build it implies — is computed in lockstep
from the same broadcast payload. The build-required route runs the
same park-build-capture sequence on every rank; a cross-rank
commit/abort barrier inside ``config_variant._aprepare_member``
makes an asymmetric build outcome (one rank OOMs) roll back on all
ranks. A build that fails (locally or on a peer) raises here, so
the admin response reflects the cluster outcome, not just rank 0.

Ephemerally drop ONE parked config variant on every rank.

Mirrors :meth:`aapply_config_overrides`: rank 0 broadcasts a
``"config_drop_variant"`` control op (carrying the variant key), then
runs the same local teardown. Workers drop the identical record + reclaim
its namespace physical + VA arena in lockstep, so the residency set stays
symmetric across ranks (a per-rank drop that ran on only some ranks would
desync the pool).

Drop EVERY non-active config variant on every rank (broadcast reset).

The ephemeral-A/B reset across the whole cluster: rank 0 broadcasts a
``"config_drop_variant"`` op with ``key=None`` and every rank returns its
residency to the single active baseline member in lockstep.

Switch the resident model on EVERY rank (fast park/wake, broadcast).

Mirrors :meth:`aapply_config_overrides`: rank 0 broadcasts a
``"switch_model"`` control op, then runs the same local park/wake
(:meth:`Engine.aswitch_model`). Workers apply the identical switch in
:meth:`_apply_worker_control_op`, so a request-triggered residency swap
(driven from the run loop when the active scheduler drains) is
rank-lockstep. The recapture-free park/wake remaps each rank's own
physical from host at its stable VAs — no cross-rank collective — so a
broadcast + symmetric local apply is coherent (same contract as the
``swap`` / ``config_override`` worker handlers).

Switch-to-or-load the named model on EVERY rank (broadcast).

Mirrors :meth:`aapply_config_overrides`: rank 0 broadcasts a
``"switch_or_load_model"`` control op carrying the full request, then
runs the same local :meth:`Engine.aswitch_or_load_model`. Workers apply
the identical op in :meth:`_apply_worker_control_op`, so the present-vs-
absent routing AND (on absent) the runtime member add — its lazy
residency engage, park/build/capture, and cross-rank commit/abort barrier
— run in lockstep from the same payload. A missing model dir / refused
build / cluster abort raises identically on every rank; rank 0's raise
surfaces to the admin caller, the workers' matches (loop stays alive).

Start / stop the live torch profiler on every rank.

Rank 0 broadcasts a ``"profile"`` control op so each worker runs
the same :func:`start_live_torch_profile` /
:func:`stop_live_torch_profile` against its local engine — so each
rank dumps its own per-rank Kineto trace. ``action`` is ``"start"``
or ``"stop"``. Returns the rank-0 trace dir (workers' dirs match).

Deliberately not drain-gated — the whole point is profiling live
decode traffic — so this broadcast can interleave with an active
slate's per-slate bridge tuples in legacy driver mode (the admin
HTTP handler and the step loop share one event loop). The worker
dispatch tolerates that: it defers the op to the tick boundary
(see :meth:`arbi_serve.distributed.worker_bridge.WorkerRankBridge.run_worker_dispatch`).

Rank-0 loop: schedule → broadcast slate-start → strategy.run_step → tick_done.

The per-slate forward / drafter / seed broadcasts ride on the
:class:`arbi_serve.distributed.worker_bridge.RankZeroBridge`
installed on the engine: every rank-0 forward / drafter call
emits a worker-bridge op before running, so workers
reconstruct + run the same op against their local engine in
lockstep. After the strategy finishes, this loop emits a
``"tick_done"`` op that returns workers to the outer recv
loop.

Worker loop: outer recv → if slate, drain bridge → repeat.

Outer ops (``shutdown`` / ``noop`` / ``swap``) handle long-
lived rank-0 actions. ``slate`` enters the per-slate worker-
bridge dispatch which consumes batch / drafter / seed ops
until rank 0 broadcasts the slate's ``tick_done``.

Apply a broadcast admin op on a worker rank.

Returns True if ``ctl`` was an admin op this method handled
(``swap`` / ``config_override`` / ``profile``), False otherwise so
the caller can route the remaining ops (``slate`` / unknown).

The single worker-side hot-swap dispatch surface: both the legacy
worker loop and the SPMD worker loop call this so a backend swap /
config override / profile toggle is applied identically on every
rank. Rank 0 already broadcast the op (``aswap_attention_backend``
/ ``aapply_config_overrides`` / ``profile_action``) before its own
local apply; the worker mirrors it here so the per-rank pools +
captured graphs are remapped in lockstep (the build-required route
rendezvouses on the cross-rank commit/abort barrier inside
``config_variant._aprepare_member``). Admin ops are dequeued
between ticks (the worker holds no request map and is idle), so the
local apply — including a park/build/capture — never races a
half-applied forward.

Rank-symmetric SPMD (``ARBI_SPMD_TP``) run loop + verify orchestration
for :class:`~arbi_serve.engine.distributed_driver.DistributedEngineDriver`.

Holds the SPMD build hook, the deferred-release drain, the
per-tick loop, and the verify / commit / stream sub-passes. The per-step
derivation + sample + agreement helpers live in
:mod:`arbi_serve.engine.distributed_driver_spmd_step`. Methods read
``self`` attributes set in ``DistributedEngineDriver.__init__``.

Build the per-rank SPMD derivation core (ARBI_SPMD_TP only).

Every rank binds a :class:`SpmdRankLoop` to its rank-local page
table; rank 0 additionally owns the :class:`RankZeroSlateTracker`
that turns each scheduler tick into a :class:`SlateDelta`. The
derivation core is engine-free and CPU-tested
(``tests/test_spmd_derivation*.py``); this hook only binds it to
the live page table so :meth:`_run_loop_spmd` can drive it.

Page-table support: SPMD mirrors the FlatPageTable alloc/free
order across ranks directly. The radix prefix cache is also
supported — the worker can't re-run the LRU ``match_prefix`` (the
match is the one non-derivable decision), so rank 0's
``prefix_match_len`` + matched page ids ride the
:class:`AdmitRow`; the worker aliases the same pages and mirrors
``commit_full_pages`` / ``remove_request`` in lockstep (see
:meth:`RadixPageTable.add_request_spmd`). Radix LRU eviction under
pool pressure is supported: rank 0 plans the victim pages up front
(:meth:`RadixPageTable.plan_evictions_spmd`) and the ids ride
``SlateDelta.evict_pages``; workers replay the same frees
(:meth:`RadixPageTable.evict_spmd`) before deriving. Scheduler
preemption is likewise broadcast (``SlateDelta.preempt``) and
mirrored. Any other page-table impl is refused (see AGENTS.md).

Queue a finished/cancelled request for the next-tick release drain.

Installed as the scheduler's ``set_spmd_deferred_release`` hook
(rank 0, radix). Deduped — ``finished`` + a later ``remove`` for the
same request release once.

Release the queued finished requests (rank 0, top of tick).

Runs, per rid in finish order: the mirror-sourced finish-commit
(:meth:`SpmdRankLoop.commit_full_pages_on_finish`), the page-table
``remove_request`` free, and the recurrent-slab free — the same
three ops, in the same relative position (before this tick's
preempts, LRU-eviction plan, and per-step allocations), as the
worker's ``apply_delta`` evict replay. Returns the drained rids;
they feed ``delta.evict`` so the worker replays them this tick.

Straggler net: a request that left ``eng.requests`` without passing
through the scheduler's hooked finish paths would leak on the
workers (never ride ``delta.evict``) and pin rank 0 pages forever —
release it here too, loudly.

Drain the deferred finish-release queue while the outgoing member
is still bound (``Engine.register_pre_model_switch_hook``).

The queue's rids index the outgoing member's page table and recurrent
pool — and both ride the parked record and come back on a later wake,
so a rid left queued is not stale bookkeeping, it is a page + recurrent
row leaked across the swap and never freed on either rank.

The drain cannot be deferred to the next tick's normal drain for two
independent reasons, which is why this hook exists at all:

  * By then ``eng.page_table`` / ``eng.pool`` point at the incoming
    member, so :meth:`_spmd_drain_pending_release`'s
    ``remove_request`` would hit the wrong table (suppressed
    ``KeyError`` — a silent no-op) and free a recurrent row that
    belongs to a different member's slab.
  * The re-bind at the top of the tick rebuilds the loop, which resets
    ``_spmd_pending_release`` to empty — the queued rids would simply
    vanish before anything looked at them.

Ordering across ranks is the other half. The drained rids are the
source of the workers' ``delta.evict``; the switch's ``config_override``
control op is broadcast during the switch itself. Both ride the same
FIFO control queue, so broadcasting the evict delta here — before the
switch op is enqueued — is what guarantees a worker replays these frees
against the old table it still has bound, then re-binds. Draining after
the switch would invert that order and diverge the ranks.

Rank 0 only: workers own no scheduler and no release queue (their frees
arrive as ``delta.evict``). A no-op when nothing is queued.

Rank-symmetric SPMD step loop (ARBI_SPMD_TP).

Replaces the per-step driver/worker materialized-batch broadcast
with the model from :mod:`arbi_serve.distributed.spmd`:

  1. Rank 0 runs the scheduler, builds a compact
     :class:`arbi_serve.distributed.spmd.SlateDelta` from the
     tick's decisions, and broadcasts it once over the shm
     control queue. Workers dequeue it.
  2. Every rank applies the delta to its
     :class:`arbi_serve.distributed.spmd.SpmdRankLoop` (mirror +
     rank-local page table) and derives the identical per-step
     batch tensors — no per-step batch crosses the wire (the
     contract proven byte-for-byte by the 2-rank gloo test).
  3. Every rank runs the same forward + sample from the derived
     batch (greedy = argmax over the post-lm_head-all_gather
     logits, which is bit-identical on every rank); only rank 0
     commits to the request map + streams.

Plain decode (``spec_mode != MTP`` or ``step_K < 2``) is fully
wired here. The MTP verify+accept+draft path is wired via
:meth:`_spmd_run_verify`: every rank runs the same verify forward
+ greedy accept + drafter chain off the replicated mirror
(rank-symmetric collectives, no token broadcast); only rank 0
streams. Plain-decode and MTP-verify stochastic sampling are both
rank-symmetric (the draws are derived from rank-agreed seeds —
``gumbel_exp_noise`` for plain decode, the lockstep ``mtp_seed``
counter for the rejection sampler + drafter chain). A mixed verify
tick (prefill / non-MTP rows alongside the MTP rows) is wired via
:meth:`_spmd_run_verify`'s split: the verify forward for the MTP
rows then a second rank-symmetric plain forward for the legacy
rows, paired 1:1 across ranks. Scheduler preemption (the
:class:`Preemptor` bumps a running victim to waiting + resets its
decode state) rides ``SlateDelta.preempt``; radix LRU eviction under
pool pressure rides ``SlateDelta.evict_pages``. Both are
rank-0-decided and mirrored deterministically — no per-rank
nondeterminism, no refusal.

Duplex rows never reach this loop: their tick needs the duplex
lane's ``prepare_tick``/``complete_tick`` hooks (which only the
single-rank ``_run_forever_inner`` calls), a per-tick
``pending_embed_override`` (which the replicated mirror this loop
derives from cannot carry), and ``retain_logits_heads`` (which
this loop's forward entry does not thread the per-row requests
for). ``admit_duplex_request`` refuses the combination — see
:func:`~arbi_serve.realtime.nemotron_voicechat_duplex_admission
.duplex_multi_rank_refusal`.

Run one SPMD step body (every rank). ``slate`` is the rank-0
scheduler slate ``[(Request, n)]`` on rank 0, ``None`` on workers.

Wrapped by :meth:`_run_loop_spmd` in a per-step try/except so a
recoverable failure refuses one slate without killing the loop.

Rank-0-only: commit the sampled tokens + run the finish/stream path.

Reuses the engine's existing post-step path
(``scheduler.commit`` + ``post_token``) so finish-reason handling,
detok hand-off, and streaming are byte-identical to the
single-GPU engine. Workers never reach here (they discard their
identical sample).

Output flush: ``post_token``'s per-token structs (TokenOut /
FinishOut) are staged onto the engine's per-step output batch,
not delivered — exactly like the single-process commit loop
(``_commit_step_result``). Fire the single ``_flush_output`` at
the end of the sub-pass so the batch reaches the consumer this
step. Without the flush the SPMD loop never delivers it (it has
none of ``run_step``'s flush sites), so a finished request's
response would sit staged until the next engine step flushed it —
the inter-request turnaround stall. Immediate-flush-per-subpass
is the SPMD flush-placement invariant.

Split an MTP-verify tick into (verify, legacy) sub-slates.

Mirrors the driver's :meth:`MtpStrategy.run_step` partition
(``mtp_strategy.py``): a row is an MTP-verify row iff it opted into
MTP (``mtp_k > 0``) AND is past prefill; every other row (prefill,
or a non-MTP decode) is a legacy row run via a plain forward.

Derived purely from the replicated mirror's per-row
:class:`SamplingDigest`, so every rank computes the identical split
— the two sub-passes below then pair 1:1 across ranks with no NCCL
desync. Slate order is preserved within each sub-slate.

Run one MTP verify step rank-symmetrically (mixed slate OK).

``slate`` is the rank-0 scheduler slate ``[(Request, n)]`` on rank
0, ``None`` on workers. ``delta`` carries the slate ids + uniform
``step_K``.

A verify tick can be mixed: MTP-opted decode rows alongside
prefill / non-MTP rows the scheduler bucketed onto the same tick
(new requests arriving while others MTP-decode). This is the common
shape at any real concurrency. The tick splits the same way the
driver's :meth:`MtpStrategy.run_step` does (see
:meth:`_spmd_split_mixed_slate`):

  1. Verify sub-pass — the MTP-opted decode rows run the same
     rank-symmetric verify core
     (:func:`run_verify_step_spmd_worker`): one verify forward
     (row-parallel all-reduce + lm_head all_gather) + greedy/sto
     accept over the post-all_gather logits (deterministic →
     identical tokens on every rank, no broadcast) + the K-step
     drafter chain.
  2. Legacy sub-pass — the prefill / non-MTP rows run a second
     rank-symmetric plain forward this tick (the exact
     :meth:`_spmd_forward_and_sample` path the SPMD plain-decode
     branch uses): every rank derives the same legacy sub-batch
     from the broadcast delta and runs the identical forward +
     sample (no token broadcast, rank-symmetric collectives).

Both sub-passes run in the same order on every rank (verify, then
legacy), so their TP collectives pair peer-for-peer and the NCCL
group never deadlocks. The verify core is driven off the rank-local
:class:`RankSlateMirror`, so the drafter-chain batch shape is
byte-identical on every rank (it never drops finished rows mid-step
— finishes are a rank-0 streaming decision that takes effect on the
next tick via the delta's evict, so dropping them here would desync
the collective shapes and deadlock the group).

Only rank 0 then streams the verify accepted tokens / commits the
legacy sampled tokens onto its real :class:`Request` objects
(``post_token`` — detok / finish / SSE). Both ranks advance the
mirror by the identical results.

Rank-0 helper: pull the ``(Request, n)`` rows for ``sub_ids``.

Preserves ``sub_ids`` order so the rank-0 stream/commit pairs 1:1
with the per-row results computed off the same id order. Rows not
on the live rank-0 slate (already-finished) are dropped.

Rank-0-only: stream a verify step's accepted tokens.

Feeds each row's accepted tokens (computed identically on every
rank by the verify core) into the engine's existing per-token
finish/stream path (``post_token``) on the real :class:`Request`.
Workers never reach here — they discard their identical accept
result after the mirror advance.

Output flush: same staging contract as
:meth:`_spmd_commit_rank_zero` — ``post_token`` stages typed
structs onto the engine's per-step output batch; fire the single
``_flush_output`` at the end of the sub-pass to deliver this
verify step's accepted tokens to the consumer now (else the
response stalls until the next engine step flushes).

Per-step SPMD derivation, forward+sample, and cross-rank agreement
helpers for
:class:`~arbi_serve.engine.distributed_driver.DistributedEngineDriver`.

Holds the rank-symmetric per-step machinery driven by
:meth:`arbi_serve.engine.distributed_driver_spmd._SpmdLoopMixin._spmd_step`:
batch derivation, forward+sample, MTP seed/draft, grammar advance, the
finished/preempt/evict planners, and the C9 step-outcome reduce. Methods
read ``self`` attributes set in ``DistributedEngineDriver.__init__``.

A peer never joined the C9 agreement reduce — it is wedged mid-step.

Raised by :meth:`_SpmdStepMixin._spmd_agree_step_failed` when THIS rank
failed its step and the peer did not arrive at the agreement within
``spmd_agree_timeout_s``. That combination has exactly one meaning: this
rank raised BETWEEN two collectives, so it skipped one the peer had
already enqueued, and the peer is now blocked inside it.

The peer cannot be rescued (see the method docstring), so this is a
terminal condition: the loop latches the health fault and lets the
exception exit the process, which is what makes ``torchrun`` tear the
whole group down in seconds instead of waiting out the NCCL watchdog.

Advance each constrained row's matcher by its committed token.

Runs on every rank (the matchers are rank-local copies built from the
same broadcast grammar), so each rank's matcher state stays identical
and the next step's bitmask is rank-symmetric. Mirrors the single-rank
``XGrammarLogitsProcessor.accept_token`` contract (skip terminated
matchers + stop tokens).

Advance request ``rid``'s matcher by one committed token.

The single place the SPMD mirror's matcher moves, shared by the
plain-decode sub-pass and the MTP verify sub-pass so both advance a
constrained row by the identical rule.

Advance each constrained row's matcher by its ACCEPTED tokens.

The MTP verify sub-pass commits several tokens per row per tick, so the
matcher has to walk all of them in order; the plain-decode counterpart
advances by the single sampled token.

Reduce the per-step outcome over the TP group (C9 invariant).

Returns True iff any rank failed this step. One small all_reduce
(MAX) per tick, issued symmetrically by every rank on both the
success and failure path — so it can never itself desync. This
turns an asymmetric step failure into a symmetric refuse: a
surviving rank learns its peer died and refuses the same tick
rather than advancing into the next collective alone.

WHAT THIS CAN AND CANNOT COVER. The guard only reaches both ranks
when the failure happened OUTSIDE the step's collective region —
i.e. the failing rank skipped no collective its peer issued. That
covers every rank-symmetric guard raise (the mirror-read
NotImplementedErrors, admission refusals, a scheduler-side error)
and any failure before the first collective or after the last.

It CANNOT convert an in-collective-region failure into a clean
refuse, and it never could — the original docstring's claim that
it did is what left arbi-serve#1302 mis-diagnosed for so long.
If this rank raises BETWEEN two collectives, the peer has already
enqueued the one this rank skipped and is blocked inside it. RTX 4090 and none
work:

  * The peer is not blocked in Python but in a CUDA device sync
    (``_spmd_forward_and_sample``'s sampled-token D2H) whose
    stream head is the unmatched collective. Only its OWN NCCL
    watchdog can break that — no message from this rank can.
  * Issuing this reduce on the forward's process group instead of
    the default one does NOT help: this rank's 5-element payload
    then collides with the peer's activation-sized collective at
    the same sequence number, which is a size mismatch — it hangs
    exactly the same way, and is what NCCL's "wrong sizes used
    across ranks" text warns about.
  * ``_abort_process_group()`` on this rank does NOT release the
    peer either; on a single node its collective keeps spinning
    with no disconnect to detect.

So for that class the guard's job is not rescue but FAST, ACCURATE
TERMINATION. When this rank failed, the wait is bounded by
``spmd_agree_timeout_s``; if the peer has not arrived by then it is
wedged, and :class:`SpmdPeerWedgedError` is raised so the loop can
latch the health fault and exit. ``torchrun`` then tears the group
down. Measured end-to-end: ~14 s to full teardown with an 8 s
bound, versus ~180 s (120 s watchdog + teardown) unbounded — and
the log names the ROOT-CAUSE exception instead of leaving only a
mismatched-``NumelIn`` NCCL dump two minutes later.

The bound applies ONLY on the failure path. A healthy tick keeps
the original single ``all_reduce`` + ``.tolist()`` with no polling
and no added latency: reaching this point on a healthy tick already
proves this rank is not the wedged one.

The reduce also carries this tick's captured-vs-eager forward
dispatch deltas (``[fail, cap, -cap, eag, -eag]``; MAX of ``x``
and ``-x`` recovers per-group max and min in one collective).
Every rank resolves each forward's replay-vs-eager branch from
the same derived batch through the same batch-keyed lookups, so
the deltas must agree tick-for-tick. Divergence means the ranks
took different dispatch branches — with equal collective shapes
that is a silent numerics/rank split (a hang would already have
surfaced), so it latches the sticky health fault and logs
critically rather than letting a bench run for hours against a
desynced pair.

Falls back to the local flag when no process group is initialized
(the single-process / CPU-test path).

Issue the agreement reduce with a deadline; raise if the peer is wedged.

Only ever called on the failure path (see
:meth:`_spmd_agree_step_failed`), so the polling cost is paid once
per failed tick and never on a healthy one.

``async_op=True`` + ``Work.is_completed()`` is the only bound that
works for both backends: for NCCL a plain ``wait()`` merely enqueues
a stream dependency and returns, so the host block would land on the
caller's ``.tolist()`` where no timeout can reach it.

A zero/negative ``spmd_agree_timeout_s`` restores the legacy
unbounded wait.

Debug-only per-tick page-state divergence probe (every rank).

Computes :meth:`RadixPageTable.lockstep_audit_digest` — (free-list
hash, tree hash, free count, node count) — and all-reduces MIN and
MAX over the TP group. Any component with MIN != MAX means the
ranks' page state diverged by this tick; log critically with the
component and the tick's delta composition so the first divergent
op is identifiable from one run. Log-limited to 5 hits.
Symmetric on every rank (same 2 collectives) — cannot desync.

Request ids that finished since the last tick (rank-0 only).

``tracker`` names WHICH slate tracker's ``_known`` set is
authoritative — the driver's own by default, or one attention-DP
set's when rank 0 is building that set's delta.

``drained`` is this tick's deferred-release drain result
(:meth:`_spmd_drain_pending_release`): under the radix deferral it
is the exact ordered set of rids whose pages rank 0 just released
at the top of this tick, so it is the evict list the worker must
replay (same rids, same order, same relative position). Under the
flat table (no deferral hook) the drain degrades to the legacy
alive-diff over the tracker's ``_known``. Either way, filter to
tracker-known rids (a request that never rode a slate has no
worker-side state) and drop any rid transiently re-emitted on this
very slate (the finished-reschedule defense) — loudly, since rank 0
already released its pages.

This tick's scheduler preemptions, as :class:`PreemptRow`s (rank-0 only).

``tracker`` / ``scheduler`` name which attention-DP set's state to
read; both default to the driver's own.

The rows come from the scheduler's own preempt log
(:meth:`Scheduler.drain_spmd_preempts`), written by
:meth:`Preemptor.preempt_for_space` at the instant it frees a
victim's pages and resets its decode state. Each row carries the
rebuilt prompt length — and, when the same admission pass put the
victim back on the slate, the prefix re-bind's match, amended onto
the row it already logged — so a worker reproduces the identical
commit + page free + re-add + recurrent-slab recycle + mirror reset
in ``apply_delta``.

The log REPLACED a post-hoc scan ("a known request, still alive,
not on this tick's slate, whose live state is reset while the
mirror still shows progress"). That scan could not see the most
common preemption of all: the victim is pushed to the FRONT of
``waiting``, so the SAME ``schedule()`` call's waiting-admit pass
re-admits it onto the very slate the preempt made room for, as a
chunked-prefill row of ``min(remaining, chunk_prefill,
token_budget)`` tokens. Being on the slate, it was skipped — no
PreemptRow was broadcast — and every rank's mirror kept the
request's pre-preempt DECODING state while the slate said the row
was a wide prefill chunk. ``resolve_row_shape`` then refused the
row ("decode row has n=2041, expected 1"), and the workers' page
tables + recurrent slabs never replayed the free/re-add, splitting
the per-rank free lists.

Two filters, both of which drop a preemption no worker can or
should replay:

  - the tracker must already have admitted the request onto a
    slate, else no rank holds worker state for it (its later
    AdmitRow carries the rebuilt prompt);
  - the request must still be alive, so a victim that finished
    before the delta was built rides ``delta.evict`` ONLY —
    ``apply_delta`` frees an evicted request's pages and would
    then RE-ADD them off a PreemptRow for the same rid.

Pre-plan the radix LRU evictions for this slate (rank-0 only).

Delegates to :meth:`RadixPageTable.plan_evictions_spmd` — rank 0
reclaims the slate's page-demand shortfall up front and the victim
ids ride ``delta.evict_pages`` so workers replay the same frees.
No-op (``[]``) on the flat table (no LRU eviction to plan).

On an MTP verify step (the same predicate ``_spmd_step`` routes to
:meth:`_spmd_run_verify` with — ``spec_mode == MTP and step_k >= 1``)
the tick is mixed: MTP decode rows allocate a tail slot plus
``step_k`` draft slots, while prefill / non-MTP rows run the legacy
sub-pass' full ``allocate_slots(rid, n)``.

Compute ``mtp_fill_enabled`` from the replicated mirror.

True iff the engine has an MTP driver and any slate row opted into
MTP. Derived from the mirror's per-row :class:`SamplingDigest`,
which is replicated on every rank — so both ranks fire (or skip)
the bundled-MTP-head all-reduce in lockstep. This is the crux of
the TP2 deadlock fix.

Build an on-device :class:`ScheduledBatch` from derived tensors.

Runs on every rank. Mirrors :func:`_reconstruct_batch_from_plan`
(the worker StepPlan rebuild) but sources the host tensors from
the rank-local :class:`DerivedStepTensors` instead of a broadcast
plan. ``mm`` / ``mrope_positions`` are populated by the replicated
vision tower + merge when a row carries image features (gated by
ARBI_SPMD_MM; inert otherwise). ``mtp_fill_enabled`` is set on both
ranks from the replicated mirror.

``req_ids`` carries the per-row request ids (slate order) so the
recurrent (GDN / Mamba) metadata builders resolve each row's real
slab row via :meth:`RecurrentStatePool.row_for` — the same route
rank 0's ``_build_batch`` uses. Both ranks now hold the same row
for the same request (the worker's ``apply_delta`` drives
``alloc_recurrent_state`` in lockstep with rank 0's scheduler), so
the GDN recurrent-state reads stay rank-symmetric.

``slate`` MUST be ``derived.slate`` — the rows the tensors describe
— not the broadcast slate the caller derived FROM. The page-demand
precheck (``precheck_step_demand_spmd``) defers rows that do not
fit, so the two can differ, and every row-indexed field built here
(``req_ids``, ``lora_assignments``) would then be a different
length than the row-indexed tensors. The concrete failure that
motivated this guard: ``req_ids`` at the UNTRIMMED length sized the
recurrent slab-row gather at B=8 while ``seq_lens`` resolved the
captured graph at B=6 — ``CapturedGraph.replay: recurrent_state_
indices[GDN] has shape torch.Size([8]); graph captured at
torch.Size([6])``.

Run the forward + sample identically on every rank.

Reuses ``ModelRunner.forward`` (the same captured-replay / live
forward + metadata path ``execute`` uses) — not a reimplementation.
The worker bridge is None under SPMD, so ``forward`` issues no
broadcast; every rank runs it directly and the row-parallel
all-reduce + lm_head all_gather fire in lockstep.

Sampling runs through the engine's real :class:`Sampler` over a
:class:`Request`-shaped view of the replicated mirror, so penalties
/ temperature / top_k / top_p / min_p are byte-identical on every
rank (deterministic functions of the shared logits + shared token
history). The stochastic Gumbel noise is supplied by the caller
from each row's ``(seed, position)`` — see
:func:`arbi_serve.sampler.rank_symmetric.gumbel_exp_noise` — so the
draw is rank-symmetric with no broadcast. A mid-prefill row (no
output token this step) samples ``None``.

Seed the MTP drafter for rows that just finished prefill (TP>1).

The SPMD counterpart of :func:`run_step._run_mtp_seed_forward`. Runs
the drafter chain over the just-sampled token's hidden and writes each
row's K drafts into the mirror, so the row's first verify tick already
holds a head-driven draft instead of missing the cache and collapsing
the whole slate to eff_k=0.

Rank symmetry (why this needs no broadcast). Every input is already
bit-identical on every rank:

  * ``slate`` / ``seed_rows`` — derived from the broadcast SlateDelta
    and the replicated mirror.
  * ``sampled`` — the rank-symmetric sampler's output (identical
    logits + rank-agreed Gumbel seeds).
  * ``hidden_last`` — the replicated forward's last-slot (B, H)
    hidden (``_seed_forward_or_replay``: eager gather or the
    captured graph's in-graph gather — same row either way).
  * the MTP head itself is replicated on every rank.

so every rank re-runs ``draft`` on its own head and derives the same
drafts into its own mirror — the symmetry is the broadcast. This is
exactly the contract the SPMD *verify* step's drafter chain already
relies on (mtp_verify_spmd), so the call is proven safe on this loop;
it was simply never made on the prefill tick.

The buckets are iterated in sorted key order so the sequence of
``draft`` calls — and therefore the sequence of collectives each one
issues, and the number of ``advance_mtp_seed`` advances — is identical
on every rank by construction, not by dict-insertion coincidence.

Rank-agreed base seed for one request's stochastic draw.

The explicit per-request ``seed`` when set, else a stable seed
derived from the request id. Both are identical on every rank (the
digest is replicated; the id is the same), so the derived Gumbel
noise is rank-symmetric.

Warm flat-dump cache for a separately-loaded spec-decode drafter.

The verifier reaches its weights through the flat dump
(:mod:`arbi_serve.loader.flat_dump`): one sequential read of a local blob,
overlapped with multi-stream H2D straight into the destination slabs. A
drafter is a second checkpoint on a second path, so it never entered that
route — it re-read its safetensors through one ``safe_open`` per tensor off
the model store and then migrated the result storage-by-storage into the
``model.drafter`` pool.

This module gives the drafter the verifier's two-route contract:

  * **warm** — build the drafter graph under ``skip_weight_load`` (module
    swap + placeholder buffers, no payload read) inside the ``model.drafter``
    pool, presize the placeholders from the dump header, DMA the blob into
    them, then fire the quant backends' rebind hooks;
  * **cold** — the checkpoint load plus the pool migration, then publish the
    dump so every later boot is warm.

The dump is keyed by
:func:`~arbi_serve.loader.flat_dump.default_flat_cache_dir` over the DRAFT
checkpoint's content and the parallel topology, so re-exporting the drafter
misses by key.
:func:`~arbi_serve.loader.flat_dump.flat_dump_compatible` gates the warm
route on an exact live-graph match and the loader's post-fill bit-signature
check refuses a blob whose bytes disagree with its manifest — so a warm
boot's drafter weights are bit-identical to a cold boot's, or the boot falls
back to cold and says so.

Cache entry holding ``draft_path``'s flat dump.

The shared content fingerprint folds in the checkpoint bytes, the parallel
topology (including rank) and the load ``dtype``. ``tp_shard`` is a
boot-time SELECTION on top of that topology — at ``tp_size > 1`` rank 0
holds a sharded drafter with the flag on and a full replicated one with it
off — so it carries its own suffix instead of sharing a key with a
differently-shaped tensor set.

Normalize the ``flat_dir`` argument to ``(path, key_or_None)``.

Callers that resolved the entry through :func:`drafter_flat_cache_key` hand
the key itself, which is what carries the fields the provenance manifest
records; a bare path (tests, older call sites) still works and simply has
no manifest to write or check.

Build the drafter warm-from-dump when possible, else cold.

``build`` constructs the drafter from its own checkpoint and is called
exactly once — under ``skip_weight_load`` for the warm attempt (topology
only, no payload read), bare for the cold one. Returns
``(model, route)`` with ``route`` one of ``"warm"`` / ``"cold"``.

Warm: the graph is built INSIDE ``pool`` so every dense parameter is
already pool-resident, the dump header presizes the quant placeholder
buffers into one slab per dtype (also inside ``pool``), and the DMA fills
allocated OUTSIDE the pool scope — a cuMem-backed ``MemPool`` never hands
a freed segment back (pytorch#145168), so a ring taken from ``pool``
would be a permanent free list on ``model.drafter``; in the default
allocator ``empty_cache`` genuinely reclaims it.

Cold: the load runs in the DEFAULT allocator, so its dense-construct and
staging transients return to the driver, and only the final persistent
tensors migrate into ``pool``.

An incompatible dump is deleted and the cold route runs — a mismatched
drafter graph is never served.

Write the drafter's flat dump so later boots take the warm route.

Called on a COLD load only, at the seam where the drafter's persistent
tensors are final and pool-resident — which is the same tensor set a warm
boot's ``skip_weight_load`` graph presents, so the dump round-trips.

Best-effort by design: a dump failure logs LOUD and leaves later boots
cold. It can never make a boot wrong — the warm LOADER checks every
filled tensor against the dump's bit signature.

Engine driver protocol — abstracts the per-step "execute one slate" op.

Two implementations:

  - :class:`Engine` — single-process driver: scheduler, sampler,
    forward, and post all run in one process at ``tp_size == 1`` and
    ``ep_size == 1``.
  - :class:`DistributedEngineDriver` — multi-rank driver: rank 0 owns
    scheduler + sampler + API + post-token, ranks 1..N-1 spin in a
    forward-only loop fed by a per-step plan broadcast from rank 0.
    Per-layer NCCL all-reduce inside :class:`RowParallelLinear` carries
    the data plane; ``dist.broadcast_object_list`` carries the control
    plane.

The driver is async-first: control-plane methods
(``add_request``, ``cancel``, ``shutdown``, ``aswap_attention_backend``)
are coroutines so the FastAPI surface composes without
``run_in_executor`` workarounds, and ``run_forever`` yields to the
event loop between steps.

Per-step execution surface.

The server's API layer holds an :class:`EngineDriver`; the driver's
impl decides whether to run all in-process or fan a forward step
out across multiple processes.

Both impls expose the same async surface — the FastAPI app wires
against this protocol so swapping single-process for distributed
is a one-line change in :mod:`arbi_serve.server.app`.

Hot-swap the backend for one :class:`StateKind`.

``spec`` is the registered ``"kind:name"`` (e.g.
``"paged_kv:tkv-k4v4"``); ``kind`` cross-checks against the
spec's declared kind. Returns a swap-result dict with
``{prev, new, kind, drain_s, alloc_s, total_s, stats}``.

Dynamic-config whitelist for ``GET / PATCH /v1/admin/config``.

Anything not in :data:`DYNAMIC_CONFIG_KEYS` is "static" and must be
changed by restarting the process. The whitelist gate runs here (not at
the HTTP layer) so internal callers (autotune, admission policy) hit
the same validated path.

**Every key WRITES THROUGH to the field its consumer reads.** A key that
only assigned the ``_dyn_`` shadow would make the endpoint accept the
value, echo it back from :func:`get_dynamic_config`, and change no
behaviour — an override that reports success and does nothing, which is
worse than one that errors. The shadow is kept only as the reporting
surface; it is never the enforcement surface. Where a key lands:

===================== ====================================================
key                   authoritative field (the one its consumer reads)
===================== ====================================================
``default_timeout_s`` ``eng.cfg.default_timeout_s``
``kv_watermark_pct``  ``eng.admission_cfg.kv_watermark_pct``
``queue_depth_max``   ``eng.admission_cfg.queue_depth_max`` + ``eng.inflight_limit``
``log_level``         the ``arbi_serve`` logger's level
``min_improvement_pct`` the calibration-reload quality floor
``prompt_prefetch``   ``server.prompt_prefetch``'s process-local override
===================== ====================================================

``prompt_prefetch`` is enforced in the process that SERVES HTTP, which is the
process this function runs in for the in-process server. Under the split
(API-child) topology the engine and the HTTP surface are different processes,
so ``/v1/admin/config`` also applies this one key locally in the HTTP process
(``routes/admin/config.py``) — the same setter, called from the only place
that is guaranteed to be on the request path.

The overlap's live state, read from the holder its consumer reads.

``None`` when the server module is not importable at all (a schema-only
process), which is the honest answer there: no request path, no state.

The values that GOVERN, read from the field each key's consumer reads.

Not the ``_dyn_`` shadows. A shadow is set only by an override, so reading
them reported ``None`` for every knob nobody had touched — on an endpoint
whose whole purpose is to answer "what is this server doing". ``None`` for
``kv_watermark_pct`` on a server admitting at 90% is not a smaller answer
than 90.0, it is a different and wrong one: a client cannot tell it from
"this knob does not apply", which is what a client seeing it concluded.

The authoritative field per key is the table in this module's docstring —
the same field the write-through in :func:`set_dynamic_config` targets, so
what is reported here and what is enforced cannot diverge.

``min_improvement_pct`` is the exception and stays a shadow read: it is an
operator FLOOR beneath a per-call argument, so unset is not a value the
server is using, it is the absence of a floor
(:func:`resolve_min_improvement_pct`).

Embedding + reranking primitives for the Qwen3-Embedding / Qwen3-Reranker
model families.

Both families are ordinary Qwen3 causal decoders, so the engine reuses
the existing prefill forward (see :meth:`Qwen3Model.forward`) and only
needs the small, model-family-specific post-processing collected here:

  * **Embedding** — last-token pooling of the post-final-norm hidden
    state (``return_hidden_state=True`` already exposes it), optional
    Matryoshka (MRL) truncation, and L2 normalisation. With causal
    attention the last token attends to the whole sequence, so a single
    prefill pass yields the correct sentence embedding — no bidirectional
    encoder is required.

  * **Reranking** — a single forward whose last-position logits are read
    at the ``"yes"`` / ``"no"`` token ids; the relevance score is the
    softmax weight on ``"yes"`` (see :func:`rerank_score_from_logits`).

Everything in this module is pure (tensor in, tensor/float out) and CPU
testable; the engine wiring that calls it lives in
:mod:`arbi_serve.engine.run_step`.

Return ``(yes_id, no_id)`` for the reranker score readout.

Mirrors the official Qwen3-Reranker reference, which resolves the
bare ``"yes"`` / ``"no"`` tokens. Supports both the HF transformers
API (``convert_tokens_to_ids``) and the Rust ``tokenizers`` API
(``token_to_id``); the latter returns ``None`` for an unknown token.
Raises if either token fails to resolve, since a silent
mis-resolution would corrupt every score.

Resolve ``(pooling_mode, l2_normalize)`` from the sentence-transformers
files a checkpoint ships, rather than assuming a strategy.

Reads ``1_Pooling/config.json`` (the ST ``Pooling`` module: one of
``pooling_mode_lasttoken`` / ``pooling_mode_mean_tokens`` /
``pooling_mode_cls_token`` is true) and ``modules.json`` (presence of
a ``Normalize`` module → L2-normalize). This is the same contract
sentence-transformers / Infinity / TEI honour, so it generalises
across the embedder family. Falls back to ``(lasttoken, True)`` —
the Qwen3-Embedding default — when the files are absent.

Pool a ragged (flat) batch of per-token hidden states to one vector
per sequence, per the sentence-transformers ``mode``.

``hidden_full`` is the engine's flat per-token post-final-norm state
of shape ``(N_tokens, hidden_size)``; ``cu_seqlens_q`` is the
``(B + 1,)`` cumulative token-boundary vector. Returns
``(B, hidden_size)`` in ``hidden_full``'s dtype.

The batch is ragged (no padding), so ``lasttoken`` is just the last
position and ``cls`` the first — none of the left/right-padding
bookkeeping the HF examples carry applies here. ``mean`` averages
each sequence's token slice via a segment reduction.

Optionally Matryoshka-truncate, then (optionally) L2-normalise.

``dim`` (the OpenAI ``dimensions`` knob / Qwen3-Embedding MRL) keeps
the first ``dim`` components before renormalising, matching the
reference recipe (truncate → normalise, not the reverse). ``None``
keeps the model's native width. ``do_normalize`` follows the model's
sentence-transformers ``Normalize`` contract (see
:func:`read_st_pooling_config`). Always promotes to fp32 — bf16
hidden states lose precision in the reduction and callers serialise
fp32 floats anyway.

Compute ``P(yes)`` relevance scores from last-position logits.

``last_logits`` is ``(B, vocab)`` — the engine's per-sequence
last-token logits (already sliced by :meth:`Qwen3Model.forward`).
Returns ``(B,)`` fp32 scores in ``[0, 1]``: a 2-way softmax over the
``no`` / ``yes`` logits, taking the ``yes`` weight, exactly as the
Qwen3-Reranker reference does.

arbi-serve engine: continuous-batching loop, multi-backend hot-swap.

The engine owns the model, the per-:class:`StateKind` active-backend
map, the :class:`MultiStatePool` and per-layer attn ops, and the
scheduler / sampler / tokenizer. It implements
:class:`EngineDriver` directly for the single-process case;
:class:`DistributedEngineDriver` is the multi-rank sibling.

Each step:

  1. Scheduler returns a slate ``[(req, n_tokens), ...]``.
  2. Engine extends each request's page table, records per-token slot
     indices, and builds a :class:`ScheduledBatch`.
  3. Each active backend's metadata builder populates the batch's
     kind-specific ``*_meta`` slot.
  4. ``model.forward(input_ids, batch, cache, attn_ops)`` runs.
  5. The :class:`Sampler` produces one token per request.
  6. Engine commits to the scheduler, emits new tokens to per-request
     async queues, and checks stop conditions.

Hot-swap (``swap_attention_backend(name, kind=StateKind.PAGED_KV)``)
sets the drain flag, waits for in-flight requests to finish on the
current backend, frees the per-kind pool view + attn ops + metadata
builder, builds the new ones (applying TKV calibration if loaded),
rebuilds the page table against the new pool, and clears the drain
flag.

CUDAGraph capture is opt-in: decode capture is on by default;
``--no-cuda-graphs`` disables.

The class :class:`Engine` declares state in ``__init__`` and binds
each method to a free function in a topic-focused sibling module:

  - :mod:`arbi_serve.engine.build` — boot orchestration, headroom
    probe, auth, active-state build, calibration apply, MTP build.
  - :mod:`arbi_serve.engine.run_step` — outer loop, MTP routing,
    terminal-state, shutdown. The actual model-execution body
    (flat-tensor build, captured-graph or live forward, sample) lives
    in :class:`arbi_serve.runtime.model_runner.EagerModelRunner` and
    is held on ``self.model_runner``.
  - :mod:`arbi_serve.engine.submission` — :meth:`submit` /
    :meth:`asubmit` + per-request build.
  - :mod:`arbi_serve.engine.swap_admin` — backend hot-swap,
    calibration reload, model reload.
  - :mod:`arbi_serve.engine.sleep` — release / resume, sleepable-state
    registration.
  - :mod:`arbi_serve.engine.cudagraph_admin` — boot-time cudagraph
    pre-capture sweep.
  - :mod:`arbi_serve.engine.lora_attach` — LoRA target-key discovery
    + admin attach hooks.
  - :mod:`arbi_serve.engine.distortion` — admin per-layer dequant
    relative-error report.
  - :mod:`arbi_serve.engine.dynamic_config` — runtime-reconfigurable
    config keys.

The method bodies themselves are grouped into cohesive mixins in
sibling ``engine_*`` modules so this file stays small:

  - :mod:`arbi_serve.engine.engine_init` — the ``_setup_*`` field
    initializers ``__init__`` invokes.
  - :mod:`arbi_serve.engine.engine_model_state` — named-pool
    construction + per-model container reset.
  - :mod:`arbi_serve.engine.engine_boot` — boot entry points, MTP
    seed/driver seams, cudagraph byte accounting + pre-capture.
  - :mod:`arbi_serve.engine.engine_admin` — admission, LoRA, critical
    section, stable-VA residency, hot-swap, distortion.
  - :mod:`arbi_serve.engine.engine_runtime` — sleep, dynamic config,
    submission, run loop.
  - :mod:`arbi_serve.engine.engine_base` — shared module-level
    constants + boot helpers (re-exported below).

Single-process top-level orchestrator.

Owns model + per-:class:`StateKind` active backends + pool + page
table + scheduler + sampler + tokenizer. Methods delegate into the
topic modules listed in this package's docstring. The method bodies
are grouped into the ``engine_*`` mixins this class inherits;
``__init__`` and the forwarding properties stay here.

Forwarding alias for ``cudagraph_pools["drafter_stoch"]``.

The true-stochastic drafter chain pool — populated at boot when
``ARBI_TRUE_STOCHASTIC_DRAFT`` is ``1``/``auto`` (the default; see
:func:`arbi_serve.engine.cudagraph_admin.precapture_drafter_chain`);
stays empty under an explicit mode ``0``, so the routing lookup in
:meth:`MtpDriver._draft_captured` misses to the live chain.

Engine admin-surface delegators.

Mixed into :class:`~arbi_serve.engine.engine.Engine`: admission checks,
LoRA attach, the critical-section factory, multi-model stable-VA residency,
and attention-backend / calibration / model hot-swap, distortion analysis.
Each method delegates into its topic module. Nothing here imports the
``engine`` module (avoids a cycle).

Public hook for HTTP handlers; thin wrapper around the helper.

``priority`` shapes the backpressure gate: a ``"batch"`` request
yields to interactive demand, while an interactive request is
not 429'd by a batch backlog. See
:func:`arbi_serve.engine.admission.check_admission`.

Return an async context manager that drains in-flight
requests, hands the caller a typed :class:`CriticalHandle`,
then resumes scheduling on exit. See
:mod:`arbi_serve.engine.critical`.

``kind`` is a :class:`StateKind` for backend / calibration
swaps, or :data:`arbi_serve.engine.critical.ALL_KINDS` for
sleep / full-engine reloads.

``drain_deadline_s`` bounds the in-flight drain: a config-swap /
backend-flip passes one so it refuses cleanly under sustained load
(``SwapDrainTimeout``) instead of wedging the engine. ``None`` (the
default) keeps the wait-forever semantics sleep / reload rely on.

Engage the opt-in multi-model stable-VA residency controller.

Call once after boot (the boot model is fully built + captured +
serving). Wraps the booted model as the active resident record so a
later :meth:`aswitch_model` to a different prepared model performs a
recapture-free park-then-wake switch. Idempotent — returns the
existing controller if already engaged.

Requires ``cfg.cumem_pools`` (the per-model park/wake unmaps physical
behind stable VAs through the cuMem allocator). Returns the
:class:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController`.

Apply a sparse config delta vs the active config (live override).

Each resulting variant is keyed + prepared/parked ONCE via the stable-VA
pool; re-selecting a known variant is an instant swap, a new
capture-affecting variant pays a one-time build. See
:mod:`arbi_serve.engine.config_variant`.

A delta that cannot be served RAISES, naming the config that is
actually serving. The residency rollback normally leaves the outgoing
member mapped and serving, and nothing is suspended; the sticky fault
is reserved for a rollback that could not put a serving member back. ``rollback_on_failure`` is the affirmative, default-off opt-in to
resume serving in that case; :meth:`arestore_last_valid_config` is the
same act performed deliberately, after the fact.

``drop_previous`` is the affirmative, default-off opt-in to fully
reclaim (not just park) the member a build-required swap supersedes —
see :func:`arbi_serve.engine.config_variant.aapply_overrides`. Set it
for a caller that walks through many genuinely distinct configs it
will never revisit (e.g. a bit-allocation search); otherwise a park's
standing VRAM tax (captured-graph exec memory + torch-MemPool
held-free segments neither ``empty_cache()`` nor the cuMem sleep path
can return) is paid again on every swap with no bound.

Resume serving the last config that actually served (affirmative).

The deliberate way back from a refused override. Raises when no config
has ever successfully served — there is nothing to return to.

Ephemerally drop ONE parked (non-active) config variant, reclaiming
all its VRAM (``DELETE /v1/admin/pool?key=…``).

The inverse of a build: return the residency to exactly the state before
the variant existed — no permanent KV/memory tax. Refuses to drop the
active member.

Register a post-switch hook run after every completed stable-VA
model switch (not on the no-op fast path, not during pool prepare).

For app-level state derived from the active model that lives outside
the engine's repointed per-model attributes (e.g. the server's
``TokenizerPool`` — a cross-arch switch replaces ``eng.tokenizer``
with a different-vocabulary tokenizer — and the cached multimodal
image processor). Hooks are zero-arg callables, run in registration
order by :meth:`StableVaResidencyController.switch_to`; a hook
exception is logged and does not abort the switch (the engine state
is already consistent — hooks own only derived state).

Register a pre-switch hook run while the outgoing member is still
the one bound to ``eng.page_table`` / ``eng.pool`` / ``eng.scheduler``.

The post-switch hooks above own *derived* state and are best-effort. A
pre-switch hook owns state that must be settled against the outgoing
member's own pools, because those objects are snapshotted into the
parked record and come back on a later wake — anything left queued
against them is not merely stale, it is leaked across the swap. The
SPMD driver's deferred finish-release queue is the motivating case: its
rids index the outgoing page table and recurrent pool, and its drain is
also the source of the workers' ``delta.evict``, so it must run (and be
broadcast) before the switch control op reaches the workers.

Consequently, and unlike the post-switch hooks, a pre-switch hook
exception aborts the switch: the engine is still fully consistent on
the outgoing member at this point, so refusing is safe, whereas
proceeding would strand the un-drained state. Zero-arg callables, run
in registration order.

Run the registered pre-switch hooks (see
:meth:`register_pre_model_switch_hook`). Called by every path that
repoints the engine's per-model attributes at a different member —
the stable-VA :meth:`switch_to` and the build-required variant path,
which repoints inside ``build_member_into_engine`` + ``set_active``
without going through ``switch_to`` at all.

Fan out a resident routing mutation.

Called by the stable-VA controller (register/switch) whenever the
routing-visible state changed. Hook failures are logged, never
raised — the mutation already completed; hooks own only derived
state (the process-mode control-event emitter). Defensive read:
a partially-constructed engine (unit tests build via ``__new__``)
simply has no subscribers yet.

Switch the resident model to ``key`` (recapture-free if prepared).

No-op when ``key`` is already active. Raises ``KeyError`` if ``key``
is not a registered resident record. Serialised + drained via the
engine critical section (no forward runs against half-mapped physical).
The registered model-switch hooks run after a completed switch.

``drop_outgoing`` reclaims the member being switched AWAY from, between
the park and the wake — the one window where its residue is both
free-able and still needed. Off by default: it destroys that member.

Switch to the named model, loading it into the residency if absent.

Present resident ⇒ the fast park/wake (:meth:`aswitch_model`); absent ⇒
build it into the residency (:func:`runtime_resident.aadd_resident_member`)
and make it active. Engages residency lazily for a single-model boot.
A distinct model loaded here carries its OWN per-member drafter source
(``enable_mtp`` / ``mtp_assistant_path`` / ``mtp_draft_model_path`` / …),
not the boot model's. The destructive full reload
(:meth:`areload_model`) stays the explicit fallback.

Prepare a single-VRAM-resident multi-model stable-VA residency set.

``members`` is ``[(served_name, builder), ...]`` in priority order. Each
``builder`` takes this (primary) engine, sets its ``cfg`` to the model's
config, and calls ``eng.build()`` — every member's per-model state is
built into the one primary engine (never a separate :class:`Engine`), so
the engine-global run-loop seams (``requests`` / ``_wakeup`` /
``_forward_executor`` / ``model_runner``) are shared and a woken member
serves through the same live run loop.

The orchestrator builds the boot model → offloads it to host → builds
the next member as the sole VRAM resident under its tag namespace →
offloads it → … → wakes the boot model. Only one model is VRAM-resident
at a time (the rest parked in host RAM); the outgoing member is
offloaded before the next is built, so no second model's physical is
mapped during a capture (removing the capture-while-parked IMA). A later
request for a parked member auto-switches to it recapture-free
(``docs/memory-accounting.md``, "Release, wake, swap").

``build_primary=False`` skips invoking the first member's builder — use
it when the boot model is already built (the server startup hook calls
``abuild()`` before engaging the pool), so the first member is just the
already-resident boot model and the orchestrator only builds + offloads
the extra members.

Returns the engaged
:class:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController`.

Shared module-level engine constants and boot helpers.

Lets the ``Engine`` mixin modules (``engine_init`` / ``engine_model_state``
/ …) reference the named-pool taxonomy, the ``--no-cumem-pools``
escape-hatch warning, the dtype resolver, and the stable-VA driver boot
guard without importing the ``engine`` module (which would be a circular
import — ``engine`` imports the mixins). ``engine.py`` re-exports every
name here so the canonical import site ``from arbi_serve.engine.engine
import NAMED_POOL_NAMES`` (and ``_torch_dtype`` /
``_check_phase2_driver_available`` / …) keeps working.

Fail loud at engine boot when the stable-VA sleep driver is missing.

The stable-VA (phase2) backend is the only production sleep path:
it preserves captured cudagraphs across release/resume by keeping
every ``data_ptr`` stable, which a drop-and-realloc approach cannot.
It requires the cuda-python ``cuda.bindings.driver`` bindings
(cuMemAddressReserve etc.), present on every production build.

On a real GPU boot, missing bindings is a hard error rather than a
silent degrade — there is no fallback backend. On a CPU /
no-GPU host (no CUDA available) we skip the check: the pool runs as
an inert stub and there is nothing to sleep.

Engine boot + CUDAGraph pre-capture delegators.

Mixed into :class:`~arbi_serve.engine.engine.Engine`: the synchronous /
async one-shot boot entry points, the MTP seed/driver seams, the
cudagraph-pool byte accounting + teardown, and the boot-time pre-capture
sweep hooks. Each delegates into a topic module. Nothing here imports the
``engine`` module (avoids a circular import).

Synchronous one-shot init. Called once at server startup
from sync call sites (CLI, smoke harness).

Runs the procedural boot in :func:`arbi_serve.engine.build.build`
— pre-flight, model load, tokenizer + xgrammar, profile,
active-state build, LoRA + MTP attach, activation arena,
cudagraph capture sweep, post-capture VRAM gate.

For async callers (FastAPI lifespan, distributed driver
startup), use :meth:`abuild`.

Advance the host-side step counter and copy_() it into ``mtp_seed``.

Returns the persistent buffer so the caller threads it into
:func:`verify_and_accept` / :meth:`MtpDriver.draft`. The
contract is "one advance per verify call"; the run-loop calls
this in the seed-drafter chain and inside
:func:`run_verify_step` once per slate.

Determinism. Two engine runs that hit the same number of
verify calls in the same order observe the same counter
sequence and so produce bit-identical samples — the basis of
the determinism tests. Tests that boot a fresh engine see
``counter`` start at zero; the first ``advance_mtp_seed``
call leaves ``mtp_seed[0] == 1``.

TP>1: rank 0 advances the local counter and broadcasts the new
value on the TP group; each worker rank's bridge dispatch runs
``mtp_seed.fill_(value)`` before any sampling kernel sees it,
so all ranks read identical Gumbel-max draws (the
determinism contract holds across TP topologies). The
broadcast is gated on the bridge's slate-active flag — boot-
time forwards (``profile_activation_peak``) advance the seed
but skip the broadcast since workers haven't entered the
dispatch loop yet.

Drop every captured graph across decode + layer + drafter pools.

Used by the shutdown / reload teardown path. The stable-VA sleep
backend keeps graphs alive across release/resume (their
data_ptrs survive the cuMem cycle), so the sleep-release path
does not call this — only a full teardown does.

Iterates :attr:`cudagraph_pools` so a future capture-pool kind
added to the dict at boot is automatically covered without
editing this method.

Simply dropping the pool dict + ``gc.collect()`` does not reliably
drive ``~CUDAGraph -> reset()``, so the capture pool's per-graph
``begin_allocate`` registrations survive teardown
(``capture.cudagraphs`` stays at ``use_count() == 1 + N`` instead of
1) and ``NamedPoolRegistry.free_all`` parks the multi-GiB pool in
``_LEAKED_POOLS`` rather than freeing it — a per-teardown VRAM leak
that compounds across every model cycle / sleep-rebuild. Release
each registration deterministically by ``reset()``-ing every
owned ``torch.cuda.CUDAGraph`` first. ``CUDAGraph.reset()`` is
idempotent, and the shared teardown clears the sleep entries before
calling this, so no graph is reset twice.

Resident cuMem bytes per CAPTURE pool, ``{}`` when cuMem is inactive.

The capture sweep costs VRAM in TWO named cuMem pools, and a ledger that
names only one of them is not a ledger:

  * ``capture.cudagraphs`` — the graph-private working set the captured
    graphs reference and can never free (in-graph ``lm_head`` logits,
    dispatch staging, the capture trampoline);
  * ``capture.io_buffers`` — the persistent kernel I/O buffers captured
    graphs bake by ``data_ptr`` (``logits_persistent`` /
    ``hidden_persistent``). Deliberately a SEPARATE MemPool (the BUG A
    address-reuse rationale in ``engine_base``), and just as resident
    and just as un-freeable as the first.

ever swept in accidentally, by the double-counted ``memory_reserved``
fallback this method replaces.

Reads :meth:`arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.
mapped_bytes` (a sum over LIVE allocations carrying the tag), NOT
``tag_mapped_bytes`` (an incremental counter maintained only for tags
that currently carry a CAP — it reads 0 both before ``arm_boot_caps``
and for any pool that was never capped, and stops tracking frees after
``clear_caps``). The tag is namespace-QUALIFIED via
:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.
qualified_tag`: under ``stable_va_residency`` the live allocations carry
``"<model-key>/capture.cudagraphs"`` and a raw-name lookup matches
nothing.

``peek()``, never ``get()``: an accounting probe must not construct the
native allocator on a boot that does not use one.

Total resident bytes across the cuMem CAPTURE pools (0 if none).

Sum of :meth:`cudagraph_bytes_by_pool`. See that method for why the
counter, the tag qualification and the two-pool sum are each
load-bearing.

True resident cudagraph-pool bytes — capture's whole VRAM cost.

  * **cuMem term** — :meth:`_cumem_graph_pool_bytes`: live mapped bytes
    under the ``capture.cudagraphs`` + ``capture.io_buffers`` tags. On
    the cuMem path (the default) this is authoritative and EXACT.
  * **torch-handle term** — ``Σ pool.total_memory_bytes()`` over
    :attr:`cudagraph_pools`, i.e. each ``CapturedGraph``'s
    ``capture_memory_bytes`` (a ``torch.cuda.memory_reserved`` delta
    bracketed across its capture). A per-graph ESTIMATE, used only when
    no cuMem allocator exists (``--no-cumem-pools`` / CPU).

The two are not disjoint — on the cuMem path the ``memory_reserved``
delta also sweeps in the transient forward-activation working set that
frees after the sweep — so they are never summed, only preferred in
order. The post-capture value is persisted into the
``graph_pool_budget`` cache, so an inflated figure poisons every warm
boot at that configuration with a spurious ``BudgetExceeded``.

HISTORY worth keeping: the cuMem term was DEAD CODE. It read
``tag_mapped_bytes`` (cap-gated) on an UNQUALIFIED tag, returned 0 on
every boot, and this method silently fell through to the torch-handle
sum — which was itself double-counted at the time (see
:func:`arbi_serve.runtime.capture._common.capture_resident_bytes`). The
doubled. Both halves are fixed; the authoritative counter now actually
runs.

True when the Phase-2 freeze has locked the allocator.

Once :func:`arbi_serve.engine.phase2_freeze.freeze_for_serving` runs,
the cuMem ``graph_pool`` is capped at its live size; a re-capture
would be denied at map time. Kernels are already captured at boot and
the canonical phase2 sleep/wake preserves them across stable VAs (no
re-capture), so a frozen capture call is a defensive no-op. A reload
calls ``unfreeze_for_reload`` before re-entering

Boot pre-sweep of every media binding's lazily-captured graph ladder.

A binding whose encoder captures per input bucket (today: the
NemotronVoiceChat perception tower's mel-frame ladder) would otherwise
mint its first graph on a live request — after the Phase-2 freeze, into
no capped tag. Capturing the whole ladder here puts it inside the boot
sweep's ``capture.cudagraphs`` tag, and :meth:`_seal_media_capture`
closes the ladder at the freeze so nothing can be added later.

Duck-typed on the binding's ``extra["capture_ladder"]``: a binding with
no ladder is skipped, so text / vision boots are unaffected.

Boot pre-sweep of the shared-KV drafter chain's ``(B, K)`` ladder.

A shared-KV head (the Gemma-4-class assistant, ``writes_draft_kv`` False)
is declined by the Qwen own-slab drafter sweep and self-captures inside
the driver, so without this its first live request at each shape captures
against the frozen layout, is denied, and serves eager for good. Duck-
typed: a drafter with no such ladder is skipped.

Close the drafter's own capture surface at the Phase-2 freeze.

The sibling of :meth:`_seal_media_capture` for the shared-KV chain: after
sealing, an uncaptured shape pads into a wider captured one or falls to a
WARNED eager chain, rather than attempting a capture the frozen pool
denies mid-request.

Close every media capture ladder so no bucket can be captured later.

Called by the Phase-2 freeze. After sealing, an input outside the
already-captured ladder takes the binding's eager encode path, whose
transient the boot profile folded into ``media_encode_peak_bytes``.

``Engine.__init__`` field-setup helpers.

The ``_setup_*`` methods here are pure state initialization for
:class:`~arbi_serve.engine.engine.Engine`, kept in a separate module so
``Engine`` stays under the per-file size ceiling. They are mixed into
``Engine`` and run against a real engine instance (``self`` is an
``Engine``); nothing here imports the ``engine`` module (avoids a
circular import — ``engine`` imports this).

Named-pool construction and per-model container reset.

These methods are mixed into :class:`~arbi_serve.engine.engine.Engine`.
``_init_named_pools`` builds every :class:`NamedMemPool` bucket;
``reset_model_state_for_build`` re-initialises the per-model containers so a
fresh model can be built into the one primary engine (multi-model stable-VA
residency); ``_reset_workspace_pool_state`` drops the previous member's
GDN-FLA / cuBLAS workspace routing. Kept in a separate module to keep each
file under the size ceiling; nothing here imports the ``engine`` module.

Create every :class:`NamedMemPool` bucket on ``self.named_pools``.

Extracted from ``__init__`` so the in-process model hot-swap
reload path can re-create the named pools after a coordinated
teardown freed them: :meth:`NamedPoolRegistry.free_all` empties
the registry (returning each bucket's VRAM to the driver), and
:func:`arbi_serve.engine.build.build` only ever ``get()``s pools —
it never makes them. Idempotent against a populated registry only
in the sense that callers must clear it first; here we assume the
registry is empty (fresh ``__init__`` or post-``free_all`` reload).

Re-initialise the per-model containers so a fresh model can be built
into this engine in place (multi-model stable-VA residency).

The build-into-the-primary residency design (see
:func:`arbi_serve.engine.stable_va_controller.build_member_into_engine`)
constructs every member's per-model state on the one primary engine so
the shared run loop (``self.requests`` / ``self._wakeup`` /
``self._forward_executor`` / ``self.model_runner`` — the forward-worker
thread + intake, all engine-global) serves whichever member is repointed
active. Before building member B we snapshot member A's per-model state
into A's record, then call this to install fresh containers so
:meth:`build` populates B's state cleanly without mutating the objects
A's snapshot still references.

Re-creates exactly the per-model containers ``__init__`` builds (named
pools, the four cudagraph pools, the per-tensor sleep registry, the
multi-group cache, the backend / metadata / attn-op registries) and
clears the per-model handles (model, pool, page table, scheduler,
sampler, drafter, verify/piecewise buffers, stores, activation profile).
The shared run-loop seams are deliberately untouched. Idempotent against
a freshly-snapshotted engine (the caller owns retaining the old objects
via the snapshot).

Drop the previous member's GDN-FLA / cuBLAS / sampler workspace state.

A live backend / model swap parks the active member and offloads its
pool physical to host (``evict_parked_to_host`` → ``sleep_all``), then
builds the new member into the same primary engine. Process-global
device-tensor caches outlive that evict and, left dangling, point into
the previous member's now-unmapped pool physical:

  * ``_fla_persistent_cache._GDN_WORKSPACE_POOL`` — routes the GDN-FLA
    chunk-prefill intermediate allocations (e.g. ``chunk_local_cumsum``);
    set on the previous member's wake / capture to its ``gdn_workspace_pool``.
  * the cuBLAS / cuBLASLt per-(device, stream, handle) workspace cache —
    pinned (for the process lifetime) by the previous member's main +
    GDN per-layer side streams, living inside its ``cublas_workspace``
    cuMem pool.
  * every OTHER module-level device-scratch memo, enumerated in
    :data:`arbi_serve.engine.member_scratch_retire.
    MEMBER_SCOPED_DEVICE_CACHES` — the tkv Turbo prefill/verify LSE and
    split-KV partial scratches, the tkv Hadamard / VQ2-LUT caches, the
    speculative-decode tree/index/seed/margin caches, the sampler
    ones-tensor, and the two ``prepare_chunk_indices`` caches that hold
    the GDN batch metadata's chunk-index tables.
    Each is keyed by ``(device, dtype, shape)`` with no member
    discriminator and is first allocated inside whichever named pool
    was ambient — for the capture sweep that is the member's own
    ``capture.cudagraphs`` pool, which the park evicts. They are
    RETIRED (pinned, then dropped from the memo) rather than freed, so
    a parked member's captured graph never loses an address it baked.
  * the stochastic-sampler Triton scratch/table caches
    (``topk_topp_triton._TRITON_BUFFER_CACHE`` / ``_TRITON_TABLE_CACHE``)
    — the per-(device, vocab) fp32 pivot-search buffer + CDF→sigma
    tables, first allocated by whichever member ran the sampler first
    and not rebuilt on a member swap. The incoming member's build-time
    sampler warmup (``warm_sampler_kernels``) can reuse the outgoing
    member's cached buffer whose physical the park had unmapped,
    faulting the ``_topk_topp_kernel`` launch with
    ``cudaErrorIllegalAddress`` even though a fresh ``torch.randn`` on
    the same line succeeds — the fault is the stale cached tensor, not
    a poisoned context.

The new member's build runs a live activation-profile forward (a manifest
miss, because a backend change moves the fingerprint) before it is woken
(which is what repoints these globals to the fresh pools). The prefill
shape's GDN-FLA cumsum and the decode shape's fused-projection GEMMs
(``GDNBlock._project_streams_raw_ba``) would then dereference the evicted
VA → a deterministic ``cudaErrorIllegalAddress`` that poisons the context
and refuses the swap.

Clearing the two globals + dropping the cuBLAS handle cache makes the
build's forwards fall back to the default allocator — exactly the
cold-boot path, where these globals are still None during the activation
profile (they are set later, at capture) and the cuBLAS workspaces are
seeded fresh and never evicted. The new member's fresh pools are wired in
on its wake / capture, so steady-state routing is unaffected. Safe to
clear here: the reset runs under the swap critical section (scheduler
drained, no live forward) and the previous member's captured graphs are
parked / slept, so no replay references the freed workspace. Mirrors the
shutdown teardown in ``engine/lifecycle.py``.

Engine sleep / dynamic-config / submission / run-loop delegators.

Mixed into :class:`~arbi_serve.engine.engine.Engine`. These are the
request-facing and run-loop surfaces: sleep (release/resume), the
hot-reloadable config keys, ``submit`` / ``asubmit`` intake, terminal-state
callbacks, cancel / shutdown, and the per-step ``run_step`` entry point.
Nothing here imports the ``engine`` module.

Async submit ``prompt`` through the message-shaped intake (P3).

Builds a ``SubmitMsg`` API-side and awaits the engine-side
RequestFactory's SubmitAck — see
:func:`arbi_serve.engine.submission.asubmit` for the full
parameter contract (auth/tenant bound BEFORE publish;
``target_resident`` routes stable-VA residency; ``client`` is the
SSE chat handler's pre-made ``ClientRequest`` with the tool stream
already attached, fully API-side).

``multimodal`` carries the per-modality preprocessed features
dict (e.g. ``{"image": MultiModalFeatures}`` /
``{"audio": AudioFeatures}``); ``prompt`` must already include
the expanded placeholder tokens.

Cancel a request BY ID — the P3 HTTP-path cancel surface.

The bridge marshals the body onto the engine loop; an unknown id
(cancel racing a not-yet-published submission) is recorded in
the cancel tombstones so the RequestFactory drops the racing
``SubmitMsg`` with a ``cancelled`` FinishOut. P4b maps this to
``CancelMsg`` over the wire. ``cancel(req)`` stays for in-proc
handle-holding callers (bench / drain).

Run one engine step from a pre-built :class:`StepPlan`.

``run_step`` is the **single public step entry point** — what
the run loop drives, what tests call to exercise one step in
isolation, and what the distributed driver dispatches per rank.
The convenience constructor :meth:`StepPlan.from_slate` carries
the ``[(req, n_tokens)]`` slate through.

Returns a :class:`StepResult` summarizing per-row sampled
tokens, accepted lengths, and finished request IDs.

Dedicated-OS-thread launcher for the engine run loop.

Default-ON (``runtime_flags.engine_own_thread``; disable with
``ARBI_ENGINE_OWN_THREAD=0`` for the legacy single-loop mode). The
engine's ``run_forever`` coroutine runs on its OWN OS thread with its OWN
asyncio event loop, so the engine's per-step Python bookkeeping never
holds the GIL/loop on the HTTP thread — uvicorn's I/O poll stays free and
a streamed first SSE chunk flushes the instant its token lands (closing
the c=1 TTFT SSE-starvation gap with no per-step forward-offload tax).

Cross-loop safety is handled by :class:`arbi_serve.engine.loop_bridge.LoopBridge`
(activated here once both loops exist): intake is marshaled onto this
engine loop; per-token output is marshaled back onto the HTTP loop. All
CUDA stays on this thread.

Lifecycle:
  * :meth:`start` — spawn the thread, create its loop, schedule
    ``run_forever`` on it, activate the bridge, and block until the loop
    is live (so the caller can rely on the bridge being armed before the
    first request).
  * :meth:`stop` — request ``eng.shutdown()`` (sets ``_shutdown`` + wakes
    the loop), wait for ``run_forever`` to return, then stop + close the
    loop and join the thread. Coordinated from uvicorn's lifespan
    ``finally`` so SIGTERM never hangs.

Spawn the engine thread + loop and arm the bridge.

Blocks until the engine loop is running and the bridge is
activated, so the lifespan can accept traffic knowing every
cross-loop handoff is armed.

True while the engine OS thread is running.

Read by the ``/health`` liveness probe: a dedicated-thread engine
whose run loop has crashed out (the coroutine raised outside the
per-step guard) leaves the thread dead while the HTTP server stays
up, and the probe returns 503. False before :meth:`start` and after
the thread joins.

The probe is the SECOND line of defence, not the first. A
healthcheck window wide enough for a legitimate admin model reload
(interval * retries) is also a window in which a dead engine keeps
answering, so a terminal run-loop fault exits the process itself
(:func:`~arbi_serve.engine.run_loop_exit.note_run_loop_exit`) and the
container's restart policy recycles it in seconds. This stays for
the cases the exit cannot cover: an unarmed process, and the
interval between the fault and the exit.

Stop the engine loop and join the thread. Idempotent.

Calls ``eng.shutdown()`` (which sets ``_shutdown`` and wakes the
loop via ``_wakeup``) ON the engine loop thread-safely, then
waits for ``run_forever`` to return and the thread to exit.

Boot-time reservation of the shared EXL3 reconstruct scratch.

WHY THIS EXISTS. The EXL3 large-M (prefill) leg dequantizes a full fp16
``(in_features, n)`` weight before each ``hgemm``. That buffer carries WEIGHT
geometry and no token dimension, so ``max_batched_tokens`` cannot bound it, and
a per-call allocation leaves the caching allocator holding one retained block
per distinct linear shape — the pool's resident total is the SUM over the
model's layer geometries. In the cuMem-backed ``scratch.forward_arena`` those
blocks also stay mapped once freed, so the KV sizer must hold every one of them
back for the process's life.

:func:`~arbi_serve.weight_quant.exl3.custom_op.reserve_reconstruct_scratch`
replaces that sum with a single buffer sized to the widest reconstruct any
linear that can REACH the leg will ask for. This module is the engine-side seam
that decides WHEN it is reserved, WHICH pool it lands in, and — through
:func:`reconstruct_row_bounds` — which linears can reach the leg at all.

REACHABILITY. The leg is row-gated: upstream routes a call to
reconstruct+hgemm only above this shape's ``auto_reconstruct_threshold``. A
linear whose widest call site stays under that threshold therefore never draws
a slab, and sizing the shared buffer to it reserves bytes nothing can spend.
The logits head is the case that matters: its call sites gather one row per
sequence (or ``K + 1`` per speculative slate row) where the body linears see a
whole prefill chunk, so a narrow serving width puts a vocab-wide head below its
own threshold. That is a property of the CONFIG, not of the checkpoint, so the
bound is derived from the batch geometry on every boot rather than assumed.

WHEN. After every EXL3 linear is bound and before the first forward, so the
activation profile, the capture sweep and the KV sizer all measure the steady
state. The alternative — growing it on the serving path — allocates out of
headroom the KV sizer has already given away; the TP gather staging arena
(:mod:`arbi_serve.engine.nccl_staging_warmup`) is the precedent, and it OOMed a
rank inside a collective before it was pre-reserved.

WHICH POOL. ``capture.io_buffers``: a captured prefill graph runs the
reconstruct leg and bakes the buffer's ``data_ptr``, which is exactly what that
pool is for (member-lifetime, never destroyed, kept out of the capture mempool
so no address reuse can corrupt a replay). The forward arena is the wrong home
for the same reason its overhang reclaim is sound — it re-homes what it holds
onto fresh addresses, which would move a pointer a graph had already baked.

WHOSE. That pool belongs to ONE member, and a capture-affecting config override
builds a second member off the same weights while parking the first. So the
reservation is per member: a linear reaches its buffer through its own inner,
never through process-wide state, and the cross-seam handle below is dropped
with the rest of the per-member state when the next member's pools are made
(``_reset_model_state``, the same seam and the same reason as the NCCL staging
buffers). A shared buffer pointed the rebuilt member at the parked member's
pool, whose physical is dropped behind its VA — the rebuilt member's first
profiling prefill then faulted inside ``reconstruct`` and poisoned the CUDA
context, latching a swap fault that suspended serving.

Book a fork-surface skew that silently turned the fused leg off.

THE SILENCE THIS REMOVES. The fused reconstruct is optional and
numerically indistinguishable from the leg it replaces, so an exllamav3
older than the pin costs prefill traffic and nothing else — the op
answers the capability question and serves the standalone leg (see
``custom_op._FUSED_RECONSTRUCT_FORK_GAP``). Every other thing about that
boot is identical to a boot that ran the leg: same banner, same readiness,
same served outputs. A degradation nobody can see is the "green means it
ran" defect wearing a different hat, so the capability answer is paired
with this row.

HERE because this is the one boot seam that runs with every EXL3 linear
bound, and only when a linear can actually REACH the reconstruct leg —
a build that draws no scratch has no fused leg to lose. Idempotent across
the two reservation seams by looking the row up rather than by a module
flag: the ledger is per ENGINE, and a process-global "already reported"
would silence the second member of a capture-affecting rebuild.

The exllamav3 commit actually installed, or a reason it is unknown.

Read from the distribution's ``direct_url.json`` — a git install records
the resolved commit there — because ``__version__`` is ``1.4.2`` on the
pinned revision AND on the PyPI wheel, so it cannot tell an image built
against one rev from a lock naming another. That is precisely the skew
this line has to name.

Rows per slate — the served width, known from the start of boot.

Both reservation seams run before anything is profiled, and both read the
same concrete number the scheduler will serve. That is what can put the
logits head BELOW its own reconstruct threshold on a narrow deployment: a
worst-case stand-in would keep the head on the reconstruct leg and hold
scratch for rows nobody configured.

Tokens one slate row can carry through a forward: ``K + 1``.

A verify step runs the anchor plus ``K`` speculative tokens per row. ``K``
is read from the attached drafter when there is one — it owns the served
ceiling — and from the configured draft width otherwise, which is what the
first seam has before the drafter is built.

``(body_rows, head_rows)`` — the widest row count each call class emits.

``body_rows`` bounds any linear reached from inside a model forward: the
step token budget, or a full speculative slate when that is wider.

``head_rows`` bounds the logits head, whose call sites gather rows rather
than tokens — one per sequence for sampling, ``K + 1`` per row on a verify
pass, and one prompt-logprobs tile when that surface is armed. A head is
therefore orders of magnitude narrower than the prefill chunk feeding the
body, which is what can put it below its own reconstruct threshold.

Both are UPPER bounds. Anything that widens a call site past them belongs
here, or the linear falls back to a per-call allocation outside the budget.

The EXL3 inners behind the model's logits head, by object identity.

Identity rather than a name match: ``model.lm_head`` is the handle every
logits call site goes through (the bundled draft head asserts it holds this
same object), so it names the head without a naming convention to drift.

Reserve the shared EXL3 reconstruct scratch for this rank's geometry.

``when`` names the boot seam for the log line. Returns the bytes reserved
(0 when there is nothing to reserve: no EXL3 linear bound, no CUDA, no
pool). NEVER raises — a missed reservation costs the per-shape retention
this removes, it does not break a boot: the op falls back to the per-call
allocation it used before.

Idempotent across seams. The base model reserves first; a later call (after
a drafter checkpoint attaches its own EXL3 linears) re-reserves only if the
drafter's geometry is wider. Both run before the capture sweep, so no
address a graph has baked can move.

``eng`` carries the handle only so the two seams of ONE member's build can
grow a single buffer; the member reset clears it, so a later member never
reuses a pool this member owned (see the module docstring).

Boot-time resolution of the EXL3 trellis-GEMM kernel-shape pin.

WHY THIS EXISTS. ``exl3_gemm`` picks its kernel by row count — a GEMV special
case at small m, an autotuner over the shape table above it — and picks that do
not share a ``(TILESIZE_K, TILESIZE_N)`` do not agree bit-for-bit. Unpinned,
one row's output therefore depends on how many other rows shared its call,
which is the property speculative decoding's losslessness rests on: MTP verify
runs ``B x (K+1)`` rows against decode's ``B``. The measurement and the
mechanism live in :mod:`arbi_serve.weight_quant.exl3.custom_op`; this module is
the engine-side seam that decides WHEN the pin is resolved.

WHEN. The same two seams the reconstruct scratch uses — after every EXL3 linear
is bound and before the first forward, then again after a drafter checkpoint
attaches its own. Before the first forward for two reasons: the selection TIMES
candidate kernels, which must not happen inside a cudagraph capture, and a
linear that reached serving unpinned would run a different kernel from its
already-pinned siblings.

WHAT THE SEAM REPORTS. Coverage — how many BOUND EXL3 linears carry a pin, and
whether the drafter's own linears are among them — not how many this pass
happened to add. The resolution is idempotent, so a seam that runs after
another which already pinned everything adds zero; read as a count of pinned
linears, that zero says the pin is dead. It said exactly that on every boot of a
checkpoint with a bundled MTP head, whose draft-head linears are bound with the
base model and pinned at the FIRST seam: the post-drafter line reported "0
linears pinned ... not bitwise lossless" over a fully pinned 409/409 model. A
verdict has to be a property of the state, not of the last delta.

``(module path, inner)`` for every BOUND EXL3 linear under ``root``.

The inner is the object the pin lands on, and a linear that has not built
one is skipped rather than counted unpinned: the custom op resolves its
kernel host through that inner and REFUSES a call without one, so an
unbuilt linear cannot serve a row and cannot make one row's output depend
on how many rows travelled with it.

``(label, root)`` of the attached drafter's own weight tree.

The bundled MTP head hangs off its driver as ``head`` and a separate draft
model (external / DFlash) as ``model``; either way it is the tree whose
linears the verify step's losslessness depends on, and the one the
post-drafter seam exists to cover. ``("", None)`` when nothing is attached.

Named rather than counted: "409 of 409 pinned" cannot distinguish a model
whose draft head is covered from one whose draft head was never bound as an
EXL3 linear at all, and that distinction is the whole question at this seam.

Say so, loudly, when a verify slate would take a different LEG from decode.

Returns the number of crossing geometries (0 = nothing to say).

The kernel-shape pin makes the trellis GEMM row-count invariant, which is
what speculative decoding's losslessness rests on. It cannot make the LEG
invariant: the leg is chosen from the row count
(``rows > auto_reconstruct_threshold``), a verify step runs
``B * (n_draft + 1)`` rows where the plain decode of the same tokens runs
``B``, and above the threshold the op leaves ``exl3_gemm`` for
``reconstruct`` + cuBLAS ``hgemm`` — a different algorithm the pin never
reaches. So a boot can hold a perfect pin and still not be lossless.

Checked HERE because this is the one seam that has both halves: the served
config (``eng.cfg``) and the set of geometries actually bound. It runs after
the pin so the two lines read in the order an operator needs them — what the
pin achieved, then what it does not cover.

WARN, NOT REFUSE. The condition is latent at every configuration we deploy
(``max_batch = 8`` with ``n_draft = 7`` gives a widest slate of exactly 64
against a threshold of 64, and the dispatch is ``rows > threshold``), and it
arms only from ``max_batch = 9``. An operator who raises concurrency is
making a legitimate choice whose cost is a numerics guarantee, not a crash,
and the fix is a fork change plus a pin bump rather than anything this
process can do. What is not acceptable is that they find out from an issue
they never read — which is the whole reason this is code and not a comment.

NEVER raises: a check that cannot run must not be able to fail a boot whose
pin is otherwise fine.

Widest row count a decode or verify step can present, or ``None``.

``None`` means "this boot cannot prove it", and it is the FAIL-CLOSED
answer: every caller must pin on it. A gate that opens when it cannot see
the inputs is a gate that removes a correctness guarantee precisely in the
configurations nobody characterised.

The slate is ``B x (K+1)``: a verify step runs the drafter's K proposals
plus the bonus token for each of up to ``max_batch`` sequences, where the
plain decode of the same tokens runs ``B``.

K AT THE POST-LOAD SEAM. This is called once before the drafter attaches
and once after, and the two must not disagree -- a gate that opens at the
first seam and closes at the second leaves the activation profile running
kernels serving will not use. ``cfg.mtp.n_draft`` is authoritative whenever
it is positive, because that is literally what the attach sites pass
(``build_dflash_drafter.py``: ``max_k = int(cfg.mtp.n_draft) if
cfg.mtp.n_draft > 0 else block_k``). When it is zero the served K comes
from the draft checkpoint's block size, which the post-load seam cannot
read -- so that case answers ``None`` rather than guessing, and an attached
driver's own ``max_k`` is folded in whenever one is present.

``(pin?, why)`` — whether this boot needs the pin for its own losslessness.

The pin costs decode throughput: measured at +3.76% of leg A's weight pass
at M=1 and +3.86% at M=8 on a 4090 (#1861). It buys row-count invariance,
which the unpinned dispatch ALREADY provides up to
:data:`UNPINNED_INVARIANT_MAX_ROWS`. Paying for a guarantee you already
hold is the only case this gate declines.

NOT A PERFORMANCE KNOB. There is no flag to force it open: the input is the
served slate width, so the only way to widen the decline is to serve a
narrower slate, which is a capacity decision an operator makes for its own
reasons. Every other answer pins.

The row counts THIS deployment presents, or ``None`` when it cannot say.

``{1, max_batch, max_batch * (K + 1)}``: a decode step runs one row per
scheduled sequence and a verify step runs the drafter's K proposals plus
the bonus token for each of them, so those are the widths the pinned family
has to be fast at. One row is in the set unconditionally — a single-stream
request is a shape every deployment presents whatever its concurrency cap.

The slate comes from :func:`widest_leg_a_slate`, which is the same function
the pin's own gate reads, so "how wide does this boot get" has one answer
at this seam rather than two that can disagree. ``None`` there means the
served K is not knowable yet, and it means the same here: a derived set
built on a guessed K would race the families over rows this deployment does
not run and then record the winner under that set's key.

``(rows to install, why)`` for this boot's kernel-shape race.

``None`` rows means INSTALL NOTHING — the module's own resolution (the
shipped row set, or an explicit ``ARBI_EXL3_PIN_SELECT_ROWS`` list) governs
unchanged. Only ``auto`` needs this seam, because only ``auto`` is a
function of the served config, which the weight-quant module never sees.

An ``auto`` this boot cannot derive falls back rather than refusing: the
fallback is the row set every deployment raced before the knob existed, and
a boot that would otherwise serve is not worth failing over which rows its
families were timed at.

Freeze one trellis-GEMM kernel-shape family per bound EXL3 geometry.

``when`` names the boot seam for the log line. Returns the number of
linears newly pinned by THIS pass — zero both when there was nothing to pin
and when everything bound was already pinned, which is why the log verdict
is read off coverage instead (see the module docstring). NEVER raises — an
unpinned linear keeps the row-count-dependent auto dispatch, which is the
pre-pin behaviour.

Logged per GEOMETRY, not per linear: a 27B carries ~400 linears over 9
distinct shapes, and the pin is a property of the shape. The full mapping,
including the selection timings each choice won by, stays readable from
:func:`~arbi_serve.weight_quant.exl3.custom_op.kernel_shape_pins`. Only the
linears a pin did NOT reach are named individually, because those are the
ones somebody has to go and look at.

``"<band>: <geom>=shape<id>, ...; ..."`` — the VERIFY line, per band.

ONE CLAUSE PER BAND, because that is the decision the boot took: the tile a
verify slate runs on follows the slate's own width, so a line that named
one tile for the class would be naming it for widths that do not run it.
An operator reads this against the captured ``(B, S)`` ladder the cudagraph
line prints.

The stable-VA member whose linears this arming owns.

Two seams reach the arming and they answer the question differently, so
both are read, in this order:

  * a member BUILD runs inside ``alloc.tag_namespace(key)`` while the
    residency controller's ACTIVE record is still the OUTGOING member, so
    the open namespace is the only correct answer there;
  * a WAKE has no namespace open, and the active record's namespace is.

``None`` (neither available) is the un-namespaced member — a boot with no
residency — for which the scoping is a no-op.

The decode row counts this boot's graphs present: the captured B ladder.

The same ladder the capture sweep records (``_decode_b_ladder``), so the
probe times exactly the widths a replay can run and nothing it cannot; an
intermediate B pads up to the next rung and takes that rung's decision.
Empty when the config cannot say, which resolves the threshold unresolved.

Everything the probed threshold depends on, for the budget-cache key.

The trellis pin's own key (card, driver, torch, exllamav3 and arbi-serve
revisions, the shape table, the raced row set, TP size) plus the served
shape: max_batch, the widest verify slate (which carries K), the prefill
chunk and the checkpoint. The row set and the int8 plans are added by the
resolver itself, which is where they are known.

Pin the int8 leg's launch shape and reserve its scratch. Returns bytes reserved.

SEPARATE FROM THE TRELLIS PIN ON PURPOSE. ``pin_exl3_kernel_shapes``
declines itself whenever this boot's widest slate is narrow enough that the
trellis GEMM is invariant without it, and the int8 leg's pin must not
disappear with that decline -- they answer different questions about
different kernels.

NO-OP UNLESS ARMED, and cheaply so: ``resolve_int8_shape_pins`` returns
before touching the kernel when the flag is off, so an unarmed boot does
not pay the nvcc build. NEVER raises -- an unarmed or unreserved leg refuses
by name and the shipped legs serve.

Two-sided by construction. What it logs is BOTH the geometries pinned and
the geometries declined with the reason, because a boot that reports six
pinned and says nothing else reads identically to one that silently served
none of the other three.

The token a request fed at every position it has forwarded.

A request's fed stream is its prompt, then its output tokens and its
consumed context injections interleaved in forward order. Injected
context never enters ``output_token_ids``, so ``context_injections``
— pairs of ``(output tokens committed before the injection, the
injected ids)`` — is the only record of where those tokens sit.

Position ``p`` is the RoPE position the model saw a token at, which
makes this the map any state keyed on token HISTORY (Qwen4-Exp PLE
n-gram context) is rebuilt from: chunk boundaries, preemption
rewrites, prefix-cache hits, and context injection all move a row's
starting position without changing the stream underneath it.

Mixed into both :class:`~arbi_serve.engine.request.Request` and the
per-rank SPMD mirror row, so every rank names the same token at the
same position. Implementors supply ``prompt_token_ids``,
``output_token_ids``, ``pending_context_token_ids``,
``context_consumed``, and ``context_injections``.

Number of tokens whose fed identity this request can name.

``len(prompt) + len(output) + context_consumed``. Differs from
``total_length`` by the optimistic advances — tokens that were fed
but whose ids are still device-only.

Fold the queued context injection into the fed-token record.

Moves ``pending_context_token_ids`` into ``context_injections`` and
``context_consumed`` in one step, so the count that advances
``total_length`` and the ids that reconstruct the stream can never
disagree. Idempotent: nothing queued is a no-op.

Periodic engine heartbeat — a CHANGE-driven operator line.

ENGINE-plane module: the heartbeat runs on the engine loop (started by
``run_forever`` / the distributed driver's rank-0 loop) and every
per-tick read is a direct, loop-local engine read; in process mode it
lives in the engine process with no HTTP surface at all.

The line reports the phase the engine is actually in::

    arbi  prefill <rate> tok/s · req#<id> <done>/<total> (<pct>), <age>
        · model <name> · <n> active · GPU <used>/<total> · KV <pct>
    arbi  prefill <rate> tok/s · req#<id> <done>/<total> (<pct>), <age>
        · <n> active · KV <pct>
    arbi  decode <rate> tok/s · req#<id> out <n>, <age> · <n> active
        · KV <pct> · MTP <pct> accepted (<len> tokens/decode-step)
    arbi  no progress for <age> · req#<id> out <n>, <age> · <n> active · KV <pct>

Three rules decide what a line says, and they are the whole point of
this module:

1. **The metric matches the phase.** A request whose prompt is still
   being consumed is PREFILLING; its rate is prompt-tokens/s and its
   progress is ``consumed/total``. A decode rate is reported only for a
   request that is decoding. A zero decode rate over a long prompt still
   being consumed is not an uninformative line, it is a WRONG one — a
   reader concludes the server is wedged.
2. **A line carries what the previous line did not.** Every tick builds
   a MATERIAL KEY of quantized fields (phase, progress bucket, rate
   octave, KV band, in-flight counts, …). The line that states the
   invariants and the stall verdict fires when that key changes, or when
   the quiet backoff below elapses.
3. **The invariants are stated once.** The model path and the GPU
   used/total are constant for the process's life (memory is locked at
   boot). They render on the first line of a busy episode and whenever
   they actually change, not on every tick.

Quiet backoff: while requests are still running but the material key has
not moved, a keep-alive fires after ``QUIET_BACKOFF_START_S`` and then at
doubling intervals up to ``QUIET_BACKOFF_MAX_S``. It says which of the
two things a reader actually needs to tell apart is happening:

  - the request's counters ADVANCED (only the buckets are unchanged) —
    the ordinary line, whose token counts and elapsed are the new fact;
  - the request's counters did NOT advance — ``no progress for 2m00s``,
    which is what a wedge looks like from the log.

Either way "is it stuck?" is answerable without an admin endpoint, and
neither is a repeat of the previous line.

Levels: the heartbeat is DEBUG. A request's progress is not an event in
the server's life — the request's own access-log pair (``→ call#N`` /
``✓ done#N``, :mod:`arbi_serve.engine.run_step.finish`) is what INFO
carries, and a periodic line between those two would be a progress bar in
a log file. ``--log-level DEBUG`` turns the cadence on for anyone tracking
a specific request.

The one exception is a STALL — a request in flight that is not advancing.
That is a condition, not progress, and it goes to WARNING, where a real
problem can speak regardless of level.

Idle suppression: a tick with 0 active requests emits nothing and resets
the interval baselines, so the next busy tick measures from the moment
work resumed. A field whose source is unavailable is omitted cleanly —
never a fabricated zero.

The formatter (:func:`format_heartbeat`) is pure and unit-tested in
isolation; :class:`EngineHeartbeat` is the asyncio driver that samples
the engine, decides which line this tick gets, and calls it.

StatsMsg feed: each tick additionally accumulates its interval deltas
(generated tokens, MTP drafted/accepted/rows) into a counter bucket
drained by :meth:`EngineHeartbeat.take_counter_deltas`;
:func:`heartbeat_stats_msg` is the ready-made ``stats_provider`` for
``EngineProcRuntime`` — ``build_stats_msg`` with those deltas attached
(docs/engine_core_process.md §5.2).

The longest-running in-flight request — the one "is it stuck?" is about.

Carried on the heartbeat so a long request is legible from the log
with no admin endpoint: WHICH request, HOW LONG it has been in
flight, and HOW FAR through its work it is.

One interval's sampled runtime stats.

Every optional field is omitted from the line (never zeroed/faked)
when its source is unavailable: ``mtp_*`` when no MTP driver is
attached, ``kv_pct`` when the KV pool is unavailable, ``queued`` when
nothing is waiting, ``model`` when the served name is unresolvable,
and the GPU-memory segment when ``gpu_*`` is ``None``. So a minimal
line is just throughput + active count.

Render the one-line heartbeat, or ``None`` when nothing to emit.

Returns

``stalled_s`` replaces the rate head with ``no progress for <t>``: for
a request that has not advanced, the elapsed IS the content and the
rate is not.

The leading segment names the phase and reports THAT phase's rate: a
prefill interval reports prompt-tokens/s, a decode interval reports
output-tokens/s. A stats object with no phase information (bare
fixtures, callers that only sample throughput) renders the plain
``N tok/s`` it always did.

The MTP segment is appended only when both ``mtp_accept_rate`` and
``mtp_acc_len`` are present (driver attached + drafted this
interval). A missing ``kv_pct`` drops the ``KV X%`` field cleanly.
``mtp_spec_disabled_by_load`` adds its own segment, independent of the
accept segment — the interval it describes is exactly the one where
there is no accept data to render.

The segments a line states once per episode rather than per tick.

The model and the GPU total are fixed once memory is locked at boot.
GPU *used* is a live probe that drifts with allocator churn, so it is
keyed at the resolution the line PRINTS it — a key that moves at a
finer grain than the text can change while the rendered line stays
byte-identical, which is a duplicate the reader cannot account for.
Every key term here must therefore be derived from what
:func:`format_heartbeat` renders, not from the raw reading.

The quantized content of a line — what a reader would LEARN from it.

Two ticks with the same key say the same thing; the second one is the
spam this module exists to stop. Every field is bucketed at the
resolution declared by the module constants above, so the ordinary
tick-to-tick jitter of a decode rate does not manufacture a change.

The key is a function of what the line RENDERS. A term the renderer
can drop (the MTP segment, which needs an acc-len) or round more
coarsely than the key buckets it contributes nothing but duplicates:
the key moves, the text does not, and the reader sees the same line
twice with no way to tell what changed.

The lead request's RAW counters — advancing or frozen, unbucketed.

:func:`material_key` is deliberately coarse so ordinary jitter does
not manufacture a line. That coarseness cannot answer "is this
request moving at all", which is the one question a long-running
request raises, so the raw counters are carried separately.

Cumulative output tokens generated by the engine so far.

Sum of finished-request output counts (accumulated in
:func:`arbi_serve.engine.run_step.on_finished`) plus the live
in-flight requests' current output lengths. Exact across the
request lifecycle: a request that starts and finishes inside one
interval still contributes via the finished accumulator.

ENGINE-SIDE read: the heartbeat runs on the engine loop
(started by ``run_forever``), so reading the engine ``Request``
map here is loop-local — it is not a consumer-boundary crossing
(the process split relocates this into ``StatsMsg``, P5).

``{request_id: prompt tokens consumed}`` over the running set.

Keyed per request rather than summed, because a summed cumulative
would drop when a request leaves the running set and turn a finished
prefill into a negative rate. :func:`_prefill_advance` diffs per key
and sums only the advances.

``(prompt tokens consumed this interval, any baseline existed)``.

A request first seen this tick has no interval to measure over, so it
contributes nothing and does not make the rate look like zero. The
second element separates "not measurable yet" from "measured, and it
really is zero" — a stalled prefill must still be able to report a zero
rate, which is a fact; an unmeasured one must report nothing.

The longest-running in-flight request, rendered for the line.

"Longest-running" is by admission time: that is the request an
operator watching a slow log is asking about. Requests with no
``submit_time`` (fixtures, duplex lanes admitted off the intake path)
sort last and still resolve, without a fabricated elapsed.

The served model's name for the heartbeat line, or ``None``.

When stable-VA model cycling is engaged this is the ACTIVE resident's
served name (it flips on a switch); otherwise the configured
``served_name`` (falling back to the model path's basename — the same
resolution the /v1/models route uses). Never raises.

(used_bytes, total_bytes) for the engine's GPU, or (None, None).

``used = total - free`` from the driver — the same number nvidia-smi
shows, so the heartbeat and the card agree. Never raises.

Periodic sampler that emits the heartbeat line while busy.

Owns its own interval baselines (tokens + prompt-consumed + MTP
counters + wall) so every rate is an INTERVAL delta, and its own
material-key memory so INFO fires on change rather than on a timer.
Start with :meth:`start` from the run loop; stop with :meth:`stop`.

Build a ``StatsMsg`` with the heartbeat's accumulated counter deltas.

The ready-made ``stats_provider`` for ``EngineProcRuntime`` (P4b
wiring: ``stats_provider=lambda: heartbeat_stats_msg(eng)``): the
defensive :func:`~arbi_serve.engine.proc.stats.build_stats_msg`
snapshot plus whatever counter deltas the running heartbeat
accumulated since the previous send. Deltas drain exactly once
(idempotent-by-delta, §5.2); with no heartbeat running (mid-build,
worker ranks) the deltas are simply empty.

Drain the accumulated counter deltas (returns-and-clears).

Feeds ``StatsMsg.counter_deltas`` (deltas keep the message
idempotent, per the §5.2 contract). Swap-based, so it is safe to
call from the proc runtime's output IO thread while ticks keep
accumulating on the engine loop.

Spawn the periodic emit task.

An engine has AT MOST ONE heartbeat. Idempotence per instance is
not enough: each instance carries its own interval baselines and
its own emission memory, so two of them on one engine each render
a full, correct line from the same state and neither can see the
other's — the reader gets byte-identical duplicates, and the
per-tick side effects (the fragmentation gauge, the spec-disable
counter delta) are applied twice. Displacing the incumbent makes
that unrepresentable rather than merely unlikely: a run-loop
handover starts the new heartbeat before the departing loop's
``finally`` has awaited ``stop`` on the old one, and that overlap
is exactly the window the duplicates live in.

Stop ticking, without awaiting teardown (idempotent).

The synchronous half of :meth:`stop`, for a caller that is not a
coroutine. ``_run`` swallows the cancellation, so the task needs
no await to reach a clean end.

Choose this tick's one line and log it.

The change-triggered line (material key moved, or the quiet backoff
elapsed) carries the invariants and the stall verdict; every other
tick emits the plain per-tick detail. Both are DEBUG, except a
stall, which is a WARNING — see :meth:`_log_progress`.

Emit one line and record what it said.

A stall is a WARNING: the engine holds a request it is not
advancing, which is a condition an operator must act on and not a
report of ordinary progress. Everything else is progress, and
progress belongs to the request's own access-log pair, so it goes
to DEBUG where someone tracking a specific request can ask for it.

Engine-step infrastructure-failure classification + unhealthy latch.

An engine step can fail for two very different reasons, and they must not
be reported the same way:

  * **Per-request failures** (a malformed sampling combination, a model
    quirk on one input) affect one request. The run loop finishes that
    request with ``finish_reason="error"`` and the server stays healthy.
  * **Infrastructure failures** (CUDA OOM, ``cublasCreate`` alloc failure,
    illegal address, device-side assert) mean the engine cannot run
    forwards at all — every subsequent request will fail in milliseconds.
    Reporting those as per-request errors produces a *silent total
    outage*: HTTP 200 + ``finish_reason="error"`` + 0 tokens, which load
    balancers and naive clients count as success.

This module classifies a step exception (:func:`is_infra_error`) and keeps
a small latch on the engine (:func:`note_infra_failure` /
:func:`note_step_succeeded`): while the latch is set, ``GET /health/ready``
reports 503 ``engine_infra_failure`` so an orchestrator pulls the pod, and
blocking HTTP responses carry the reason. The latch self-heals: it clears
on the next successful engine step, or decays after
:data:`INFRA_UNHEALTHY_WINDOW_S` without a new failure so a drained pod can
re-admit probe traffic (which either succeeds → clean, or fails → re-latch).

This module owns the classifiers and the latches. It does NOT own the order
they are applied in — that is one seam shared by every run loop, and it lives
in :mod:`arbi_serve.engine.step_outcome`, which is the module a loop calls.

Logging is once-per-burst: the first infra failure logs ERROR with the full
traceback; while the burst persists, a counting summary is emitted at most
every :data:`INFRA_LOG_INTERVAL_S` (a sustained infra fault errors every
request — per-request exception logs would be pure spam).

Raised when an OOM landed with NOTHING LEFT TO NARROW.

Not a defect and not a poisoned context: every forward would be correct if
the card had the memory for one. It is raised only from the one state
:mod:`arbi_serve.engine.memory_pressure` defines as the end of its ladder —
an OOM while admission is already closed for memory pressure and no step has
succeeded since it closed — by which point :func:`note_memory_exhausted` has
already latched, so ``/health/ready`` is 503 before this reaches any loop.

It exists so a path that is not the run loop's own step (today: the drafter
seed draw, which cold-paths its bucket rather than failing the step) can
still propagate that verdict out of the loop for an orchestration restart,
without borrowing ``DrafterChainBreakerError`` — which would name a drafter
defect for a card that is simply too small.

True iff ``exc`` poisoned the CUDA context (irrecoverable).

An illegal memory access / device-side assert leaves the context
corrupt — subsequent forwards on it cannot be trusted and no
"successful" step can clear the damage. The caller latches a sticky
fatal fault (:func:`note_fatal_fault`) so the engine fails fast
instead of serving guaranteed-broken requests from a poisoned context.

Classified by message marker only: torch surfaces these as
``torch.AcceleratorError`` / ``RuntimeError`` with well-known
substrings (the device error is sticky on the CUDA stream).

True iff ``exc`` is a CUDA out-of-memory / allocator-exhaustion error.

A strict subset of :func:`is_infra_error` — an illegal-address / device-side
assert is infra-fatal but is NOT an OOM, and a memory-history snapshot of a
poisoned context is useless. Used to gate the dump-on-OOM hook so it fires
only where a snapshot names an escaping allocation, and to let a boot phase
catch NARROWLY: "this allocation did not fit" is a different fact from "the
code that was measuring is broken", and only the first may be absorbed.

Flag-gated, best-effort dump-on-OOM at the fail-loud detection point.

Called from the infra/fatal latch funnels — which every serving-forward OOM
BEFORE the loud teardown proceeds, so the snapshot captures the
escaping allocation. No-op unless the exc is an OOM and
``ARBI_SERVE_CUDA_MEMORY_HISTORY`` is on; NEVER raises (the dump can never
mask the original OOM).

Log the in-process memory ledger ONCE for a failed step, at ERROR.

Every OOM triage starts by asking which pool held what when the allocation
failed, and the answer is already in this process: the cuMem allocator's
per-tag ledger, cap rows and breach counters (pure bookkeeping, readable on
a poisoned context), the driver's free/total, and the live card. This
prints all of it in one record next to the failure, so the state does not
have to be reconstructed from a restarted process.

Best-effort and synchronous: each section is guarded on its own, so a CUDA
context that can no longer answer ``mem_get_info`` or the allocator
snapshot still leaves the bookkeeping sections printed, and nothing here
can raise into the latch that called it. Fires once per exception object
(the latch funnels that call it are each once-per-episode already), so a
failure that latches two funnels prints one card. The once-guard is a mark
on the exception itself rather than a reference held on the engine: holding
the exception would pin its traceback, and with it every tensor the failed
step had in scope.

Latch a sticky fatal fault — the CUDA context is poisoned, fail fast.

Set after an engine step hit a context-poisoning error
(:func:`is_context_fatal_error`: illegal address / device-side assert).
While set, ``/health/ready`` is 503 and inference requests are refused
(both route through :func:`engine_unhealthy_reason`). Unlike the infra
latch this never decays and is not cleared by a later step — a poisoned
context can falsely appear to "succeed". It clears only on process
restart. Idempotent: re-latching keeps the first reason and stays loud
once.

Latch a sticky memory-exhausted fault — the card cannot serve, fail loud.

Set by :mod:`arbi_serve.engine.memory_pressure` at the one point where its
ladder runs out: an OOM landed while admission was already closed for
memory pressure and no step had succeeded since it closed, so the reaction
was applied and did not work and there is no narrower slate to build.

Its OWN latch rather than :func:`note_fatal_fault`'s, because the two say
different things to an operator and both go on to be quoted verbatim: a
fatal fault says the CUDA context is corrupt and no forward on it can be
trusted, while this says every forward would be correct if the card had the
memory for one. Reporting the second as ``engine_context_poisoned`` would
send whoever reads it to look for a kernel bug instead of at the size of
the card. Like the fatal latch it is STICKY — a card is not made larger by
a later step appearing to succeed — so it clears only on process restart,
which is where a smaller ``--max-batch`` / ``--max-context`` /
``--gpu-memory-utilization`` can be applied.

Idempotent: re-latching keeps the first reason and stays loud once.

Latch an infrastructure step failure on ``eng`` (unhealthy until clear).

Callable from anywhere holding ``exc``, inside its ``except`` block or
not: the burst-opening log passes ``exc_info=exc`` explicitly rather than
relying on the ambient exception, because the SPMD driver classifies its
saved step exception after the cross-rank agreement reduce. Sets
``_infra_failure_reason`` / ``_infra_failure_ts`` /
``_infra_failure_count``; ``/health/ready`` reads them via
:func:`engine_unhealthy_reason`.

Clear the infra latch after a successful engine step (self-heal).

Does not touch the sticky fatal-fault latch (:func:`note_fatal_fault`):
a poisoned CUDA context can falsely appear to "succeed", so a fatal
fault clears only on process restart.

Latch a sticky admin-swap fault — serving is suspended until corrected.

An admin backend swap that cannot succeed (e.g. TKV with no calibration,
or it does not fit) must make the engine refuse to serve, not keep quietly
serving the previous backend: a remote bench that requested the swap would
otherwise attribute the old backend's results to the one it asked for. So
while this latch is set, ``/health/ready`` is 503 and inference requests are
refused (both route through :func:`engine_unhealthy_reason`).

Unlike the infra-failure latch this does not self-heal on a successful step
(the engine can still step on the rolled-back backend) and does not decay —
it clears only via :func:`clear_swap_fault`, i.e. when the admin corrects
the config and re-issues a swap that succeeds.

THE INSTRUCTION HAS TO MATCH THE STATE. "Re-issue the swap to resume" is
only true while a residency member is still mapped: the re-issued swap
parks it, builds beside it and switches. With NOTHING resident there is no
member to park, so the retry cannot even start — it refuses on the
residency guard with a message about inconsistent state, which reads as a
second, unrelated bug and sends an operator hunting one. When no member is
resident this says so and names the restart, because that is the only
thing that works.

Whether a residency member is mapped and could serve a re-issued swap.

Best-effort and deliberately permissive: an engine with no stable-VA
residency (``--no-cumem-pools``) never parks anything, so its member is
always the resident one and the ordinary retry instruction is the right
one.

The latched unhealthy reason, or ``None`` when healthy.

A sticky fatal fault (:func:`note_fatal_fault`, CUDA context poisoned)
takes top precedence and never decays — the engine refuses to serve until
the process restarts. A sticky memory-exhausted fault
(:func:`note_memory_exhausted`) is next and also never decays — the
narrowing ladder ran out, and a card is not made larger by waiting. A
sticky admin-swap fault (:func:`note_swap_fault`) follows, likewise without
decay — an impossible swap leaves the engine refusing to serve until an
admin corrects it. Otherwise the infra-failure
latch applies, with the :data:`INFRA_UNHEALTHY_WINDOW_S` decay so a pod
pulled from the load balancer (no traffic → no successful step to clear the
latch) eventually re-admits probe traffic instead of staying 503 forever on
a transient.

Option B (B2) — in-process single-capture, no-OOM at high GMU.

The problem this replaces: the cold-boot no-OOM path measured the cudagraph
capture pool in a subprocess (:mod:`arbi_serve.engine.capture_sizing`), which
loads the model twice, purely to learn how big the capture pool is so the
parent can size KV to leave room.

The fix — reserve-max-VA + partial-map + grow, in one process. Built on the B1
:class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion` primitive:

  1. Build the KV pool at the profiled max ``num_pages`` but GrowableRegion-
     backed: the slab VA is reserved at max (costs ~0 physical) while only a
     small capture-valid prefix of pages is physically mapped. Lots of free
     VRAM remains for the capture transient.
  2. Capture the full sweep once in-process. The captured graphs bake the KV
     slab's stable base VA + block_table indices. The synthetic capture inputs
     index only low pages (page 0 throughout — verified per backend: decode /
     prefill / layer all ``block_table.zero_()`` + ``slot_mapping`` zeros;
     recurrent ``state_indices = arange(B)`` index the separate recurrent_pool,
     not the paged KV slab) so they never touch an unmapped page.
  3. Measure the persistent capture pool consumed (free-VRAM low-water delta).
  4. Grow KV physical into the reclaimed space at the same reserved VA: map
     ``(free - floor - margin) // per_page`` more pages (capped at the profiled
     max). Captured graphs see the newly-mapped pages — the stable VA is the
     same property sleep/wake's ``wake_all`` remap relies on; here it is the
     grow axis. Post-capture free is >= floor by construction.

When cuMem is unavailable this path is not taken — build.py keeps the
subprocess measure. Not gated on ``vram_mode``: correct capture-then-grow KV
sizing runs identically for bench and serving (the bench-vs-serve axis is only
fail-loud vs warn on an incomplete engine, in
:func:`arbi_serve.engine.build.verify_post_capture_headroom`). Boot-only; zero
hot-path cost.

``value`` reduced with ``op`` across every rank (int64 all_reduce on
the default, world process group); identity at world_size 1 or when
distributed is unavailable (unit tests).

ONE reduce for every rank-collapsed boot quantity. A boot verdict is a
GROUP verdict — the slate rank 0 builds is executed by every rank — so any
input that varies per rank has to be collapsed before it decides anything.
Credits and debits take opposite ops off that one premise, which is why the
op is a parameter rather than two hand-rolled reduces: see
:func:`_all_rank_min_int` and :func:`_all_rank_max_int`.

Used for the KV grow target (block-table indices are rank-independent —
a divergent per-rank mapping leaves another rank's page index physically
unmapped → a GPU fault under load; mirrors the all-rank min floor in
:func:`arbi_serve.engine.profile.profile_and_size_kv_pool`) and for the
grow-verdict inputs (every rank must reach the same broken/constrained
outcome — a per-rank verdict could have one rank warn-and-serve while
another refuses the same boot, killing the server after "serving now").

Gated on ``world_size``, not ``tp_size``: every rank of the world
executes the same grow, so the agreement must span every rank, not
just the TP group. Behaviour at ``tp_size == world_size`` (``ep_size ==
1``) is unchanged — the TP group is the world there.

The smallest ``value`` across every rank.

The collapse for a CREDIT — free VRAM, an arena residency credit, a KV page
target — where the binding answer is the least any rank can offer.

The largest ``value`` across every rank.

The collapse for a DEBIT — a reserve, a hole, anything a step may NOT spend
— where the binding answer is the most any rank must hold. Binding a
reserve on the SMALLEST rank would let every other rank serve against bytes
it does not have, which is the same error as binding free VRAM on the
largest (``feedback_conservative_direction_flips``).

True iff the PAGED_KV backend's slab layout can be GrowableRegion
prefix-backed (rank-3 packed-bytes per layer, i.e. the TKV layout).

:class:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool` only honours
``growable=True`` for rank-3 layouts (pages on dim 0 → a byte-prefix is a
page-prefix); a rank-4 K/V pair layout has pages on dim 1 of the
stacked slab, so a "growable" request would materialize the full slab.
B2 must therefore never drive a rank-4 pair boot: its whole premise — the
capture transient lives in the unmapped tail, no graph-pool reserve
needed — is false for a fully-physical pool.

True iff the in-process single-capture (B2) path should drive this boot.

cuda_graphs + cuMem driver available + paged-KV with a known per-page cost
+ a slab layout that actually supports prefix-backing
(:func:`kv_growable_layout_supported` — a rank-4 K/V pair layout does
not, so those boots keep the subprocess-measure / reserve-from-plan path).
This supersedes the subprocess measure when on (no double model load); the
subprocess remains the fallback for the cuMem-unavailable case. Not gated on
``vram_mode``: bench and serving size KV identically (the only mode axis is
fail-loud vs warn on an incomplete engine).

Pages the GrowableRegion KV slab must physically map before the capture
sweep so no captured forward touches an unmapped page.

The synthetic capture batch indexes page 0 throughout — every backend zeroes
``slot_mapping`` and ``block_table`` at capture (verified in
:mod:`arbi_serve.runtime.capture.{decode,prefill,decode_graph}`:
``slot_mapping.zero_()`` + ``block_table.zero_()``), so the forward reads /
writes only the page-0 slot of the KV slab regardless of the captured
block-table width (``max_pages_in_table`` columns, all valued 0). The
determinism-gate / replay shapes likewise stay in the low range. So a
handful of low pages — the page-0 null sentinel plus one page per ``B`` row
plus a round-up margin — covers every captured backend.

It must not include a full ``max_context`` window (``ceil(max_context /
block_size)``): that term physically maps the entire per-sequence context
capture, which on a tight TP1 boot consumes the whole KV pool and leaves the
capture sweep no headroom — :func:`assert_capture_sweep_fits` then sees a
starved ``free`` and refuses, even though the capture only ever touches page
0. The capture transient + the post-capture KV grow
(:func:`grow_kv_after_capture`) handle the real KV sizing after capture; the
prefix only has to keep every captured/replayed page index mapped, which is
page 0 (+ a tiny margin), not a context window.

``arbi_serve/realtime/_duplex_admission.py``), re-derived here rather
than imported (that module is realtime-layer, this is engine-layer boot
sizing) from the same two ``BatchConfig`` fields other reserves in this
module already read directly off ``eng.cfg.batch``. Falls back to the
same 400s/0.08s defaults ``BatchConfig``/``_duplex_admission.py`` use
when either field is absent (a bare test double) or non-positive, so
this is never zero/negative.

Exact TTS paged-KV capacity reserved by both KV-sizing paths.

* :func:`serving_floor_for_grow` (this module) — the B2 in-process
  capture-then-grow path, gated on cudagraphs actually being captured
  (:func:`~arbi_serve.engine.phase2_freeze.should_freeze_for_serving`).
* :func:`arbi_serve.engine.profile.profile_and_size_kv_pool` — the
  generic pre-capture sizing path every boot runs, cudagraphs or not.

The hook prices the configured duration, page size, connection capacity,
speaker prefix, and two CFG rows. Zero means the served model has no TTS
paged runtime.

What the serving floor holds for ``scratch.forward_arena``, and why.

``bytes_`` is the reserve; ``provenance`` is one of
:data:`~arbi_serve.runtime.pool_taxonomy.PROVENANCE`. The pool's own
physical is the cover for what serving re-maps, so this row is 0 unless
the widest step admission can build needs MORE free VRAM than the base row
and the pool's residency hold between them — see
:func:`forward_arena_regrow_plan`. ``derived_bytes`` and ``resident_bytes``
are carried so the boot line and the operator card can show the two
measurements the row was resolved from rather than the result alone.

The arena's share of the widest slate admission can build. DERIVED.

:mod:`arbi_serve.scheduler.activation_budget` calibrates one step's
activation bytes as ``const + per_token x prefill_tokens + per_row x rows +
per_gather_kv_token x gather_kv_tokens`` against the boot profile's
measured shape peaks, lifting ``const`` until the fit reproduces every
profiled shape at or above its measurement. ``worst_admissible_bytes`` is
that model at the widest slate the scheduler's limits allow.

WHY THAT IS A BOUND AND NOT AN ESTIMATE. The three arguments are the ones
``Scheduler.schedule`` caps STRUCTURALLY, before any byte accounting: it
seeds ``token_budget = max_batched_tokens`` and ``batch_left = max_batch``
and stops appending rows the moment either reaches zero, and the gather
argument is clamped by the ceiling past which tkv takes the zero-copy
varlen route. The model is monotone non-decreasing in all three, so no
admissible slate evaluates above the corner. This does not rest on the
byte gate itself, whose one escape hatch — a solo row admitted over budget
because one row is the minimum unit of progress — is still a slate inside
that same corner.

The FOURTH argument is capped structurally too, and by the same kind of
limit. ``max_query_tokens`` is ``max_batch x K`` — every row opted into
speculation at this boot's resolved draft depth — and admission cannot
exceed it because it cannot exceed ``max_batch`` rows and no row's ``K``
exceeds the depth the drafter resolved. It is read through the one
expression ``arm_activation_admission`` uses, so the corner this reserves
for and the corner the boot log reports are the same number.

The profile's peaks are ALLOCATED bytes over every pool, so this bounds the
step's whole live set. ``base_floor_bytes`` is the part of it the floor's
base term holds (the profiled peak with the arena's own high-water taken
out), so the remainder is the arena's, and ``base + this`` is the step
bound exactly once. The pairing is what the guarantee is stated over: the
two terms are drawn from ONE pool of free VRAM, so a step that lands more
in the arena and less outside it is covered by the sum either way.

0 when the model cannot be calibrated — fewer than two profiled token
widths, or no profile at all — which leaves the reserve to the candidates
that are observations.

What the serving floor must hold FREE for ``scratch.forward_arena``.

THE POOL IS ITS OWN COVER. The EAGER serving forward runs inside that cuMem
pool (``EagerModelRunner._activation_arena_pool_ctx``), and boot runs the
same forwards through it — the activation profile probes at the widest
admissible shapes, the verify-width probe, the drafter pre-sweep. So by the
time the KV grow reads free VRAM the pool is already holding the physical
those forwards realized, at the block sizes they asked for. Those bytes are
RESIDENT: they are outside the free VRAM the grow divides up, KV was never
offered them, and a serving step of an admissible shape re-uses them rather
than taking them from this floor.

That is why this row is normally 0, and why the 0 is a measurement rather
than an omission. Boot does not forecast the pool's serving size; it
REALIZES it, and then keeps it.

The one thing residency cannot answer is whether the widest step admission
can build needs more FREE VRAM than the base row and the pool's residency
hold between them. :func:`forward_arena_step_bound_bytes` is that bound —
admission's own enforced model of one step, fitted to this boot's profile —
and the floor's pair closure raises this row by any shortfall against it
(see :func:`serving_floor_for_grow`). Both inputs are measurements of THIS
boot; neither is a forecast of another one.

WHAT THIS ROW DELIBERATELY DOES NOT DO is predict the pool's serving
high-water. A private ``MemPool`` keeps a block per distinct allocation
size it has served, so its reserve is a union over the shape mix a workload
presents, and the union has no closed form. Two things were tried and are
on record as failures: the pool's boot mark handed to KV and held back by a
bound (the bound is not one — under concurrency several rows are live at
once and the pool can map past any single probe's high-water), and a
reserve carried across processes in the budget cache (a reading taken by a
process that served one short request is indistinguishable from one taken
under load, and it under-reserved a production boot to the point of OOM).
Keeping the realization is the version that needs neither.

The serving arena watch
(:func:`~arbi_serve.engine.arena_watch.observe_forward_arena_reserved`)
reports the live pool against ``bytes_ + resident_bytes`` on every served
step, so a pool that does grow past what boot realized says so LOUDLY while
it happens rather than being predicted at a boot that could not know.

This configuration's recorded serving driver-side growth, or ``None``.

A configuration that has been watched serving has a number here, one that
has not gets ``None`` and the floor holds nothing rather than a stand-in.

What it covers is physical the DRIVER takes after the KV layout is frozen —
a cubin materialised the first time a served step dispatches a kernel no
boot phase launched, the kernel-stack pool re-grown for a deeper frame, a
graph instantiated post-boot. ``driver.modules_loaded`` brackets the BOOT
window and is then held, so none of it is inside that term.

The card REPORTS the first of the three as ``driver.modules_serving`` and
names the other two beside it, but a report is not a reserve and that one
only reads a number while something re-measures the serving bracket. This
is the reserve, and it is the superset: for a floor the three are one
quantity — free VRAM the grow handed to KV that the driver then took. See
:func:`~arbi_serve.engine.memory_budget.graph_pool.
predict_serving_driver_growth_bytes` for why it is measured rather than
derived.

The measurements behind the arena row, for the boot line.

The row is a number to take on faith without them: a reader has to be able
to see that the pool's own physical is what covers the step bound, and by
how much. The third clause appears once the boot has measured what the pool
costs SERVING rather than what the boot probes left in it
(:class:`~arbi_serve.engine.serving_realization.ServingArenaMark`) — the
residency this row is stated against is then a drive's reading rather than
an inheritance, and the two are indistinguishable from the row alone.

The arena row's operator-card note: why the row is what it is.

A 0 here is the normal answer and it is a MEASUREMENT — the pool realized
its serving working set during boot and kept it, so those bytes are
resident rather than free and this floor, which holds FREE VRAM, has
nothing left to hold for the pool. That is a different statement from "no
measurement was available", which this row can no longer be in: nothing it
reads comes from another boot or another process.

Bytes the serving floor holds back for ``scratch.forward_arena``.

The resolved reserve of :func:`forward_arena_regrow_plan`, which says where
the number comes from. Called without a base term, so the admission-derived
candidate is charged in full here; the floor itself passes its base term so
the two rows split the step bound instead of double-booking its overlap.

Measure the stochastic verify tail's per-step peak on demand.

The serving floor is never sized from a config-dims estimate: a grow that
lands before the activation profile exists (or against a profile that
predates the shape) measures the tail HERE instead of falling back. The
synthetic runs the rejection sampler on its own operands, not
``model.forward``, so it needs no KV pool, metadata builders or captured
graphs and is safe at any point a grow can happen.

Memoized on the engine: every later grow reuses the first measurement
rather than re-running a vocab-scale synthetic per call.

Measure the verify slate's ``lm_head`` epilogue peak on demand.

Runs the REAL head — the model's own ``logits_from_hidden``, epilogue terms
included — at :func:`verify_logits_slate_rows`, through the same two-pass
seam the boot profile uses. Nothing in the activation profile covers this
call: the decode and mixed probes gather ONE row per sequence before the
head, the verify-forward probe is ``model.forward`` (whose head is also the
per-sequence gather), and all three run inside ``scratch.forward_arena``
while the served epilogue allocates from the default pool.

Needs no KV pool, no metadata builders and no captured graph — one GEMM
over a zero hidden block — so it is safe wherever a grow can happen.
Memoized on the engine, like the verify tail's own measurement.

Rows ONE context-assemble can be handed, from the engine's own limits.

:meth:`~arbi_serve.spec_decode._dflash_driver_stash._DFlashStashMixin.
observe_seed_forward` projects each slate row's span SEPARATELY
(``feats[lo:hi]`` per request, in a loop), so the peak is the WIDEST single
span, not their sum — concurrency packs more rows into a step without
widening any one of them. Two limits bound that span and both are read
here rather than assumed:

  * ``chunk_prefill`` — the scheduler's per-row cap. Every prefill row
    takes ``min(remaining_prompt, chunk_prefill, budget_left)``
    (:mod:`arbi_serve.scheduler.scheduler`), so no row exceeds it.
  * ``max_batched_tokens`` — the step's whole token budget, which one row
    can spend alone (this GDN hybrid runs prefill depth-first FCFS).

The MIN of the two is the widest span that can exist. Admission's
activation gate may narrow a chunk BELOW it at serve time, which is a
tightening this reserve must not be sized from: the gate can be inert (no
profile to calibrate) while the scheduler's cap still stands.

Every input :func:`~arbi_serve.engine.memory_budget.
dflash_context_projection_peak_bytes` takes, read off the LOADED drafter.

``None`` when there is no DFlash drafter to read (no MTP driver, a native
MTP head, a worker rank's weightless stub) — which is a reserve of 0, not
an unmeasured one.

``num_kv_heads`` comes from ``local_num_kv_heads``, the TP-LOCAL count the
projection actually runs, and ``fused_kv_gemm`` from whether
``build_fused_kv_buffers`` produced a one-GEMM weight — an EXL3 drafter's
trellis tensors cannot be concatenated, so it keeps the per-layer loop and
a wider live set. Both are properties of the checkpoint that is loaded,
never of the checkpoint we expect.

The row count below which an EXL3 K/V projection changes GEMM leg.

A tile is a MEMORY decision and must not silently become a DISPATCH one:
``exl3_linear_gemm`` picks the trellis-inline leg or the reconstruct leg on
the row count, and the two differ in both speed and numerics. Flooring the
tile here keeps every tile on the leg the untiled call would have taken.

Zero (no floor) when the projection is dense — there is no leg to cross —
or when the threshold cannot be read, which only ever lets the tile be
smaller than needed rather than wrong.

Pin the DFlash assemble route and report whether the FALLBACK is dead.

The seam between the drafter's proof and the row that spends it. Returns
True only when :meth:`DFlashDrafter.freeze_graph_assemble` has refuted
every condition under which a draft step allocates a fresh per-step
context destination — the lever, the capture arming, the pool's shape
coverage and the empty slate — for the whole life of this member.

It decides BOTH halves of the row, which is why it is resolved once here
rather than asked twice: the probe runs the route the reserve covers (so a
proved member's probe keeps its graph pool and measures the direct route),
and the closed-form bound that stands in where the probe cannot run drops
the same destination. The drafter is told the answer
(``_draft_dest_unreserved``) so the branch the row stopped paying for
fails LOUD if it is ever reached, rather than allocating device memory the
KV grow has already handed out.

A drafter that cannot answer (no DFlash driver, a test double, a
``max_batch`` the config does not carry) is False: the destination stays
reserved and the row is byte-identical to before.

The closed-form bound on one DFlash draft step's transient peak.

ONE expression of the bound for every caller — the floor row and the tile
target both read it through :func:`_dflash_draft_step_reserve`. It is the
number the row holds when the step could not be RUN; it is not what the
row holds when it could.

The draft STEP's peak — the reserve the draft row holds anyway.

The tile target. The serving floor's draft row is a MAX over the draft
step and the context-assemble, so tiling the projection below this buys
no VRAM back: these bytes are reserved either way. Tiling TO it is what
makes the projection stop being the binding term.

Reads the SAME resolved reserve the floor row holds
(:func:`_dflash_draft_step_reserve`) — measured where the step could be
run, closed-form where it could not — so the tile is aimed at the bytes
actually reserved rather than at a second opinion about them.

The row tile the observe path runs in — MEASURED on this card.

WHY THIS IS NOT A DERIVATION. The obvious closed form is "tile to a byte
budget the floor already holds", and it is wrong physics. What actually
sets the optimum is where the projection's working set stops fitting the
GPU's cache and where per-call overhead (an EXL3 tile re-runs the trellis
reconstruct) starts to dominate — properties of THIS card, this drafter's
per-row bytes and this dtype. None of them move with ``max_batch``, and
every byte-budget candidate in this engine does. GPU-measured on one RTX
4090 at the served drafter: the optimum sits at 2048 rows, and the
103,072 B = 2046) doubles to 4101 rows at ``max_batch`` 8 — off the
optimum, for four times the reserve, with no mechanism connecting the two.
That agreement was a coincidence of magnitude and is not load-bearing.

So the tile is measured the same way the reserve is: time the FULL row
bound at each ladder rung and take the fastest. The whole sweep is one
projection's worth of work per rung and costs ~2 s at boot.

THE TIE-BREAK FREES VRAM. Rungs within the sweep's own observed spread of
the best are indistinguishable in time, so the SMALLEST of them wins — the
reserve is linear in the tile, so a tie is bytes for free. The spread comes
from the measurement (the widest min-to-max seen across rungs), never from
a chosen tolerance.

Falls back to the closed-form tile, loudly, if the sweep cannot run; a
boot with no timing is still better off tiled than not.

Bytes the serving floor must hold for ONE draft context-assemble.

THE ROW THAT WAS A CONSTANT. ``project_context_stacked`` runs after every
target forward, outside the forward arena and outside every named pool, at
a PREFILL CHUNK's row count — so it is the one DFlash allocation that
scales with ``chunk_prefill``, and nothing in the budget scaled with it.
A reserve calibrated at one chunk width silently under-reserves at a wider
one, and the draft then OOMs INSIDE a serving forward.

MEASURED first, always. The drafter is loaded and the projection is pure,
so the honest number is the one the loaded drafter produces at
:func:`_dflash_observe_row_bound` rows — the same standing as every other
row in the itemisation. The analytic bound sizes the row ONLY when that
measurement could not be taken at all, and the row then says ``DERIVED``
and names the reason.

THE TWO ZEROS ARE DIFFERENT FACTS and are reported as such:

  * no drafter to read (``geometry is None``) — the reserve is genuinely
    0, there is no context-assemble on this boot;
  * the probe ran out of memory — the reserve is UNKNOWN and the card is
    too full to find out. That is not a 0 and it is not a smaller
    measurement: it RAISES, naming the shape, the row bound and the knobs
    that move it, so the boot refuses cleanly instead of OOMing inside a
    forward later. Sizing the row from the analytic there would let an
    estimate stand in for a measurement the card actively refused.

The measurement cannot be blocked by its own reserve: it is taken at the
post-capture grow, BEFORE the KV pool is grown into the free VRAM, and
nothing consults this number until it returns.

Keep the first resolved reserve for the life of the engine.

ONE measurement per boot, and it has to be the FIRST call's, because that
is the emptiest the card ever is. :func:`serving_floor_for_grow` is called
again from ``reclaim_kv_before_freeze`` — AFTER the KV pool has been grown
into the free VRAM this very reserve sized — where a chunk-scale probe
cannot fit by construction and would refuse a boot the same measurement had
just approved. Re-measuring at a seam the reserve itself narrowed is a
forecast that invalidates its own input; memoizing is what stops it.

Run the LOADED drafter's context projection at ``rows`` and return its
peak allocated bytes; ``None`` when the probe could not be run at all.

The features are allocated OUTSIDE the profile window on purpose. A served
observe reads them as a VIEW of the persistent tap slab
(``_collect_tap``), so counting a fresh copy of them here would charge the
reserve for bytes the step never allocates — the same double-book class
the activation floor was just fixed for.

Raises :class:`~arbi_serve.runtime.activation_profile.ActivationProbeFailed`
when the probe OOMs. See :func:`_dflash_observe_reserve` for why that is a
refusal rather than a fallback.

The WIDEST draft step this boot can be handed: ``(B, max_full_len,
per-layer context caps)``; ``None`` when there is no DFlash drafter.

Not a guess at the worst case — the SAME shape the drafter's own capture
pre-sweep enumerates as its top bucket
(:func:`~arbi_serve.spec_decode.dflash_capture.dflash_presweep_shapes`
takes ``max_batch`` and ``slots._layer_len``), so the step this reserve is
measured at is a step the boot has already proven it can build. Both
inputs are read off the objects that decide them rather than recomputed:

  * ``B`` is ``batch.max_batch`` — admission never builds a wider slate,
    and the draft forward captures at the EXACT slate width.
  * the caps are ``DraftKVSlots._layer_len``, each layer's whole slab.
    ``_assemble_fixed_context_device`` narrows a FULL-attention layer to
    the caller's ``max_full_len``, so the widest step is the one handed
    the widest slab; passing it explicitly is also what keeps the probe
    off ``_host_width_bound``, whose page-table read has no answer for a
    request that does not exist.

Why a draft step must NOT be run as a boot probe here, or ``None``.

A draft step's embed and lm_head legs are the TARGET's vocab-parallel
modules under TP>1, and the sharded drafter build carries row-parallel
all_reduces of its own. Every one of those rendezvouses with peer ranks
that are NOT running this probe, so issuing one here is a half-issued
collective at boot — an NCCL wedge, not an exception. The floor keeps the
closed-form bound there and says so; a conservative row is a cost, a
wedged boot is an outage.

Bytes the serving floor must hold for ONE DFlash draft step.

MEASURED first, always — the sibling discipline of
:func:`_dflash_observe_reserve`, and for the same reason: the drafter is
LOADED, the step is a pure function of its context buffers, and the widest
shape it can be handed is already enumerated by the capture pre-sweep. A
closed form standing in front of a step that can simply be run is a
forecast where there is an observation.

WHY THE CLOSED FORM EXISTED. The bound prices the FALLBACK route — the
step that finds no covering graph and so allocates its own ``(B, nkv,
cap_i, hd)`` destination — at a shape no boot used to run. It is
enumerative (destination + assemble scratch + per-layer K/V concat +
folded mask + MLP pair + lm_head), and an enumeration of a forward's
temporaries prices every tensor a reader can name, whether or not the
kernels allocate them. That is the half this measurement replaces.

The probe runs THE ROUTE THE RESERVE COVERS. A step whose shape a captured
graph covers assembles into that graph's own buffers — resident before the
grow reads free VRAM, and deliberately outside this row — so the bytes
this row holds are the fallback's. The probe therefore detaches the graph
pool for its window, which is not a different code path but the same
:meth:`~arbi_serve.spec_decode._dflash_driver_context._DFlashContextMixin.
_assemble_and_draft` taking the branch it takes whenever no graph covers.

NOT CACHED ACROSS BOOTS, and that is the keying answer rather than a gap
in it: the reading is a function of the drafter checkpoint, the slab
geometry, the batch width, the dtype and the card, and every one of those
is already resolved on the object being run. Measuring the loaded drafter
each boot keys the number on the boot itself; a persisted entry would add
a way for the key to disagree with the drafter.

The two zeros are different facts: no drafter (or no DFlash path) is a
genuine 0, while a step that could not be run leaves the row at the
closed-form bound with ``DERIVED`` provenance and a stated reason.

Keep the first resolved draft-step reserve for the life of the engine.

The sibling of :func:`_memo`, and for the same reason: the first call is
the emptiest the card ever is, and ``serving_floor_for_grow`` is called
again after the KV pool has been grown into the VRAM this very reserve
sized.

Run ONE draft step on the loaded drafter at the widest shape and return
its peak allocated bytes; ``None`` when the step could not be run at all.

The step is the SERVED step, not a re-enumeration of it: the same
``_assemble_and_draft`` entry point, the same embed / lm_head modules
(:meth:`~arbi_serve.spec_decode.dflash_driver.DFlashDrafter.
draft_head_modules`), the same ``(B, block)`` mask block, the same
``dhid_all`` gather and the same lm_head width. What it does not do is
realize a block — sampling reads the logits this window already counted.

The slots are BORROWED, at the widest committed length the slab can hold:
a draft step's destination width comes from the slab and its key positions
from ``lengths``, so a step measured against empty slots would measure a
step whose context is all pad. They are released on every exit path.

Raises nothing: a probe that cannot run is a refusal to LOWER the row, and
the closed-form bound (an upper bound over the same tensors) keeps the
floor safe. That is the opposite of the context-assemble probe, whose
closed form is the same order as its measurement and whose OOM therefore
means the card cannot afford the step at all.

A profiled shape's peak with the forward arena's share taken out.

Delegates to :meth:`~arbi_serve.runtime.activation_profile.ShapeProfile.
default_peak_bytes`. A shape object that predates the split — a restored
profile, an engine double — has no such method and keeps its WHOLE peak in
the base term, which is the conservative direction: the floor then reserves
the arena's bytes twice rather than not at all.

Physical ``scratch.forward_arena`` already holds. MEASURED, at this seam.

The pool's own reserved-bytes reading. It is the one number that says how
much of a step's activation cost is payable from memory the process has
ALREADY mapped: a private cuMem pool never returns a freed block
(``empty_cache`` skips it), so what the boot probes mapped stays mapped for
the process's life.

Two callers need it and must read it once: :func:`_arena_covered_bytes`,
which caps how much of a profiled peak may leave the floor's base term, and
:func:`serving_floor_for_grow`, which credits the same bytes to admission's
per-step budget. Two reads could disagree, and then a byte would be shed
from one term without being credited to the other.

Arena bytes the serving floor holds OUTSIDE its base term.

The cap on what :meth:`~arbi_serve.runtime.activation_profile.ShapeProfile.
default_peak_bytes` may take out of the base term. A byte only leaves that
term if something else already holds it, and for ``scratch.forward_arena``
that is one of two things: the pool owns the physical its boot forwards
realized, so those bytes are resident and not in the free VRAM this floor
measures; or the floor's ``arena_regrow`` row holds free VRAM on top of
that residency, which it does only where the enforced step bound exceeds
what the base row and the residency cover (:func:`forward_arena_regrow_plan`).
Those two are DISJOINT, so the pool's capacity is their SUM, and it is the
same sum the serving arena watch checks the live
pool against. A MAX here named only the larger half and left the base term
holding the smaller one a second time.

The map round-up slack the serving floor holds, from the driver.

Asks the cuMem allocator for its device allocation granularity and bounds
the slack at the two boundaries a round-up can move. Falls back to the
stand-in constant only when no driver answers.

What ONE serving step costs on a boot with no per-shape profile.

``bytes_`` is the base-term input the floor uses in place of the per-shape
max; ``provenance`` is one of
:data:`~arbi_serve.runtime.pool_taxonomy.PROVENANCE`, and ``note`` says on
what authority — the row is read by an operator deciding whether this
boot's KV size is a fit or a guess.

``bytes_ == 0`` is the answer "nothing about this step was measured", and
it is not a small floor: it is the state in which no floor can be derived
at all, and :func:`grow_kv_after_capture` declines to grow rather than
claim VRAM against it.

OPTIMISTIC ONE. It is smaller than a real decode step, so a floor that
collapses to it lets the post-capture grow claim VRAM a serving step needs,
and the shortfall surfaces as a client-visible 500 after the layout freeze
rather than as a boot refusal. Its own definition scopes it to "the
granularity / fragmentation floor", which is a different quantity from a
step's activations and cannot stand in for one.

So the floor falls back to a MEASUREMENT, not to a bigger constant and not
to a closed form. ``boot_state.profiled_activation_peak_bytes`` is the
Phase-3 activation probe: a real forward at ``max_batched_tokens``, run by
``profile_and_size_kv_pool`` and MAX-reconciled across ranks, recorded

  * It is at least as WIDE as any step the scheduler can assemble. The
    probe runs at ``max_batched_tokens``, and
    :func:`~arbi_serve.runtime.activation_profile.reachable_step_tokens`
    bounds a real step at ``min(max_batched_tokens, max_batch x
    chunk_prefill)`` — so its token count is a bound on the serving
    shapes' token counts, and token-scaled GEMM scratch is what dominates
    a forward's activations.

    WHAT THAT IS NOT: a proof it dominates in BYTES. ``SHAPE_MIXED``
    carries the same tokens across MORE ROWS, and per-row metadata scales
    with rows, so an equal-width mixed step can hold a little more than
    this single-sequence probe. The margin that covers the difference is
    the arena double-book below, which on an arena-heavy config is worth
    far more than the row residue — and it is a consequence of what was
    measured, not a factor anybody chose. This is a same-width
    measurement standing in for a per-shape max, and the row says so
    rather than claiming a bound it does not have.
  * It is parametrised on everything it depends on by being a measurement
    of the real model at the real width — not a formula with a factor in
    it. The closed form
    :func:`~arbi_serve.engine.memory_budget.predict_activation_reserve_bytes`
    keeps for the pre-profile path is not usable here: at the reference
    geometry it comes out several times UNDER the measured mixed peak, so
    substituting it would replace an optimistic constant with an
    optimistic formula.

THE WHOLE PEAK STAYS IN THE BASE TERM. Nothing measured the arena's share
of this probe, and :meth:`~arbi_serve.runtime.activation_profile.
ShapeProfile.default_peak_bytes` already fixes what an unknown arena share
means: the shape keeps its whole peak, and the floor reserves the arena's
bytes twice rather than not at all. The over-reserve is the point — this
boot trades KV for a floor it can stand behind — and the boot says so.

Returns ``bytes_ == 0`` when even that probe left no number. There is then
no measurement of any width to size a floor from, and the honest move is
not to grow.

``(bytes, provenance, note)`` for the free VRAM serving consumes.

MEASURED once boot has RUN the widest step this configuration can issue
(:func:`~arbi_serve.engine.serving_realization.realize_under_serving_pressure`):
what that step was watched take is what the floor holds free for the next
one. The drive's own bytes do not stay resident and cannot be netted off —
the settle that follows it hands the caching allocator's freed blocks back
to the driver, so they are inside the free VRAM the grow divides up and a
served step takes them again.

WHICH MAKES THE CONDITIONS OF THE READING PART OF THE READING, and the note
carries them. A drive on an empty card takes fresh segments a step under
serving pressure never gets; the realization therefore re-drives against
the card the grow is about to leave and it is THAT number that arrives
here.

Otherwise the labelled pre-measurement bound, and the note says which boot
could not realize and why. The note is what the ledger row and the boot log
both print, so an operator reads the same sentence in both places.

The true serving free-VRAM floor the post-capture KV grow must leave.

Derived from the max over the non-prefill profiled serving shapes (decode /
mixed / the stochastic MTP verify tail) — not the capture-time prefill peak
(``activation_reserve``), which is reclaimable. The stochastic verify tail is
vocab-scale and scales with concurrency (see
:func:`arbi_serve.engine.memory_budget.serving_activation_floor_bytes`); the
profiler measures it directly so the floor self-corrects.

Fails loud when stochastic MTP is enabled but the verify tail was not
measured: sizing KV against a decode-class floor would let a concurrent
stochastic verify OOM the card at serve time (a client-visible 500), which
is exactly the silent degrade we refuse to ship.

Pre-capture prediction of the post-capture realized servable KV pool.

``boot_state.kv_serving_ceiling_pages`` is a deliberately generous VA cap:
it subtracts the serving activation floor folded as a max
(:func:`arbi_serve.engine.memory_budget.profile_time_serving_floor_bytes`),
so the growable slab reserves enough address space to fill genuinely-free
VRAM. But the post-capture grow (:func:`grow_kv_after_capture`) holds the
same peak-serving-step transients out additively
(:func:`serving_floor_for_grow`), so the realized pool lands well below that
ceiling. The pre-capture viability gate
(:func:`arbi_serve.engine.build_memory_sizing.assert_kv_context_window_fits`)
that reads the raw ceiling therefore never sees the shortfall and its
``max_context=auto`` narrow never fires — the boot then dies at the
post-capture :func:`assert_kv_pages_floor` with no narrow escape.

De-rate the ceiling by exactly the additive-minus-max floor delta the grow
will hold out, so the gate predicts the realized pool. The ``gpu_memory_
utilization`` headroom the grow keeps free (``serving_floor_for_grow``'s
``gmu_floor``) is already represented in the ceiling's ``usable = total ×
gmu`` base, so difference against ``max(plan_floor, gmu_floor)`` — never
double-count gmu (a low-gmu boot would otherwise be spuriously narrowed).

Returns the ceiling unchanged when the delta cannot be priced (no per-page
cost, no stored plan floor, or the grow floor is unavailable) — this only
ever lowers the ceiling, and only for the auto-narrow decision, so an
unpriceable prediction can never false-refuse.

Return every reclaimable byte to the driver, so the KV sizing that
follows reads a free-VRAM value that nothing is about to change.

Three distinct mechanisms hold reclaimable physical after the capture
sweep, and each needs its own release — an ``empty_cache`` alone returns
none of the first two:

  1. **The forward arena's idle overhang.** The
     ``scratch.forward_arena`` MemPool ballooned to the boot prefill-width
     peak; at serving it re-maps only what a served step puts there.
     ``empty_cache`` does not visit a private ``MemPool``'s block pools
     (pytorch#145168), so only destroying the pool frees them —
     :func:`release_idle_forward_arena` gives up the declared per-step
     memos (:mod:`~arbi_serve.engine.arena_step_caches`), then destroys the
     now-empty pool and recreates it. It REFUSES on anything else still
     live, by name: re-homing a buffer a captured graph may have baked is
     not sound, so the physical stays resident instead. This does NOT run
     at boot: the pool's boot high-water is the serving working set, so
     boot keeps it (:func:`forward_arena_regrow_plan`).
  2. **Uncollected Python cycles.** The capture sweep runs with the
     collector paused (see ``_gc_paused`` in ``build.py``), so CUDA tensors
     dropped during the sweep can still be reachable from unbroken cycles
     and keep their allocator segments alive.
  3. **The default allocator's free segments.** ``empty_cache`` returns the
     wholly-free ones.

It also clears the lazily-pinned cuBLAS workspaces first, so that residency
does not depend on which cuBLAS handles the capture sweep happened to touch.

Ordering is load-bearing: the arena release and the ``gc.collect`` both
hand blocks back to the caching allocator, so the ``empty_cache`` must come
last to return them to the driver.

Idempotent (the second call finds nothing to release) and latches
``eng._kv_sizing_settled`` so :func:`grow_kv_after_capture` can assert that
it ran. Best-effort per stage: a release that fails leaves fewer bytes for
KV, never an incorrect engine.

What the post-capture KV sizing holds the KV page ceiling against.

``planned_bytes`` is the plan's input (``boot_state.
serving_torch_caching_overhang_bytes``) and ``measured_bytes`` is what this
boot's settled card holds; ``bytes_`` is the monotone bound over the two.
``measured`` says whether a reading was taken at all — a boot that could
not measure keeps the plan and must not be reported as having agreed with
it.

MEASURE the serving torch-caching overhang on THIS boot's settled card.

The quantity the plan booked as the activation footprint's overhang form,
read where it is directly observable rather than forecast. The Phase-2
freeze reads the same meter and persists it for the NEXT boot at this
configuration; that record is a plan seed and a reconcile, and it is
structurally too late to size anything this boot maps
(``freeze_for_serving`` runs after the grow).

WHICH SIDE OF THE RECLAIM. Called from :func:`grow_kv_after_capture` AFTER
:func:`settle_vram_for_kv_sizing`, on the same settled card the single
``mem_get_info`` is read from. That is the side that makes the reading a
bound on bytes KV can never have:
:meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.
torch_caching_reserved_overhang_bytes` subtracts the default pool's
FULLY-free segments, which are exactly what the settle's ``empty_cache``
hands back to the driver for the grow to map as KV. Measured before the
settle, the reading would carry those bytes and the ceiling re-cut would
hold out physical the grow is about to convert — silently shrinking KV on
every boot.

Returns the plan's value with ``measured=False`` when no registry / no
reading is available. Never raises: a boot that cannot measure keeps the
plan it already had.

The page ceiling BOTH post-capture KV sizing passes bound their target by.

One definition, called by :func:`grow_kv_after_capture` and by
:func:`reclaim_kv_before_freeze`, because a ceiling that lives in two
places is two ceilings: the second pass re-derives the same number and any
cut applied to only one of them is undone by the other.

Three terms, in order:

1. **The slab's OWN reserved VA** (``pool.num_pages``), not
   ``kv_serving_ceiling_pages``. That ceiling is FORECAST-derived — it
   subtracts a profile-time serving floor whose ``caching_overhang`` term
   is not measurable at the moment it is computed — so bounding the grow
   with it lets a guess decide how many pages serving gets, and a
   pessimistic guess costs KV that is physically free and already
   VA-backed. The reserved VA is the sleep/wake/swap-safe quantity, and it
   is the bound ``grow_kv_to_fit`` (the swap / sleep-wake reclaim) already
   uses for the reason it documents: both ``boot_state`` ceilings are
   engine-GLOBAL and a residency member's variant build clobbers them DOWN
   without restoring them, while ``pool.num_pages`` is per-pool and fixed
   at boot. Co-residence stays safe by construction — a genuine pool-mate
   never gets a generous reservation in the first place (``profile.py``
   skips the solo VA widening when ``pool_cap_pages`` binds the member to
   its slot), so the slab's own VA is already the reduced bound.
2. **The deferred-calibration re-cap.** On a deferred-calibration boot
   (``eng.cfg.inprocess_calibration``) the pre-capture resize deliberately
   shrank the pool to the calibration basket's footprint so the centroid
   fit gets working VRAM instead of a pool filled to the serving ceiling.
   Growing back to the serving ceiling here would silently undo that and
   OOM the very next Lloyd-Max fit. A no-op on every normal serving boot.
3. **This boot's caching-overhang bound**
   (:func:`measure_serving_caching_overhang` →
   :func:`~arbi_serve.engine.memory_budget.kv_budget.
   kv_ceiling_pages_for_overhang_bound`), which takes back the pages the
   pre-capture plan cut against a forecast of a quantity the boot has since
   MEASURED. Identity until that measurement exists, and identity on every
   boot whose budget cache already carries it.

4. **The OSCAR basis's measured residency**
   (:func:`~arbi_serve.engine.build_memory_sizing.
   kv_ceiling_pages_for_attn_codec`). ``scratch.attn_codec``'s plan row
   prices that pool's TQ buffers and not the per-KV-head rotation matrices
   an OSCAR install puts beside them, so the plan handed KV those bytes.
   Identity on every boot with no rotation configured.

None of the four can over-commit on its own: the target is separately
bounded by MEASURED ``free_bytes - serving_floor_bytes``
(:func:`~arbi_serve.engine.memory_budget.kv_budget.kv_grow_target_pages`),
so a ceiling only decides how much of that measurement KV may keep.

Grow the GrowableRegion-backed KV pool into the VRAM the capture
transient freed, at the same stable VA the captured graphs baked.

Runs in Maps ``(final_free - serving_floor) // per_page``
more pages (capped at the profiled max), so the gate's ``free >= floor``
holds by construction. No-op when the KV pool is not growable.

THE single post-capture KV sizing: one ``mem_get_info`` read, one
``grow_kv_to_pages``. :func:`settle_vram_for_kv_sizing` runs first and
completes every release that returns physical to the driver (the
forward-arena overhang, the swept cycles, the caching allocator's free
segments), so the free this reads is the final one — nothing after it
hands memory back, and the realized page count is a function of the
config alone rather than of when a release happened to land.

Calling the settle here, rather than at the call site, makes that ordering
an invariant of the function instead of call-site discipline. The passes
that follow (:func:`reclaim_kv_before_freeze`) are safety nets that must
find nothing.

The serving floor held out is the decode-class serving peak, not the
reclaimable capture/prefill peak.

Emit the single realized-KV summary + capacity ledger, and learn the
prediction residual, against the MEASURED servable pool.

Runs once after the last post-capture grow (grow_kv_after_capture plus the
pre-freeze reclaim re-grow), so the pages reported and the residual learned
are the true serving pool, never an intermediate grow.

Safety net over the single post-capture KV sizing: assert that no
reclaimable VRAM appeared after it, and say so loudly if any did.

:func:`grow_kv_after_capture` sizes KV once, immediately after
:func:`settle_vram_for_kv_sizing` has completed every release that returns
physical to the driver. This pass re-reads ``mem_get_info`` and re-runs the
identical arithmetic against the identical floor
(:func:`serving_floor_for_grow` — never redefined) AND the identical
per-target price list
(:meth:`~arbi_serve.cache._growable_kv.GrowableKvMixin.kv_mapped_bytes_for_pages`).
Identical is the whole detector: a pass that priced a page at capacity
while the grow paid the per-layer map cost reports the gap between two
arithmetics as a release, on any boot whose page count happens to land far
enough from a granularity boundary. It MUST find nothing.

Finding growable pages here is a PHASE-ORDERING DEFECT, not a windfall:
some release ran outside the settle barrier, which means the realized KV
size depends on when that release happened to land rather than on the
config. So this fires the must-not-fire ``kv_sizing_late_growth`` counter
and logs at ERROR, naming the byte count. It still claims the pages —
stranding VRAM as frozen-free would be a second bug on top of the first —
but it never does so silently.

The claim itself is safe at any point: ``grow_kv_to_pages`` →
``GrowableRegion.map_to`` maps more physical at the RESERVED base VA, not a
realloc, so captured-graph KV pointers stay valid.

No-op when KV is not growable or when capture did not drive this boot.

Every live tensor whose storage sits in ``pool``'s mapped segments,
grouped by storage.

Returns ``{storage_data_ptr: (storage_nbytes, [tensors sharing it])}``. The
grouping is by STORAGE (not tensor) so a view and its base migrate together
onto one new buffer, keeping their shared-offset invariant. Read via a
``memory._snapshot`` walk (the active-allocated block ranges of the pool)
plus a ``gc`` scan for CUDA tensors landing inside them — the same seam
:meth:`NamedPoolRegistry._debug_dump_pool_leak_owner` uses, here to MOVE the
owners rather than just name them.

One rendered allocation stack per live block in ``pool``, or ``[]``.

THE REFUSAL'S OTHER HALF, and the one that works where the first does not.
:func:`~arbi_serve.engine.arena_tenants.owner_paths` answers "what still
REFERS to this storage" by walking the heap, and it correctly refuses to do
that on the engine loop thread — so the serving re-tightening's refusal
names every tenant ``<unattributed>``, which is the shape of a report that
cannot be acted on. This answers the other question, "where was this block
ALLOCATED", off the allocator's own per-block trace: no heap walk, no gc,
nothing that holds the GIL for longer than formatting the strings.

Empty unless the allocation-history recorder is armed
(``ARBI_SERVE_CUDA_MEMORY_HISTORY``, which the boot arms once at the
serving transition). That is not a limitation to work around: the recorder
hooks every device alloc/free and must stay off in every benchmark arm, so
the honest behaviour is to report nothing rather than to arm it on a
refusal — by which time the allocation is long past and no frame exists to
recover. A refusal that says ``<unattributed>`` and lists no sites is
telling the operator which flag to boot with.

Blocks are GROUPED by their stack, because the shape that matters is "N of
these, one per attention layer" rather than N copies of one stack.

Log the refusal to release an arena that still holds live storages.

Names every live storage AND its gc owner path, so the buffer that should
not be arena-resident can be fixed from the boot log alone.

``memos_dropped`` is how many declared arena memos
(:mod:`~arbi_serve.engine.arena_step_caches`) gave up an entry on the way
here. It is on the refusal because the two readings answer different
questions and only one of them is visible anywhere else: a refusal with a
non-zero count says the drop RAN and something else still holds the pool,
while a refusal with zero says nothing was declared for what is holding it.
Without the count the two are the same line.

Hand ``scratch.forward_arena``'s physical back, if it holds nothing live.

THE SERVING RE-TIGHTENING'S ONE MECHANISM
(:func:`~arbi_serve.engine.arena_watch.note_engine_idle`). ``None`` when
the release did not happen — a guard refused, or there was nothing
meaningfully idle to hand back.

It does NOT run at boot. The pool's boot high-water is the working set the
eager serving forward re-maps, not an idle tail, so handing it back before
the KV grow would promise KV and the step budget bytes serving takes
straight back (:func:`forward_arena_regrow_plan`).

``empty_cache`` skips a private ``torch.cuda.MemPool`` (pytorch#145168) and
the driver reuses a segment unmapped out from under it (illegal access), so
the only sound release is to destroy the pool once it is EMPTY
(:meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.
release_empty_pool` — release + recreate fresh under the same name).

SAFETY — the arena must hold NOTHING live, and that is a property of the
MOMENT, not of the phase. A captured graph bakes the device pointer of
every tensor it reads into its recorded launch parameters; anything still
live in this pool survived a step boundary, i.e. is persistent, i.e. is
exactly the class a graph can have baked. Re-pointing it (``Tensor.set_``
onto a fresh storage) or unmapping the pool beneath it leaves those graphs
dereferencing freed virtual addresses, which surfaces as an illegal access
on the first replay. So a non-empty arena is REFUSED, loudly
(:func:`_refuse_reclaim_live_arena`), never migrated — and the refusal is
the same check at boot and at serving, which is what lets the serving
caller run at all: an engine between steps holds no arena tensor unless
something parked one there, and the check is what says which it is.

``keep_bytes`` is the level the serving floor priced the arena's row
against — the pool's physical is part of that cover, so the release
RE-REALIZES it into the fresh pool before returning. Only the overage
above it is a net hand-back, which is the whole of what this repair
ever claimed to reclaim; releasing below it hands the driver the bytes
the floor's arena row is the complement of, and the next wide step has
nothing to re-map from.

The cuBLAS workspaces are cleared first: cuBLAS pins one per
``(device, stream, handle)`` on the first GEMM inside a pool and holds it in
C++ with no Python tensor behind it, so an un-cleared workspace alone keeps
the arena non-empty. A later GEMM re-pins one lazily.

THE SAME ARGUMENT COVERS THE PYTHON-SIDE MEMOS. A consumer that memoises an
allocation it made in this pool holds a block across the step boundary, and
one block forfeits every wholly-free segment in the pool — the release is
whole-pool by construction, so the memo's SIZE is irrelevant to what it
costs. Those memos are declared in
:mod:`~arbi_serve.engine.arena_step_caches` and dropped here, before the
live set is read. Each entry states why the next step cannot hit the entry
it gives up; a buffer whose contents a later step reads back is NOT one of
them and still refuses below, by name.

Map ``keep_bytes`` back into the fresh arena and hold it. Returns bytes.

A cuMem pool keeps a freed block MAPPED — the property that strands the
idle overage in the first place — so allocating once and freeing leaves
the physical resident and available to the next step without any tensor
outliving this call. That is what restores the residency the floor's
arena row was priced as the complement of, while the overage above it
stays with the driver.

A failure to map it back is the dangerous direction and is logged as an
error rather than swallowed: the pool is then below the level the floor
accounted for, which is the state that turns a wide step into an OOM.

Engine lifecycle: model-GPU teardown, request cancel, and reload.

This module handles the teardown / cancel concern so the per-step loop
(:mod:`arbi_serve.engine.run_step`) stays focused on dispatch. Nothing
here runs on the decode critical path:

  * :func:`cancel` / :func:`_cancel_inner` — drop a request mid-flight
    (client disconnect / SIGTERM drain), marshaled onto the engine loop.
  * :func:`release_lora_ref` — per-request LoRA refcount drop.
  * :func:`_release_model_gpu_state` — the coordinated GPU teardown shared
    by full :func:`shutdown` and the in-process hot-swap reload.
  * :func:`shutdown` — terminate the run loop + free GPU state.
  * :func:`release_for_reload` — teardown + pool re-creation for an
    in-process model swap (run loop stays up).

The per-token output glue (``_emit_finish`` / ``_flush_output`` /
``on_finished``) stays in :mod:`arbi_serve.engine.run_step`;
:func:`_cancel_inner` reaches it via a function-local import to keep the
two modules acyclic (run_step does not import this module).

Remove ``req`` from the queue/scheduler and tear down its detok task.

``cancel`` is invoked from the HTTP side (client disconnect / SIGTERM
drain). It mutates engine-owned scheduler state, so under
``ARBI_ENGINE_OWN_THREAD`` it must run on the engine loop — otherwise
``scheduler.remove`` races the run loop's ``scheduler.schedule()`` on
the engine thread. The bridge marshals the whole body onto the engine
loop (inline when single-loop / already on the engine thread).
``getattr`` default tolerates a partial / stub engine; a real serving
Engine always has the bridge.

Cancel a request by id — CancelMsg semantics (P3 §5.1).

The HTTP-path cancel surface: routes hold only the request id (in
process mode the engine ``Request`` never crosses), so the lookup
happens engine-side. Marshaled onto the engine loop exactly like
:func:`cancel`. An unknown id is legal — request ids are allocated
API-side, so a cancel can race a not-yet-published ``SubmitMsg`` —
and is recorded in the engine's cancel tombstones; the
RequestFactory's publish consumes the tombstone and drops the racing
submission with a ``cancelled`` FinishOut. Double-cancel is
idempotent (second call finds no request and leaves a tombstone that
expires unused).

Engine-loop body of :func:`cancel` (see its docstring).

The terminal sequence itself is
:func:`arbi_serve.engine.run_step._finish_request` — the ONE path
every "drop a request out of the engine" site runs, so a cancel frees
exactly what a completion, an error finish and a timeout free
(graceful-queue entry, scheduler pages / recurrent slot, LoRA ref,
detok task, CPU-prep arrays) and emits the same terminal
``FinishOut``. ``scheduler="remove"`` drops the request from whatever
queue holds it WITHOUT a radix commit; the detok worker is torn down
before the ``FinishOut`` so no further text lands behind the finish
(its undecoded tail is deliberately dropped on cancel — the
``FinishOut``'s authoritative token count still reflects every
committed token).

Drop every GPU allocation tied to the current model + free the
named pools — the coordinated teardown shared by full
:func:`shutdown`, the in-process model hot-swap reload
(:func:`arbi_serve.engine.swap_admin.areload_model`), and a failed
:func:`arbi_serve.engine.build.build` (a fresh boot OR a stable-VA
residency member build re-entered on an already-serving, SHARED
primary engine — see :func:`arbi_serve.engine.stable_va_pool_builder.
build_member_into_engine`).

Does not touch the run loop, the ``_shutdown`` flag, or the detok /
prep / forward executors — those are model-independent infrastructure
owned by the caller. It does, in the proven ``shutdown`` order:
tear down the sleep pool, clear every captured cudagraph, drop
attn_ops / metadata_builders / pool / page_table / scheduler / model,
clear the module-global registries that pin per-layer blocks + cores
(custom-op layer dispatch, EXL3 inner registry, the release
registry), drop the cuBLAS/cuBLASLt workspaces + Dynamo cache, and
``free_all`` every NamedMemPool so its VRAM returns to the driver —
then empties the caching allocator. Best-effort + idempotent: any
single failure is logged and the rest proceeds. This is what lets
``shutdown`` reach ``memory_reserved() ~ 0``; a reload calls it and
then rebuilds (re-creating the named pools + sleep pool first).

Snapshots the ``id()`` of every ``torch.cuda.MemPool`` reachable from
``eng.named_pools`` BEFORE any teardown mutates it, so the
:func:`reap_leaked_pools` call at the end can scope its
``_CUDAGRAPH_POOL_PINS`` release to exactly THIS engine's own pools
(see that function's docstring for why an unscoped clear is unsafe on
a shared, multi-member engine).

Signal the run loop to exit and drop GPU state cleanly.

Sets the shutdown flag, then proactively drops every GPU
allocation the engine still owns:

  - drops the captured-graph pool (releases the cudagraph mem
    pool back to the driver)
  - releases sleepable state (param.data + KV slabs + TKV
    scratch via :meth:`SleepableTensorPool.release`)
  - tears down the SleepableTensorPool (cuMemAddressFree each VA
    + frees pinned host pool)
  - empties the caching allocator so reserved bytes return to
    the driver

Idempotent — calling twice is safe (second call short-circuits).
Test harnesses can rely on ``torch.cuda.memory_reserved()``
returning to ~0 after this call without manual ``gc.collect()``
/ ``torch.cuda.empty_cache()`` chains.

Tear down every stable-VA resident record before the final release.

:func:`_release_model_gpu_state` frees the current model's state, but a
multi-model pool holds N records whose snapshots (``rec.extra["state"]``)
pin each member's named pools, sleep pool (VA reservations + pinned host
backups) and instantiated cudagraphs. Left alone, those references keep
the pools above ``use_count() == 1`` so ``free_all`` parks them in
``_LEAKED_POOLS`` (still resident), which can OOM the next engine boot in
the same process (e.g. two live-GPU tests in one pytest run).

For each parked member: install its snapshot on the engine and run the
same coordinated :func:`_release_model_gpu_state` (its sleep pool is
already released — the physical lives in host backups — so the teardown
frees its VAs + pinned host pool and resets its instantiated graphs),
then drop the record's snapshot references. The active member's snapshot
is reinstalled last so the caller's normal release acts on the live
model. Best-effort: any failure is logged and the rest proceeds. No-op
when residency is not engaged.

Tear down weight pools deferred by the drain's weight-share ordering.

Runs after the final ``_release_model_gpu_state`` (the active sharer's
module is dropped), so each deferred donor ``weights_pool`` now holds
only dead allocations. Mirrors ``NamedPoolRegistry.free_all``'s per-pool
release-then-empty_cache sequence; a pool that stays pinned parks in
``_LEAKED_POOLS`` (the established safety net — reclaimed at process
exit). No-op when nothing was deferred.

Tear down all model GPU state for an in-process hot-swap reload, then
re-create the engine-lifetime pools so :meth:`Engine.build` can rebuild.

Unlike :func:`shutdown` this leaves the run loop + ``_shutdown`` flag
untouched (the server stays up across the swap), but it does recycle the
forward-offload worker thread. That thread caches per-thread CUDA state
(current stream + cuBLAS handle) and the last forward's transient refs;
if it survives the teardown it can pin a just-freed NamedMemPool below
``use_count() == 1``, forcing :meth:`NamedPoolRegistry.free_all` to park
the pool in ``_LEAKED_POOLS`` (still resident). The post-rebuild headroom
gate then sees a doubled budget and can refuse the reload. A fresh
thread starts from clean CUDA TLS, so ``free_all`` reaches every pool.

Steps, in order:
  1. recycle the forward executor (join old, spawn fresh);
  2. :func:`_release_model_gpu_state` — the coordinated GPU teardown that
     frees the model + every NamedMemPool (same path ``shutdown`` proves
     reaches ``memory_reserved() ~ 0``);
  3. re-create the named MemPools (:meth:`Engine._init_named_pools`) and a
     fresh :class:`SleepableTensorPool` — both are ``__init__``-built and
     destroyed by the teardown, but ``build()`` only ``get()``s pools, so
     without this rebuild every cache pool would fall through to the
     default allocator and the sleep-pool registration would hit a closed
     pool.

Returns GPU free bytes after teardown (logging only).

The LIVE card, as rows that close — the serving-time twin of the boot ledger.

:mod:`arbi_serve.engine.vram_ledger` closes the card ONCE, at the phase-2
freeze. This closes it on every read of ``GET /v1/admin/memory/live``, from the
same rule and the same taxonomy:

    device total = Σ resident rows + Σ free rows

with nothing left for a consumer to derive. That is the whole point of the
module. A consumer that has to compute ``total − free − Σ rows`` itself is
computing a quantity the server can name and the client cannot: every row the
payload omits lands on it, silently, under whatever label the client picks. The
rows the live payload used to omit — the state pools' sentinel-alias arenas,
and the driver sub-identity's own residual — are exactly the ones a client-side
subtraction booked as "driver residual" on a server reporting residual 0.

Three groups come out, and the split is the contract:

``rows``
    Resident physical. The caching-allocator pools, the two cuMem regions that
    drive the driver directly (``state.attn_kv.mapped``,
    ``state.gdn_recurrent.mapped``), the measured ``driver.*`` brackets, and
    ``driver.residual`` — computed HERE, from the card down, exactly as the boot
    ledger computes its own.
``free_rows``
    The unallocated remainder, itemised: the serving-step reserve the KV grow
    held back (``transient.serving_step.*``), whatever a live step has already
    taken back out of it (``transient.serving_step.in_use``, the one negative
    row), and genuinely-idle headroom (``unclaimed.no_owner``).
``ghost_rows``
    Rows that hold NO physical and take part in no sum — the address space of a
    pool this process destroyed. Carried in the payload rather than left for the
    client to know about, because a client that does not know books it twice.

When any contribution is UNRESOLVED the identity is not closed and must not be
made to look closed: no ``driver.residual`` row is emitted, ``closes`` is False,
and the gap is reported as ``unaccounted_bytes``. A breakdown missing its
largest row still sums to 100% if something absorbs the difference, which is
the failure this module refuses to reproduce.

Pure arithmetic + naming: no CUDA, no engine state, fully CPU-testable. The
caller (:func:`arbi_serve.engine.proc.admin_registry.memory_live_snapshot`)
does the measuring.

One rendered row: name, category, TENURE, bytes, live subset, PROVENANCE.

``provenance`` defaults to the taxonomy's value for the row and is
overridden only where the runtime knows better than the static default
(a reserve that fell back from a measurement to an analytic bound).
``tenure`` says what KIND of number the row is — physical that exists, a
reservation nothing has taken, address space that is not memory — and is
NOT overridable: it follows the category, so a row cannot claim a kind its
category does not have. A row whose category this build does not know
carries an empty tenure rather than a guessed one.
``leak_signal`` marks the one row whose non-zero value means an allocation
with an owner nobody declared, so a reader never has to remember which of
the residual-shaped rows is accounting and which is a defect. It is keyed
on that row's LIVE bytes, not on its name and not on its reserve: a row
holding zero live bytes has no allocation, so nothing about it is
undeclared. Its reserve can be non-zero and mean nothing -- the caching
allocator keeps empty segments outside every named pool, and on a healthy
boot this row is exactly that. Flagging on the name alone lit the glyph on
every boot, which is the failure the ``unexpected_provenance`` note below
describes: a mark that is always on is one a reader learns to skip.

``unexpected_provenance`` is the one provenance fact worth a GLYPH rather
than a hover: a row the taxonomy declares MEASURED that is reporting a
guess on this boot. It is deliberately narrower than "the provenance is
flagged" — ``unclaimed.allocator_slack`` and
``transient.serving_step.token2wav`` are declared constants and carry SEEDED
on every healthy boot, so marking those teaches a reader to skip the mark,
which costs the marks that matter. Computed here, once, because a console
that re-derived it from the taxonomy would be the second copy that comes to
disagree.

Assemble the closing row set for one device reading. Pure arithmetic.

Args:
    total_bytes: ``cudaMemGetInfo`` total — the card the rows must sum to.
    free_bytes: ``cudaMemGetInfo`` free, from the SAME call.
    pools: the caching-allocator snapshot, ``{name: {allocated_bytes,
        reserved_bytes, ...}}``. ``address_space.released_pool_va`` is routed to
        ``ghost_rows`` and left out of every sum.
    kv_mapped: :meth:`~arbi_serve.engine.memory_budget.driver_residency.
        MappedKvRow.as_dict`, or ``None`` from a backend that serves no such
        block. Only the ``mapped`` state contributes a row; ``pooled`` and
        ``absent`` are measured zeroes whose bytes are elsewhere.
    state_arena: :meth:`~arbi_serve.engine.memory_budget.driver_residency.
        StateArenaRow.as_dict`, the same contract for the state pools'
        sentinel-alias arenas.
    unmapped_tracked_bytes: ``CuMemPoolAllocator.tracked_bytes -
        mapped_bytes`` — the MEASURED total of allocations this process has
        ``cuMemUnmap``'d and ``cuMemRelease``'d while torch still counts
        their VA as reserved. It is the evidence behind ``ghost_rows``
        holding no physical: a ghost row larger than this is address space
        whose pages the allocator does not agree it gave back, which is real
        resident memory wearing a label that says it is not.
    driver_measured: the measured ``driver.*`` rows. ``driver.residual`` is
        IGNORED here and recomputed from the card down, so the row is the
        residual of the WHOLE identity rather than of the driver
        sub-identity — the two differ by every non-driver row, and reporting
        the sub-identity's would leave that difference unnamed.
    step_reserve_bytes: the serving-step reserve the KV grow held back
        (``serving_floor_for_grow``). It is carved off the free bytes so
        ``unclaimed.no_owner`` names only what nothing has claimed, which is
        the operator gate.
    step_terms: that reserve ITEMISED, from
        :attr:`~arbi_serve.engine.boot_state.EngineBootState.serving_step_terms`.
        Empty gives one un-itemised row rather than a fabricated breakdown.
    unresolved: ``[{name, reason}]`` for every contribution the caller could
        not establish. Non-empty means the card does NOT close and no
        residual row is emitted.

Every way ``card`` fails to close, one line each; empty means it closes.

The verdict a reader of the ``card`` block reaches by hand, written once so
a gate (the e2e lane's boot check) and a console agree on what "closes"
means:

* nothing ``unresolved`` and ``closes`` set -- every contribution was
  established;
* ``rows_bytes + free_rows_bytes == total_bytes`` exactly -- the identity
  holds with no gap;
* ``driver.residual`` inside :data:`RESIDUAL_ALARM_BYTES` -- the remainder
  is the measurement's noise floor, not an unnamed resident;
* no row of any group holding bytes under ``UNBUDGETED`` provenance -- a
  measured allocation no plan line predicts, or a ghost row the allocator
  does not agree holds nothing.

Pure: reads the payload and computes nothing the payload does not already
carry, so it runs wherever the JSON does.

Logprobs reduction for the ``logprobs`` / ``prompt_logprobs`` surface.

Turns hidden states into per-position ``LogprobEntry`` records without
ever holding a ``(positions, vocab)`` tensor for more than one tile: the
lm_head runs on ``tile`` rows at a time and each tile is reduced to the
requested top-k plus the target token's own logprob before the next tile
is computed. Every allocation happens inside the ``scratch.logprobs``
named pool, whose reserve is predicted by
:func:`~arbi_serve.engine.memory_budget.predict_logprobs_pool_bytes`.

Nothing here runs unless a request opted in.

The ``scratch.logprobs`` pool context, or a nullcontext.

A missing pool means the engine was built without the surface armed;
the caller is only reached when a request opted in, and allocating
outside a named pool would land in ``unpooled.torch_default_pool``
which the VRAM ledger reports as unbudgeted.

Reduce one ``(rows, vocab)`` logits tile to what was asked for.

Returns ``(target_logprobs, target_ranks, top_ids, top_logprobs)``
with shapes ``(rows,)``, ``(rows,)``, ``(rows, k)``, ``(rows, k)``.

The log-softmax is taken in fp32 over the raw logits — the full
distribution, NOT the post-top-k/top-p masked one the sampler works
in, because a reported logprob has to be the model's probability for
that token and not a probability conditioned on the sampling filter.

``rank`` is the token's 1-based position in the full vocab ordering,
computed as ``1 + #{v : logprob[v] > logprob[target]}``, which needs
no sort.

``dtype=`` on the log_softmax rather than ``logits.float()`` — the
fused form upcasts inside the kernel, so the tile's fp32 image exists
once instead of twice, which is what the pool reserve is sized for.

The model's arbitrary-rows lm_head seam.

Archs that do not route their head through the shared layer-stack
mixin have no way to score a non-last position, and a silently
skipped prompt would read as a short logprobs array rather than as
an error.

A cached ``token_id -> str`` callable for the engine's tokenizer.

Detokenization stays engine-side (the API child never re-tokenizes),
and a prompt scores every position, so the same id recurs constantly
— the cache turns a per-position decode into a dict hit.

Score this prefill chunk's slice of ``req``'s prompt.

``hidden`` is the chunk's ``(chunk_tokens, hidden_size)`` post-final-
norm activation; ``start`` is the flat row where this request's rows
begin in it and ``count`` how many rows are the request's.

The hidden state at prompt position ``p`` predicts the token at
``p + 1``, so this scores prompt positions ``consumed+1 ..
consumed+count`` and drops the request's final row, whose prediction
is the first GENERATED token and belongs to the sampled-token lane.
Results append to ``req.prompt_logprobs_acc`` in prompt order across
chunks, so a chunked prefill accumulates one contiguous list.

Score the token each row actually sampled.

``logits`` is the ``(rows, vocab)`` last-position logits already
computed for sampling and ``sampled[i]`` the id row ``i`` emitted, so
this is a pure reduction — no second forward. Rows that did not ask
for logprobs get ``None`` and cost only the list slot.

Score the tokens a speculative step committed, each at its own position.

A verify step commits a whole accepted prefix at once, so a lane that
scores one token per step is short by the acceptance length — which the
completions route rejects rather than serialise as a silently shifted
array. ``rows[i]`` is the flat slate position that PRODUCED ``tokens[i]``,
so each committed token is scored against the distribution at its own
position rather than at the step's last one.

Reads hidden states rather than the verify logits: the sharded verify path
never materialises a full-vocab row, so scoring from logits would serve
only the gathered path. The head runs over the committed rows alone, in
tiles, inside the ``scratch.logprobs`` pool.

Two-loop bridge for the dedicated-engine-thread mode: intake marshaling.

Background
----------
The legacy single-loop mode (``ARBI_ENGINE_OWN_THREAD=0``) runs the
engine's ``run_forever`` loop as an ``asyncio.create_task`` on the same
event loop as uvicorn (main thread). The engine's per-step Python
bookkeeping holds the GIL/loop and starves uvloop's I/O poll, so a
streamed first SSE chunk sits buffered until the step yields (the c=1
TTFT gap).

By default (``runtime_flags.engine_own_thread``) the engine loop runs on
a dedicated OS thread with its own asyncio loop; uvicorn keeps the main
loop. Two boundaries exist between them:

  * **Intake** (HTTP loop → engine loop) — owned here. An ``asubmit``
    coroutine runs on the HTTP loop but its publish step
    (``scheduler.add`` + ``_wakeup.set``) mutates engine-owned state and
    must run on the engine loop. :meth:`LoopBridge.run_on_engine`
    schedules it there via ``call_soon_threadsafe``. Under a burst (a
    c64 arrival storm lands N publishes in a ~1 ms window)
    :meth:`LoopBridge.run_on_engine_coalesced` buffers them on a
    thread-safe intake deque and fires a single coalesced wakeup that
    drains the whole burst — N eventfd wakeups collapse to ~1.

  * **Output** (engine loop → HTTP loop) — owned by the typed output
    pipeline (:mod:`arbi_serve.engine.output_bus` staging the
    :mod:`arbi_serve.engine.proc.messages` structs into one
    ``OutputBatchMsg`` per step, applied by the
    :class:`~arbi_serve.engine.client_request.OutputApplier`). The
    bridge only binds the bus's thread marshaler to the HTTP loop at
    :meth:`activate` and reverts it (flushing residuals) at
    :meth:`deactivate`; it never carries output callbacks itself.

Invariant: all CUDA work stays on the engine thread. The bridge only
moves cheap intake publish callbacks (``scheduler.add`` +
``_wakeup.set``) across the loop boundary — never tensors, never a
forward.

Off path (legacy single-loop, ``ARBI_ENGINE_OWN_THREAD=0``):
:attr:`enabled` is ``False`` and the intake helpers call the callback
inline on the current thread — byte-for-byte the single-loop behaviour,
zero added overhead, no extra task/hop. The output bus stays on its
inline marshaler for the same reason.

API-side pending-SubmitAck futures, keyed by ``request_id`` (P3).

``asubmit`` registers a future on its loop before marshaling the
submission onto the engine loop; whoever observes the factory
outcome calls :meth:`resolve` — thread mode calls it from the engine
loop (this class owns the one cross-loop hop, keeping the AST-linted
``call_soon_threadsafe`` inside the marshaler seam), inline mode
resolves on the same loop, and the process transport (P4b) calls the
same method from its recv thread when the ack message arrives. The
resolution is transport-agnostic by construction: nothing but the
``request_id`` and an ok/:class:`SubmitRejected` outcome crosses.

Thread-safe: ``register``/``resolve``/``discard`` may run on
different threads (HTTP loop vs engine loop vs recv thread).

Holds the engine + HTTP event loops and marshals intake between them.

A single instance lives on ``engine._loop_bridge``. It is created
disabled (single-loop mode); the dedicated-thread launcher calls
:meth:`activate` once both loops exist. ``output_bus`` (when given)
is bound/unbound alongside so the typed output pipeline switches to
its thread marshaler in lockstep with intake.

Resolve every pending ack with ``error`` (engine death, §7).

Process mode (P4b): when the engine process dies, submissions
whose acks are still in flight would otherwise hang forever —
the wire will never carry their ``SubmitAckMsg``. Callable from
any thread; each future resolves on its own loop.

Resolve the pending ack from any thread (ok xor error).

A missing entry is logged loudly (an ack for a request nobody is
awaiting means the intake plumbing lost track) but never raises —
the resolver may be the engine loop's intake drain, which must
not be poisoned by one bad ack.

Drop loop references (shutdown). Idempotent.

Reverts the output bus to its inline marshaler first — its unbind
flushes any residual staged batch inline so a trailing finish
emitted by the last step still reaches an attached consumer (the
HTTP loop may already be tearing down; best-effort). Then drains
any intake publishes buffered but not yet handed to the engine
loop (a burst that arrived as the loop tore down), inline, so the
request is published rather than silently dropped.

Run ``fn(*args)`` on the engine loop (admission publish).

OFF / single-loop: call inline.

ON: if already on the engine loop, call inline; otherwise
schedule via ``call_soon_threadsafe``. The callback is the
scheduler publish (``scheduler.add`` + ``_wakeup.set``) which
mutates engine-owned state and must run on the engine thread.

``is_closed()`` is a check, not a guarantee: the engine thread can
close its loop in the window between that read and the call
(:meth:`~arbi_serve.engine.engine_thread.EngineThread._thread_main`'s
``finally``). :meth:`notify_http` already handles exactly this
race; this path did not, and would have raised ``RuntimeError``
into whichever HTTP handler happened to be submitting. Falling
back to the inline call is the same thing the ``eng is None``
branch above already does.

Publish an intake callback, coalescing a burst into ONE wakeup.

The INTAKE analogue of the per-step output batch. A c64 burst
lands 64 ``asubmit`` publishes on the HTTP loop in a ~1 ms window;
routing each through :meth:`run_on_engine` fires 64 separate
``call_soon_threadsafe`` engine wakeups (eventfd write + callback
schedule). Instead, append ``(fn, args)`` to a thread-safe intake
deque and fire a single ``call_soon_threadsafe`` only when no drain
is already scheduled — so N near-simultaneous arrivals collapse to
~1 engine wakeup that drains the whole deque
(:meth:`_drain_intake`).

OFF / single-loop: run inline (the single-loop behaviour). ON but
already on the engine loop (a publish re-entrant from the engine
thread, e.g. a resume): run inline — buffering would defer it past
this turn and there is no cross-loop hop to save.

Lock-free handoff / race analysis
---------------------------------
``_intake_lock`` is held only for the cheap append + compare-and-set
here, and for the swap-out + flag-clear in :meth:`_drain_intake` —
never while a publish closure runs (``scheduler.add`` runs outside
the lock on the engine thread). The two critical orderings:

  * **No double wakeup.** The first appender flips ``_intake_pending``
    False->True and schedules the hop; later appenders in the same
    burst see it already True and only append. One hop per burst.

  * **No stranded publish.** The drain swaps the deque out and clears
    ``_intake_pending`` under the lock before running the swapped-out
    publishes. A producer that appends after the swap sees
    ``_intake_pending`` False and schedules a fresh wakeup, so its
    request is drained by the next hop — never lost to a "pending"
    flag that referred to an already-drained batch. The swap also
    means a publish appended during a drain accumulates into a fresh
    deque flushed by its own hop, never split mid-iteration.

  * **No wedged flag.** ``call_soon_threadsafe`` raises
    ``RuntimeError`` if the engine thread closes its loop in the
    window after the Left
    unhandled that is worse than a lost request: the append has
    already happened and ``_intake_pending`` is already ``True``,
    so every LATER publish sees a drain "already scheduled" that
    no one will ever run, and the whole intake path goes silently
    dark. The fallback runs :meth:`_drain_intake` inline, which
    clears the flag and publishes the buffered callbacks — the
    same recovery :meth:`deactivate` performs.

Run ``fn(*args)`` on the HTTP loop — the engine→consumer
NOTIFICATION hop, for a consumer whose readiness signal is not a
typed output-bus message.

``OutputBus`` remains THE path for a request's tokens/audio/finish
(one typed batch per step, gap-checked ``seq``, applied by
``OutputApplier``). This method is deliberately narrower: it
carries no payload, only "go look" — the wakeup half of a
consumer that owns its own bounded buffer on the engine side. The
duplex lane (:mod:`arbi_serve.realtime.duplex_lane`) is the first
such consumer: its per-tick events are raw waveform tensors that
the ``TokenOut``/``AudioOut(codes)`` structs do not model, and
teaching ``OutputApplier`` about them would add branches on the
hot path for all ordinary traffic to serve a shape no ordinary
request has.

The boundary invariant (``tests/test_transport_boundary.py``) is
preserved because the crossing itself happens HERE, inside the
seam, not at the caller. Callers must keep the "one hop per step,
never per item" discipline the output bus already enforces.

Returns ``True`` when the callback was scheduled (or run inline),
``False`` when there is no live HTTP loop to notify — the caller's
buffer keeps the data for a final drain either way.

OFF / single-loop, or already on the HTTP loop: call inline.

Await ``coro`` ON the engine loop; return its result here.

The coroutine-marshaling primitive for the admin/utility plane
(docs/engine_core_process.md §6): admin RPCs — including
``CriticalSection`` drain/mutate ops — must execute on the
engine loop so their lock/drain semantics match process mode,
while the HTTP handler awaits the outcome on its own loop.

OFF / single-loop (or the bridge's engine loop is gone): await
inline — one loop, nothing to marshal. Already on the engine
loop (an engine-side caller): await inline. Otherwise schedule
via ``asyncio.run_coroutine_threadsafe`` and await the wrapped
future; cancellation of the HTTP-side await propagates to the
engine-loop task (``wrap_future`` is two-way), so an abandoned
admin request does not leak a running drain.

Drain the whole intake deque on the engine loop (one hop/burst).

The single ``call_soon_threadsafe`` callback fired by
:meth:`run_on_engine_coalesced`. Swaps the deque out + clears the
pending flag under the lock (so a concurrent producer re-arms a
fresh wakeup for anything appended after the swap), then runs every
buffered publish in FIFO order outside the lock. Each publish is
isolated so one failure cannot strand the rest of the burst.

Engine-side LoRA wiring: target-key discovery + admin attach hooks.

Each parallel linear advertises a ``lora_target_key`` which the per-step
batch builder uses to pack ``(A, B)`` tensors only for relevant linears.
Admin endpoints route ``POST /v1/loras`` and ``DELETE /v1/loras/{name}``
through :func:`aload_lora` / :func:`aunload_lora`.

Map each LoRA target module name → its TP parallel kind.

``"row"`` for input-sharded row-parallel linears (o_proj / down_proj),
``"column"`` for output-sharded column-/merged-column linears
(q/k/v/gate/up). The PEFT loader uses this to slice ``weight_a`` /
``weight_b`` to each rank (see ``peft_loader.shard_lora_target``).

Admin: load + register a LoRA adapter (POST /v1/loras).

At TP>1 each rank loads its own shard of the adapter (the store on
every rank carries that rank's ``tp_rank``). Gated on the
``ARBI_LORA_TP`` flag: off by default it refuses loud rather than
serve an unvalidated per-rank shard split at TP>1.

Returns the loaded adapter's JSON-ready metadata dict (the
``POST /v1/loras`` response body) — never the live adapter object,
which must not cross the engine boundary
(``docs/engine_core_process.md`` §5.3).

Cross-rank commit/abort confirmation for a TP>1 member build.

A live config-override that needs a new captured member runs the same
park-build-capture-switch sequence on every rank in lockstep (the delta
is broadcast rank-0-first; both ranks route to build-required
deterministically). The build itself runs NCCL collectives (capture
warmup all-reduce), so both ranks reach them together by construction.

The remaining hazard is an asymmetric outcome: one rank's build fails
(per-rank OOM, a build-policy refusal) while the other succeeds. A
locally-succeeded rank would commit the switch while the failed rank
rolls back — leaving the deployment desynced. This module is the
all-or-nothing barrier that closes that gap: after each rank attempts
its build, every rank reports a local success flag and all ranks learn
the cluster verdict; either all commit or all roll back.

It is an ``all_reduce(SUM)`` of a per-rank ``1``/``0`` flag over the TP
group. This is not on the per-step hot path — it fires only on the rare,
admin-triggered build-required config-override window — so a control
collective here is fine (the same window already model-loads + captures
under NCCL). The function is injectable so the CPU two-fake-rank tests
drive a scripted verdict without a real process group.

Return whether every rank's member build succeeded.

Single-process / TP=1: the verdict is just ``local_ok`` (no peers).

TP>1: ``all_reduce(SUM)`` a per-rank ``1`` (ok) / ``0`` (failed) flag
over the TP group; the cluster commits iff the sum equals the world
size. Every rank must call this exactly once per build attempt — the
success path and the failure path both reach it — or the collective
desyncs. The reduce tensor lives on the engine's device so it matches
the group backend (NCCL needs CUDA; gloo needs CPU).

Pre-prepare VRAM feasibility gate for a build-required config-override.

A capture-affecting config override builds a new captured member beside the
currently-active one. When the variant differs only in capture/runtime knobs
the new member shares the active member's immutable weight storage — dense or
weight-quantized (EXL3 / AWQ / FP8 / NVFP4), see
:func:`build_helpers.find_weight_share_donor` — so it pays no second weight
copy and this gate lets it through. A variant that also ADDS structure (the
bundled MTP head) shares the body and pays only for the delta, which the gate
sizes from the checkpoint headers. Only a genuinely different model
(``model.path``) needs a full second residency, which on a memory-tight large
midway — after it has already started the SHM-coupled park/build sequence on
every rank.

When that OOM strikes mid-build, one rank dies inside model-load/profile while
the other is still running, and the cross-rank commit/abort barrier (which only
fires after the build returns) is never reached. The ranks desync and the
rank-0/rank-1 SHM control queue blocks forever (``shm writer waited 60s for a
free block``) — the server hangs and never recovers.

This module closes that gap with a check that runs deterministically on every
rank from the same broadcast delta, before the park/build sequence is entered.
If the new member's estimated residency does not fit the local free VRAM (with
headroom), it raises a clean :class:`RuntimeError` on every rank in lockstep —
so the server returns a 4xx/5xx to the admin caller and stays alive, instead of
deadlocking. The estimate + free-VRAM read are the same computation on every
rank (TP shards are symmetric), so every rank takes the same branch.

The gate is intentionally conservative-but-cheap: it inspects only the active
member's weight-pool residency + a checkpoint-quant probe (tensor-name only,
materializes nothing). It never reads weights or runs a profile. A driver
hiccup that makes the free-VRAM probe fail is treated as "cannot prove
infeasible" → the build proceeds (the in-build gates + barrier remain the
backstop); only a successful probe reading genuinely-insufficient free VRAM
refuses.

Best-effort live bytes the active member's weight pool holds.

A new member that cannot share weights pays approximately this much again
(the donor weights stay VRAM-mapped through the park — the park frees pool
physical, not weights), so it is the dominant term of the second residency.
Returns 0 when the pool / snapshot is unavailable (CPU bookkeeping, no
weights pool) so the gate degenerates to "cannot prove infeasible".

Local free VRAM (bytes), or ``None`` when the probe is unavailable.

Symmetric across TP ranks (each rank reads its own device). ``None`` on a
driver hiccup / CPU so the gate never refuses on a failed probe.

Bytes the park + host-evict of the active member returns before the new
member allocates.

Every named pool EXCEPT ``model.weights`` (a park keeps weights mapped by
the donor-share invariant). Best-effort: an unreadable pool contributes 0,
so the estimate is never optimistic about memory it could not see.

Estimate the weight bytes a SHARING member pays because the donor does
not cover its whole graph.

A donor is matched on ``(model.path, dtype)`` alone, so a member can share a
27B body and still add structure of its own. The one structural delta that
ADDS parameters today is the checkpoint's bundled MTP head: the new config
asks for it (:func:`bundled_mtp_head_wanted`) and the active member never
built one, so every ``mtp.*`` tensor is loaded from the checkpoint into this
member's own weight pool.

Sized from the safetensors HEADERS (``prefix_byte_size`` — names and
``data_offsets`` only, nothing materialized), divided by ``tp_size`` because
the head's projections are TP-sharded like every other linear. Returns 0
when the delta does not apply or the probe fails — the caller then behaves
exactly as it did for a full share.

Whether a member built from ``new_cfg`` could share the active member's
weight storage (so it pays no second weight copy).

Mirrors :func:`build_helpers.find_weight_share_donor`'s gate without
mutating the engine: a same-``(model.path, dtype)`` checkpoint — dense or
weight-quantized (EXL3 / AWQ / FP8 / NVFP4) — shares the donor's mapped
weights (the donor-share binder aliases quant weight buffers exactly like
the warm flat-dump reload does; see
:func:`build_helpers.find_weight_share_donor`). Only a different
``model.path`` (a genuinely different model) pays a full second weight copy.
Conservative on any probe error: assume it does not share (so the gate is
stricter, never falsely-optimistic).

Refuse a build-required override that cannot fit a second residency.

Runs on every rank deterministically (the same broadcast delta routes
every rank here before the SHM-coupled park/build), reads the local free
VRAM + estimates the new member's incremental residency, and raises a clean
:class:`RuntimeError` in lockstep when it does not fit. Raising before the
park/build is entered keeps the SHM control queue uncoupled, so a refusal
returns an error to the admin caller and leaves the server alive — instead
of one rank OOMing mid-build and deadlocking the rank-to-rank queue.

No-op (returns) when the free-VRAM probe is unavailable (CPU / driver
hiccup) — the in-build sizing gates + the cross-rank commit/abort barrier
remain the backstop; this gate only adds an early, symmetric refusal for
the provably-infeasible case.

Retire the process-global device-scratch caches at a member-build seam.

A handful of hot paths memoise a device scratch tensor in a MODULE-LEVEL dict
keyed by ``(device, dtype, shape)`` — never by which stable-VA member allocated
it. The tensor is created with a bare ``torch.empty(..., device=...)``, so it
lands in whatever :class:`~arbi_serve.runtime.named_pool.NamedMemPool` routing
is ambient at first use. During a boot capture sweep that routing is the
member's ``capture.cudagraphs`` pool
(``runtime/capture/decode.py`` wraps the warmup forward in ``graph_pool.use()``);
during serving it is the member's activation arena.

A live config-variant rebuild
(:func:`arbi_serve.engine.config_variant._aprepare_member`) parks the active
member and calls ``evict_parked_to_host``, which sleeps every named pool except
``model.weights`` — the pool physical is unmapped, the stable VAs stay reserved.
The incoming member is then built into the same primary engine and runs its own
warmup + capture forwards. Those forwards hit the memo, get the OUTGOING
member's tensor back, and read or write the now-unmapped VA:
``cudaErrorIllegalAddress``. The fault is asynchronous, so it surfaces at the
next synchronising call — typically the first capture bucket — and the sticky
context error then fails EVERY remaining bucket, which
``_aprepare_member`` correctly reads as a poisoned context and latches as a
fatal fault. A cold boot never hits it: the memos start empty.

The sampler's scratch memos are carried the same way. Their tensors are
allocated through a module-level ``_POOL_ALLOCATE`` closure over the BUILDING
member's ``scratch.sampler_triton`` / ``scratch.penalty_accum`` pool, so a memo
hit and a fresh allocation are both member-scoped. The closure itself is not a
tensor, so it is not carried here — the engine holds it as a per-model
attribute and
:func:`arbi_serve.engine.stable_va_controller._rearm_sampler_scratch_pools_on_wake`
re-arms it at the woken member's pools.

Per-member, not per-process
---------------------------
The memos are MEMBER state, so they follow the member:

  * :meth:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController.
    _park_mechanics` calls :func:`snapshot_and_clear_device_memos` and stores
    the result on the parked record, leaving the live memos EMPTY;
  * :meth:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController.
    _wake_mechanics` calls :func:`restore_device_memos` so the woken member
    gets ITS OWN entries back — at the same stable VAs its captured graphs
    baked, now remapped;
  * ``Engine._reset_workspace_pool_state`` calls
    :func:`retire_member_scoped_device_caches` as the build-seam net, for
    anything a park did not claim (an orphaned rollback member, a
    never-parked boot member).

The entries are never FREED. Dropping the last reference would return the
block to that member's (slept) pool free list, and its captured verify/decode
graphs carry the address in their kernel args — a later allocation reusing the
block would corrupt the replay after a switch back. The park/wake pair keeps
the reference on the record; the build-seam net moves it to :data:`_RETIRED`,
a process-lifetime pin (the same contract ``_CUDAGRAPH_POOL_PINS`` in
:mod:`arbi_serve.runtime.named_pool_keepalive` gives a capture pool).

Only modules ALREADY in ``sys.modules`` are touched: a memo that was never
imported cannot hold a tensor, and importing (say) the tkv prefill dispatch
just to clear an empty dict would be a gratuitous build-time import.

A memo the roster cannot NAME is a memo it cannot carry, and a closure cell has
no name: the roster is walked with ``getattr``, so a cache kept in a cell is
invisible both to the carry and to the test that catches a rename. A module
that memoises device scratch therefore holds it in a module-level mapping —
:mod:`arbi_serve._fla_persistent_cache` is the worked example, and its two
``prepare_chunk_indices`` caches are the GDN batch metadata's chunk-index
tables.

One module-level memo that may hold member-scoped device memory.

``module`` / ``attr`` name a ``dict`` whose VALUES are (or contain) device
tensors. ``attr`` may be a DOTTED path (``g_tensor_cache.cache``) for a
memo a module holds on an object rather than directly — the mapping is
still the thing carried; only the way to reach it differs. ``why`` records
what the memo caches and which forward reads it, so a reader can tell at a
glance whether an entry still applies. ``optional`` marks a memo whose
module may legitimately be absent (the vendored tkv package on a CPU-only
install, exllamav3 on a build with no EXL3 weights).

Yield ``(spec, mapping)`` for every registered memo whose module is loaded.

Only ``sys.modules`` is consulted: a memo whose module was never imported
cannot hold a device tensor, and importing it here would add a build-time
import for nothing. A loaded module that lacks the named attribute is
logged and skipped — a rename would otherwise silently reintroduce the
fault this module prevents.

Take the parking member's memo entries off the module globals.

Returns ``{"<module>.<attr>": {key: value, ...}}`` — the caller (the park)
stores it on the member's residency record and hands it back to
:func:`restore_device_memos` on that member's wake. The live memos are left
EMPTY, so whichever member is built or woken next allocates its own scratch
in its own mapped pool instead of writing the parked member's unmapped VA.

Put a woken member's memo entries back on the module globals.

Anything currently in a memo belongs to a DIFFERENT member (the one this
wake is switching away from, or an orphaned rollback build) — it is pinned
into :data:`_RETIRED` and evicted, never merged: two members' entries share
a key space with no member discriminator, so merging is exactly the
cross-member aliasing this module exists to prevent. Returns the number of
entries restored.

Drop the scratch-allocator closures that capture a member's pool OBJECT.

``sampler/topk_topp_triton`` and ``sampler/penalty_accumulator`` each hold a
module-level ``_POOL_ALLOCATE`` closure installed by the member's build over
THAT member's named pool, and ``weight_quant/exl3/kernel_scratch`` holds the
member's ``capture.io_buffers`` pool the EXL3 descriptors' input rows are
allocated in. None of them is a tensor, so retiring buys nothing — but
leaving one armed means an allocation made between
``reset_model_state_for_build`` and the incoming member's re-install lands in
the parked member's evicted pool. Uninstall so the fallback (default
allocator) is used until the new build re-installs its own.

Returns the number of hooks uninstalled. Best-effort: a missing module is
never a reason to fail a build.

Empty every :data:`MEMBER_SCOPED_DEVICE_CACHES` memo. Returns the count.

Call from the member-build reset seam — BEFORE the incoming member's build
runs any forward — so no warmup or capture forward can dereference the
outgoing member's unmapped VA. In the normal (stable-VA) case the park has
already claimed the entries onto its record and this finds nothing; what it
catches is the member a park did not claim.

``pin=True`` (the default, the stable-VA seam) keeps the values alive in
:data:`_RETIRED`: the outgoing member's pool physical is only UNMAPPED, its
captured graphs still carry the addresses, and a freed block could be
handed out again. ``pin=False`` is for the hot-swap RELOAD teardown, where
the pools are genuinely freed and the captured graphs destroyed — there
holding a reference would keep dead storage off ``free_all``'s books.

Never raises: a build must not fail because a memo moved. A memo whose
module is loaded but whose attribute is gone is logged at WARNING (the
rename would silently reintroduce the fault the module prevents) and the
regression test fails on it.

Deterministic GPU memory budget — every named pool has a predicted byte cost.

Picking three independent knobs by hand and hoping the tuple fit the
card is fragile:

  - ``--num-pages`` (KV pool depth)
  - ``--max-batched-tokens`` (workload prefill / chunk-prefill width)
  - ``--activation-headroom-gb`` (a guess at activation peak)

Drift between picked values and reality manifests as mid-request OOM,
and named pools beyond weights + activation (``recurrent_pool``,
``mtp_snapshot_pool``, ``tkv_scratch_pool``, ``activation_arena``,
``rope_cache_pool``, ``lora_pool``, NCCL workspace, captured
``graph_pool``) are easy to leave unbudgeted.

This package computes that budget deterministically. The contract:

  1. Every named pool exposes a module-level ``predict_<pool>_bytes(*,
     ...)`` function. The function takes ONLY config / shape inputs (no
     CUDA, no engine state) and returns an ``int``. Each predictor's
     docstring shows the formula.
  2. :func:`compute_kv_budget` sums every prediction into
     ``non_kv_total_bytes``. ``kv_bytes = usable_bytes -
     non_kv_total_bytes - extra_safety``.
  3. When ``Σ predictions > usable``, :class:`BudgetExceeded` raises
     with the per-pool MiB breakdown so the OOM explains itself.
  4. The graph-pool predictor is two-pass measure-and-cache:
     :func:`predict_graph_pool_bytes` reads
     ``<budget-cache-dir>/<sha>.json`` (cache key =
     SHA256(model || tp || backend || shapes || gpu_arch ||
     torch_version || turbo_attn_version || profile_cache_version));
     :func:`measure_and_persist_graph_pool` writes it after the real
     capture sweep. Cold boot persists once; every later boot at the
     same configuration uses the cached value. Set
     ``ARBI_SERVE_BUDGET_CACHE_DIR`` to override the cache location.
  5. ``compute_kv_budget`` takes ONE ``activation_reserve_bytes``
     argument (always present, always profiled-peak * safety_factor).
     There is a single activation budget line — whether the runtime
     arena buffer is allocated or not, the number is the same.

This module is kept separate from :mod:`arbi_serve.engine.build` so it can be
unit-tested with mocked CUDA introspection — every predictor takes
scalars / dataclasses and returns scalars; nothing here touches the
engine.

Every public name of :mod:`arbi_serve.engine.memory_budget` is
re-exported here:

  - :mod:`._common` — exceptions, :class:`KvBudget`, GPU-util guards,
    the GDN MTP-rollback resolver, byte-size helpers.
  - :mod:`.pool_predictors` — the ``predict_*_pool_bytes`` predictors and
    the DFlash / tkv-bypass / TKV-prefill transient peaks.
  - :mod:`.reserves` — RoPE / LoRA predictors + the serving-reserve
    estimators (NCCL workspace, cuMem-invisible, torch-overhang,
    activation reserve).
  - :mod:`.graph_pool` — the two-pass captured-graph-pool predictor.
  - :mod:`.kv_budget` — :func:`compute_kv_budget` and the KV sizing /
    grow helpers.

Shared foundation for the :mod:`arbi_serve.engine.memory_budget` package.

Exceptions, the :class:`KvBudget` result dataclass, the GPU-utilization
range guards, the GDN MTP-rollback mode resolver, and the byte-size
helpers used by every predictor. Imports none of the sibling submodules
so it can seed them without a cycle.

Resolve the GDN MTP partial-accept rollback to its internal mode
(``"replay"`` | ``"recompute"``) from ``ARBI_GDN_MTP_INGRAPH_ROLLBACK``.

CAPTURE-AWARE. The SAME resolution runs in the boot memory budget
(``predict_mtp_snapshot_pool_bytes``) and at pool attach
(``RecurrentStatePool.attach_mtp_snapshot_buffers``) so the reserved
and allocated modes always agree.

* in-graph ON (default) → ``replay`` (the in-graph masked replay) on the
  normal cudagraph-CAPTURED decode path. Under degraded-eager
  (``ARBI_ALLOW_DEGRADED_CAPTURE=1`` / capture-impossible) the replay side
  graph cannot be built, and this RAISES :class:`MemoryBudgetError` rather
  than resolving ``recompute``: substituting it would change both the
  rollback mechanism and the pool reserved for it under a caller that
  budgeted for replay. There is no opt-in to the swap.
* in-graph OFF → ``recompute`` always (host-side replay; debug / eager).
* ``ARBI_ACCEPT_INVARIANT`` → ``recompute``, to pair with the invariant
  verify forward's per-token chaining.

The degraded-eager signal is ``allow_degraded_capture`` — the flag that
permits an eager decode-bucket downgrade — read here off ``runtime_flags``,
the same source this resolver consults for the knob.

Whether ``cfg.vram_mode`` asks a memory gate to RAISE rather than warn.

``bench`` refuses to report numbers from an engine whose memory posture is
not what a benchmark claims (degraded capture coverage, an inert cap, a
budget row standing on a seed); ``production`` serves and says so. One
predicate for every gate, so no two of them read the mode differently.

Raised when the profiled VRAM budget can't fit the requested pool.

Subclass of :class:`RuntimeError` so the existing call sites that
catch ``RuntimeError`` continue to work; we tag a distinct type so
tests can assert on the precise failure mode.

Σ of per-pool predictions exceeds the usable VRAM budget.

Carries the per-pool breakdown dict so the operator (and crash log)
sees exactly which pools are oversubscribed: every OOM dumps the
full per-pool breakdown WITH which budget line was breached.

The numbers :func:`compute_kv_budget` returns.

All sizes in bytes (binary). ``num_pages`` is the floor result that
the engine hands to :class:`MultiStatePool`.

``per_pool_predicted_bytes`` is the breakdown the budget computed
from. ``non_kv_predicted_bytes`` is its sum.

``activation_reserve_bytes`` is the single activation reserve line
item — always ``profiled_peak * safety_factor`` (or the closed-form
upper bound at cold-boot fail-fast). There is one activation
budget line, not separate peak / arena lines.

Return bytes-per-element for ``dtype``.

Accepts ``torch.dtype`` (when torch is imported by the caller) OR
a string name (``"bf16"`` / ``"fp16"`` / ``"fp32"`` / ``"int32"`` /
``"int64"`` / ``"uint8"``). Predictors take ``dtype`` as ``Any`` so
they stay importable without a torch dependency in CPU-only test
contexts.

SM count of ``device`` (default: current CUDA device), or 0.

The split-K reserves need it to price an occupancy-targeted kernel pick,
which is a per-DEVICE quantity. Returns 0 rather than raising when CUDA is
absent so the callers stay usable on a CPU-only budget pass; every consumer
treats 0 as "no device geometry" and falls back to a batch-shaped bound.

Phase-bracketed measurement of the DRIVER-resident VRAM no pool can see.

Every byte of VRAM on the card is in exactly one of four places:

  1. a cuMem-mapped named pool (``CuMemPoolAllocator.mapped_bytes(tag)``),
  2. the growable KV region (cuMem-driven directly, not a torch ``MemPool``),
  3. the torch caching allocator OUTSIDE those pools (the default pool), or
  4. **nowhere any allocator we own can see it** — the CUDA primary context,
     the cubin modules the driver loads on first kernel launch, the driver's
     ``cudaGraphInstantiate`` exec allocations, NCCL's internal workspace, and
     any raw ``cudaMalloc`` a C++ extension makes.

(4) is the only term with no counter behind it. ``torch.cuda.memory_stats``
cannot see it; ``mapped_bytes`` cannot see it; only ``cudaMemGetInfo`` can, and
absorbed whatever the reconcile could not explain, which makes the residual
line balance by construction and tells an operator nothing.

This module replaces the constant with a **bracket**. ``mark(phase)`` samples
the four quantities at a boot-phase boundary and stores them; the DIFFERENCE in
term (4) between two consecutive marks is driver-resident physical that appeared
*during that phase*, measured, attributable to the work that phase does. Because
the marks telescope, the phase deltas plus the first sample sum EXACTLY to the
driver residency at the last mark — there is no slack for a fudge constant to
hide in, and no phase's growth can be silently reassigned to another.

What that buys, concretely:

  * ``driver.cuda_context`` is the FIRST sample (taken right after CUDA init,
    before a single weight is read) — the irreducible per-process context.
  * ``driver.modules_loaded`` is the sum of the later phase deltas MINUS the
    independently metered ``cudaGraphInstantiate`` bytes, and the ledger prints
    the per-phase breakdown, so "which library grew it" is answered by *which
    phase grew it* — compile/autotune, capture sweep, serve-kernel warmup.
  * ``driver.cudagraph_exec`` comes from
    :data:`arbi_serve.runtime.capture.graph_exec.instantiate_meter`, which
    brackets the real ``cudaGraphInstantiate`` calls.
  * ``driver.residual`` is what is left. On an honest boot it is ~0 because the
    phases are exhaustive; a non-zero value means driver-resident physical
    appeared BETWEEN the brackets. It is a RESIDUAL, not a leak signal: a
    configuration that legitimately loads more cubins than its held

    Not to be confused with ``unpooled.unregistered_pool``, a DIFFERENT quantity on a
    different surface: the torch caching-allocator row for a private
    ``MemPool`` that is neither registered nor released (see
    :func:`arbi_serve.engine.memory_budget.pool_residency.residency_by_pool`).
    That one IS a leak signal — it names an owner nobody declared. This one
    absorbs the difference between two measurements of the whole card.

Boot-only. Each ``mark`` is two CUDA queries plus a few Python ints; the serving
hot path never touches this module.

Ordered boot-phase brackets around the driver-resident residual.

One instance per engine, created at the top of the build and marked at every
phase boundary. Holds only the samples; the attribution arithmetic is the
free function :func:`attribute_driver_residency` so it is CPU-testable.

``growable_kv_reader`` is a zero-arg callable returning the growable KV
region's mapped bytes (``pool.growable_mapped_bytes_total``). It is a
CALLABLE rather than an engine reference so this module stays import-cycle
free and trivially stubbable; it may be installed after construction (the
meter's first marks happen before the KV pool exists, when it reads 0 —
which is correct, not missing).

Read the four accounting quantities off ``device`` right now.

Returns ``None`` when CUDA is unavailable or a probe raises.

``settled`` is for a probe taken between boot phases, where the device is
quiescent and the reading must be exact: it synchronizes first so
``cudaMemGetInfo`` sees a settled device, and it queries the cuMem
allocator for :attr:`ResidencySample.cumem_mapped_bytes`.

Pass ``settled=False`` for a probe taken WHILE THE ENGINE SERVES. Both of
those steps are unsafe there and neither is needed:

  * ``torch.cuda.synchronize`` waits for the whole device, and a saturated
    engine keeps the queue non-empty, so it has no bound;
  * ``CuMemPoolAllocator.mapped_bytes`` takes the allocator's lock, which
    an in-flight allocation holds with the GIL RELEASED while its Python
    hook waits to re-acquire the GIL — a query issued from another thread
    deadlocks the engine (the hazard
    :func:`arbi_serve.runtime.named_pool.pool_allocation_in_progress`
    names). ``cumem_mapped_bytes`` reads 0, which no term of the driver
    attribution consumes.

The ``state.attn_kv.mapped`` row AND whether its bytes are known.

The mapped KV slab is the largest single row of a serving card, and it is
the one row no allocator snapshot can see: the growable region drives cuMem
directly, so it is outside ``torch.cuda.memory_reserved`` and outside every
named-pool tag. A consumer that renders the card as a closed breakdown
therefore has to be told the difference between "these are the KV bytes",
"the KV bytes are on another row", and "the KV bytes are UNKNOWN" — booking
the third case as zero is what lets a panel close its identity by folding
the whole slab into its residual.

:attr:`state` is one of:

``mapped``
    The growable region was read; :attr:`bytes` is its mapped physical.
``pooled``
    The KV slab is not growable, so its physical is inside the
    ``state.attn_kv`` named pool and is reported on THAT row.
    :attr:`bytes` is 0 here by definition, not by failure.
``absent``
    This engine holds no KV pool (before the KV phase of boot, or while
    released). :attr:`bytes` is 0 and the card has no KV row.
``unresolved``
    The read failed. :attr:`bytes` is meaningless and no consumer may
    close the card from this snapshot.

Read the mapped-KV row off ``eng``, saying WHICH of the four states holds.

The single in-process source of ``state.attn_kv.mapped``: the boot VRAM
ledger, the driver-residency brackets and the admin memory snapshot all
take the row from here, so no surface can carry a second accounting of the
same bytes.

Physical the growable KV region has mapped right now (0 when absent).

The region drives cuMem directly, so these pages are resident physical
OUTSIDE the torch caching allocator. Every bound that separates the two
must be taken from the same reading — see ``real_reserved_bytes``. Thin
over :func:`resolve_mapped_kv_row`, which is the one source of this number;
a caller that must distinguish "0 bytes" from "unknown" calls that instead.

The ``state.gdn_recurrent.mapped`` row AND whether its bytes are known.

Sibling of :class:`MappedKvRow`, for the same reason: the state pools'
sentinel-alias arenas drive cuMem in their OWN VA reservation, so the bytes
are outside ``torch.cuda.memory_reserved``, outside every named-pool tag and
outside every allocator snapshot. A consumer closing the card has to be able
to tell "this engine has no recurrent state" (0 by definition) from "the read
failed" (0 by accident), because the second one silently moves a hybrid
model's whole recurrent slab onto whatever row absorbs the remainder.

:attr:`state` is ``mapped`` (bytes are the arenas' physical), ``absent``
(this engine holds no arena-backed state pool) or ``unresolved``.

Physical the state pools' sentinel-alias arenas hold (0 when absent).

Same class of bytes as :func:`growable_kv_mapped_bytes`: cuMem-mapped in
the pool's own VA reservation, outside the torch caching allocator and
outside the pluggable allocator's per-tag counters, so every bound that
separates torch's residency from the driver's must subtract it. Thin over
:func:`resolve_state_arena_row`; a caller that must distinguish "0 bytes"
from "unknown" calls that instead.

Bytes THIS process holds on ``device`` per NVML — 0 when unreadable.

``cudaMemGetInfo`` is device-wide; NVML's per-process walk is the only
reading that separates us from the card's other tenants, and without it the
ledger books their bytes against us. The walk does not always find us —
it reports pids in whatever namespace the driver picks — which is why
:func:`~arbi_serve.engine.boot_vram_drain._self_used_bytes` owns the
identification and this is a thin adapter onto it: one implementation of
"which of these processes is us", shared with the settle gate.

The device's pre-context mark, or ``-1`` when it was never taken.

Thin adapter onto
:func:`~arbi_serve.engine.boot_vram_drain.foreign_bytes_at_context_creation`
that speaks the ``-1``-for-absent convention
:attr:`ResidencySample.foreign_at_context_bytes` carries (the sample is a
plain frozen dataclass, so "not measured" has to be a value, not a type).

``((pid, bytes), …)`` for every compute process NVML sees on ``device``.

The per-pid detail behind :func:`nvml_process_used_bytes`. Diagnostic only:
it answers "whose bytes are these" when the device-wide reading and our own
share disagree, which is the one question a device-wide number cannot.
Empty when NVML is unavailable.

``(cudaLimitStackSize bytes, max-resident-threads)`` — ``(0, 0)`` on error.

Split out so the boot ledger can name the local-memory pool as a MEASURED
term with a knob attached, instead of leaving it inside an opaque "CUDA
context" lump. The driver raises the stack limit by itself when a launched
kernel needs a deeper frame, so this must be re-read at every phase, not
cached from boot.

The driver-resident residual, split into MEASURED named terms.

The four fields sum to exactly ``total_bytes``: :attr:`residual_bytes` is
defined as the remainder, so it can never be quietly folded into a
neighbour, and a non-zero value is a real finding rather than a rounding
artefact of a model.

:attr:`residual_bytes` is the DRIVER-side residual (``driver.residual``).
The allocator-side ``unpooled.unregistered_pool`` — an unregistered private
``MemPool`` in the torch segment walk — is a different quantity on a
different surface and is a leak signal; this one is not.

Split the measured driver residency into named terms. Pure arithmetic.

``measured_total_bytes`` overrides the meter's last sample (the freeze takes
its own ``mem_get_info`` read a moment later; pass it so the ledger balances
against the SAME numbers the freeze reports rather than a near-duplicate).

``cudagraph_exec_bytes`` is the independently metered
``cudaGraphInstantiate`` total. It is a SUBSET of the phase growth (the
capture phase's bracket already contains it), so it is carved OUT of
``modules_loaded`` rather than added on top — booking both would double the
same physical, the exact failure mode this ledger exists to prevent. The
local-memory pool is carved out the same way and for the same reason: it is
inside both the baseline context read and (when the driver raises the limit
mid-boot) the phase deltas.

``expected_modules_bytes`` is the configuration's established
``driver.modules_loaded``, from
:func:`~arbi_serve.engine.memory_budget.graph_pool.predict_driver_modules_bytes`.
When it is given, that value is BOOKED and this boot's bracketed growth is
held against it, so ``driver.residual`` comes out as the DEVIATION. When it is
``None`` the boot is seeding: the bracketed growth is booked, the residual
is 0, and the caller must report the boot as ungated.

That distinction is the whole point. Booking the bracketed growth always —
the previous behaviour — made ``modules_loaded`` absorb every byte the boot
phases grew, including an unregistered pool or a raw ``cudaMalloc``, and
pinned ``driver.residual`` at 0 by construction. CUDA offers no way to
enumerate module residency, so the term cannot be derived from the driver;
it can only be measured once and held.

After :meth:`DriverResidencyMeter.seal_boot` the samples split in two and
so does the arithmetic: the boot brackets give ``modules_loaded`` exactly as
before, and the marks taken while SERVING give ``modules_serving`` by the
same construction over the post-seal window. Without that split a live
re-measure has no term for driver residency that appears after the boot
brackets close, and the whole of it lands on ``driver.residual`` — which is how
a ledger that closes at boot decays under load. ``local_memory`` is re-read
from the LATEST sample either way, so a stack-limit raise while serving is
named directly rather than through a modules term.

With no samples at all every term is 0 except ``driver.residual``, which takes
the whole measured total: an unmeasured boot must read as unexplained, never
as explained-by-nothing.

Re-measure the ``driver.*`` terms against the device NOW.

The boot terms are bracketed and then HELD, which is correct for what they
name and wrong for what an operator reads them as: the local-memory pool
grows whenever the driver raises ``cudaLimitStackSize`` for a kernel first
launched under load, the exec meter keeps counting, and the driver keeps
loading cubins on first dispatch. Held, every one of those shows up as
``driver.residual`` — a ledger that closes at boot and decays exactly
when someone looks at it.

So the live view re-reads the measurable terms (local memory from the
device, cudagraph exec from the meter) and brackets the rest over the
serving window ``seal_boot`` opened, exactly as the boot phases are
bracketed. ``modules_loaded`` stays BOOKED at the boot value: it names the
boot's module residency and nothing about serving can change what boot did.

Returns ``None`` when the meter is absent, the boot is unsealed, or the
probe fails — the caller then reports the held boot rows, which is the
honest fallback.

Every input behind ``driver.cuda_context``, on one line.

The term is a residual of residuals — device-wide used, our own share,
torch's reservation, the kernel-stack pool — and each of those moves for a
different reason. Printing the derivation makes an inflated context term
self-diagnosing instead of a number an operator has to reproduce.

Log the ``driver.cuda_context`` derivation and flag what can corrupt it.

``driver.cuda_context`` is the bare CUDA primary context: a per-device,
per-driver constant that does not move with our configuration. Nothing
verified that. It is also the FIRST thing measured on a boot, before a byte
of ours is allocated — so whatever else is resident at that instant is the
only thing that can move it, and when the term is derived from a
device-wide reading that is exactly what it absorbs. The KV pool is the
residual of the card, so every byte booked here is a byte KV does not get,
for the life of the boot.

Three checks, each naming what it found:

  * the split we got — ``device_wide`` means the term is unverifiable and
    carries every other tenant;
  * the two independent splits AGREEING, when both are available — they
    measure the same quantity by different routes, so a disagreement means
    one of them is lying and the term is not what it claims;
  * foreign residency at the baseline — legitimate on a shared card, and
    the operator still has to be told what it costs in served context.

Diagnostic only: it never fails a boot. A card that really is shared must
still boot, and it must boot saying so.

Resident physical THIS process holds.

Three readings, in decreasing order of directness:

``nvml``
    NVML's per-pid entry for us. Direct, and the only one that stays
    right once other tenants come and go.
``pre_context``
    ``device_used − foreign_at_context``. NVML reports pids in the
    namespace the driver picks; where that is not ours, no entry ever
    matches and the walk attributes us NOTHING. The pre-context mark
    still splits the card, because we held none of it when it was taken.
``device_wide``
    The whole card. Last resort, and a LIE whenever the card is shared —
    it books another tenant's bytes as our driver residency, which is
    the failure :attr:`non_torch_bytes` documents.

See :attr:`own_attribution_source` for which one this sample got.

The least this process can possibly hold on the card.

Its own caching-allocator reservation, and the kernel-stack local-memory
pool the driver backs INSIDE its primary context. Both are measured on
this sample, and any reading of our own share below their maximum is not
a small number — it is a wrong one, because a process taking this sample
necessarily holds the context those bytes live in.

Our share per the pre-context mark, or None when the mark is unusable.

The mark is only foreign-by-construction if it was taken before this
process created its context. Nothing can prove when it was taken, but
:attr:`own_floor_bytes` can prove when it was NOT: a mark that implies
we hold less than our own floor contained our own bytes, so it splits
the card wrong and must not be used.

Which reading :attr:`own_used_bytes` came from.

``"nvml"`` / ``"pre_context"`` / ``"device_wide"`` — see there. A
consumer that sizes anything off this sample has to know: only the last
one can carry another process's bytes.

Resident physical no allocator we own can account for.

``own_used − torch_reserved − the cuMem regions we drive directly``
(the growable KV slab and the state pools' sentinel-alias arenas).
The cuMem named pools
are NOT subtracted separately: they are torch ``MemPool``s and are
already inside ``torch_reserved_bytes`` — subtracting them again is the
double-count that drove this bucket artificially toward 0 and hid real
module-load growth.

Built on :attr:`own_used_bytes`, not the device-wide reading: a
co-tenant's VRAM is not our driver residency, and letting it in put
another process's allocations into ``driver.cuda_context`` (whatever was
resident at the baseline mark) and its churn into
``driver.modules_loaded`` (every later phase delta).

Device local memory the primary context backs for kernel stacks.

``stack_limit × max-resident-threads``. MEASURED-derived, not modelled:
on a 4090 (128 SM × 1536 threads) the CUDA-default 1024 B limit implies
straight back to ``cudaMemGetInfo`` — i.e. this term is real, it is
roughly HALF the bare context floor, and it is REDUCIBLE (see
:func:`arbi_serve.engine.sleep.shrink_thread_stack`, which already
exploits it while asleep). 0 when the limit could not be read.

Sample the device NOW and file it under ``phase``.

Returns ``None`` (and records nothing) when CUDA is unavailable or any
probe raises — the meter is diagnostic, and a boot must never die
because an accounting read failed.

Close the boot brackets. Idempotent.

Splits the sample list in two: everything up to here is BOOT (its
growth is ``driver.modules_loaded``, held against the configuration's
baseline), everything after brackets SERVING. Without the split a live
re-measure has no way to tell a serving-time cubin load from a boot
one, and the whole post-freeze growth falls into the residual.

Re-sample the OPEN serving bracket (one sample, replaced in place).

The serving bracket is open for the engine's whole life and is read
every time an operator asks where the card went, so appending would
grow the sample list without bound and print a phase breakdown of
identical rows. There is exactly one serving sample: the latest.
No-op (returns ``None``) before :meth:`seal_boot`.

``[(phase, Δ non-torch bytes)]`` between consecutive marks.

The label is the phase the delta is CHARGED TO — i.e. the label of the
LATER sample, which closes that phase. Deltas may be negative (a phase
that freed driver memory); they are reported signed, because clamping
them would break the telescoping identity this whole design rests on:

    baseline.non_torch + Σ deltas == latest.non_torch

:meth:`phase_deltas` from the boot close to the latest sample.

Empty before :meth:`seal_boot` — an unsealed meter is all boot, and a
bracket that has not been opened must read as absent rather than as
zero growth.

How much of ``driver.modules_loaded`` the CAPTURE phases took.

``modules_loaded`` is the union of distinct CUDA cubins the process
ever dispatches, and on this engine most of it is not the model: the
capture sweep front-loads kernel code that a first request would
otherwise pay for. Without this split the row is a 400-MiB black box
and an operator weighing the captured ladder against KV cannot see that
the ladder is most of it.

MEASURED, not apportioned: the capture phases' own bracketed deltas,
minus the ``cudaGraphInstantiate`` bytes already metered separately as
``driver.cudagraph_exec`` (those live inside the same phase deltas, and
counting them here would book them twice).

Clamped at 0. A negative would mean the metered exec exceeded the
phases that contain it, which is an accounting fault, not a share --
and reporting a negative share as if it were one would hide it.

Taxonomy-named rows, zero-valued ones dropped except the residual.

``driver.residual`` is ALWAYS emitted, including at 0: an operator must
be able to read "the unexplained term is zero" off the ledger rather
than infer it from a missing line.

Two-pass measure-and-cache predictor for the captured ``graph_pool``.

:func:`predict_graph_pool_bytes` reads ``<budget-cache-dir>/<sha>.json``;
:func:`measure_and_persist_graph_pool` writes it after the real capture
sweep. Cold boot persists once; every later boot at the same
configuration uses the cached value.

Where the budget cache lives on this boot, and whether an operator said so.

``(directory, came_from_ARBI_SERVE_BUDGET_CACHE_DIR)``. The second half is
what tells a cold-cache report whether the directory is a deliberate,
persistent mount or the ``~/.cache`` fallback — which on a container root
filesystem is gone at every restart, so every boot reads an empty cache,
sizes memory from seeds, and fails identically. Reported by
:func:`~arbi_serve.engine.build_memory_sizing.cold_budget_cache_note`.

Pick a writable budget-cache directory.

Resolution order:
  1. Explicit ``build_dir`` argument (tests).
  2. :func:`resolved_budget_cache_dir` — the env var, else the fallback.

The second step is not re-derived here: the cold-cache report names the
directory this reader used, and a report naming a different path than the
one that was read is worse than no report.

Write ``entry`` to ``final_path`` atomically: tempfile → fsync → rename.

A crash mid-write leaves a ``.tmp`` behind but never a half-written cache
file a later boot would size memory from. Shared by all three budget-cache
entries (pool size, capture residual, serving overhang) so the durability
property cannot hold in two of them and quietly not in the third.

SHA256 of the key-sorted JSON dump of ``cache_key_inputs``, truncated.

The components (built by ``_graph_pool_cache_key_inputs`` in
:mod:`arbi_serve.engine.build`) are: model_id, tp_size, kv_backend,
shape_buckets, gpu_arch (compute capability), the kernel-library identity
(``torch_version`` — carries the CUDA build tag — and
``turbo_attn_version``), and ``profile_cache_version``. These are the
inputs that actually move the measured graph-pool size, so the cached
budget is reused iff it genuinely matches.

Hashed over a JSON-serialised, key-sorted dump so two boots at the same
config produce identical keys.

Predicted bytes for :data:`graph_pool` — two-pass measure-and-cache.

The cudagraph private allocator overlays per-layer kernel
intermediates within a bucket; there is no closed form.
The canonical source is the measured
``eng.total_cudagraph_bytes()`` from a cold-boot capture
sweep, persisted to ``<budget-cache-dir>/<sha>.json``. The cache
key (see :func:`_graph_pool_cache_key`) covers model/tp/backend/shapes/
gpu_arch + kernel-library versions + profile_cache_version; the same
configuration on the same hardware reads the same cached bytes.

Returns:
    ``int`` (cached graph-pool bytes) when the cache file exists
    and parses cleanly. ``None`` when no cache hit — the caller
    must run :func:`measure_and_persist_graph_pool` against the
    actual capture sweep, then re-call this function (or use the
    return value of the measure pass directly).

Args:
    cache_key_inputs: dict the cache key is derived from. Required
        keys: ``model_id``, ``tp_size``, ``kv_backend``,
        ``shape_buckets``, ``gpu_arch``, ``torch_version``,
        ``turbo_attn_version``, ``profile_cache_version``. Order
        does not matter (the dict is JSON-dumped key-sorted).
    build_dir: override cache directory. ``None`` (default)
        reads ``ARBI_SERVE_BUDGET_CACHE_DIR`` env, then falls
        back to ``~/.cache/arbi-serve/budget-cache``.

MEASURED serving-resident torch-caching overhang, or ``None`` when cold.

The prefill-width activation scratch PyTorch's caching allocator keeps
RESIDENT after the capture grow (per-MemPool segments are not returned by
``empty_cache``). The KV serving ceiling must hold it out, and the ceiling is
computed at PROFILE time — before the capture sweep that grows most of it —
so it was PREDICTED as ``profiled_activation_peak x 1.0``.

on Qwen3.5-0.8B TP1 bf16 mb64, and its own docstring demanded the two agree
within ~5%. The miss is most of the PREDICTED-vs-REALIZED page gap, and it
is exactly what the cumulative learned residual was absorbing — a learned
residual standing in for a modelling error is opaque provisioning, not a
prediction.

The term is not unknowable, only not knowable YET at the moment the ceiling
is computed: the Phase-2 freeze MEASURES it directly
(``NamedPoolRegistry.torch_caching_reserved_overhang_bytes``). So measure it
once per configuration and read it back — the same contract
``graph_pool_bytes`` and the instantiate rate use, and for the same reason.

Persist the freeze-MEASURED serving torch-caching overhang, MONOTONICALLY.

Written from the Phase-2 freeze, where the quantity is directly observable.
Read by ``profile_and_size_kv_pool`` on every later boot at that
configuration, replacing the ``peak x factor`` model.

MONOTONIC — keep ``max(new, cached)`` — because this is a BOUND, not a point
estimate, and the difference is not pedantry. Two consecutive boots of the
IDENTICAL configuration (Qwen3.5-0.8B TP1 bf16 mb64, RTX 4090)::

    boot            growable KV      overhang      sum
    cold cache        14.660          0.862       15.522
    warm cache        14.062          1.435       15.497

every byte the budget does not hand to KV, the caching allocator keeps as
two boots that ran the same config — because a cold boot measures the graph
pool in a CHILD process while a warm one does not, so the parent's allocator
history differs. It is a property of the boot PATH, not of the model.

A last-sample point estimate therefore OSCILLATES: predict the low sample,
hold out less, and the allocator takes more; predict the high one and the
reverse. Taking the max converges on the worst case and stops. And it errs
in the safe direction by construction — an over-large overhang costs
FORECAST pages (it can trip a spurious ``max_context`` narrow) but never
realized KV, which ``grow_kv_after_capture`` sizes against measured free.
Under-predicting is the dangerous direction: that is the forecast promising
KV the allocator is already holding.

Same monotonic discipline, same reason, as ``graph_pool_bytes``.

One stamped provenance line, or ``""`` when the file is not there.

The build writes these; a dev bind-mount has no ``/etc`` stamp and reads
``""``, which keeps the key stable for a developer iterating on a checkout
rather than re-seeding on every edit. That is the deliberate trade: the
guard gates images, where the stamps exist and staleness actually ships.

``"<n>:<geometry>,…"`` — the EXL3 geometries the shape pin resolved.

The GEOMETRY SET, not the family each one won. The set is a property of the
model and of which linears are bound, so it reproduces across boots at one
configuration; the winning family is chosen by timing and can flip on a
marginal geometry, and a key that flapped would re-seed every boot and
leave the guard permanently inert — the failure
:func:`persist_driver_modules_baseline` names. The winner is safe to omit
because selection TIMES every legal family, so all of their cubins are
resident either way and the choice moves no module bytes (measured: pinning
401 linears over 8 geometries versus pinning none moved bracketed module

What the set does carry is every change that moves the cubin COUNT: the pin
going from inert to active (0 geometries to 8) and a drafter contributing
its own linears (8 to 16 on a DFlash boot against a native-MTP one).

Read at freeze, after both pin seams have run, so predict and persist —
which both run there — see one value.

What decides WHICH CUBINS LOAD, beyond the graph-pool key's inputs.

The graph-pool key carries ``torch_version`` and ``turbo_attn_version`` and
stops there, which covers a toolkit or a kernel-library bump and nothing we
do ourselves. These three close that gap:

  * ``arbi_serve_rev`` — our own build. Loading is decided by our code:
    making the EXL3 shape pin actually resolve in a built image changed the
    kernel set with no dependency bump at all.
  * ``exllamav3_rev`` — the fork holds the EXL3 GEMM kernel map, and two
    different pins both report version ``1.4.2``, so the dist version is
    blind by construction (see the Dockerfile's provenance stamp).
  * ``exl3_shape_pin`` — the resolved geometry set, which moves with the
    model and drafter even when neither revision does.

Deliberately NOT the JIT-compiled kernel count: it is the largest single
driver of module residency (201 of 237 kernels on a DFlash boot are Triton)
but it varies with compile-cache state, and a key that moves on a warm-vs-
cold cache re-seeds forever and disables the guard.

The ``<key>.modules.json`` identity: the graph-pool key WIDENED.

A separate key rather than widening :func:`_graph_pool_cache_key` itself.
That key is deliberately blind to our git SHA so a cached graph-pool budget
survives unrelated code edits, and putting a build revision into it would
force a full capture re-measure on every commit. The two caches want
different identities: the graph pool is sized by SHAPES, the modules
baseline by WHICH KERNELS LOAD.

The established ``driver.modules_loaded`` for this configuration.

``None`` when this configuration has never been measured — the boot then
SEEDS the baseline and says so, because a fabricated expectation is worse
than an admitted gap.

CUDA exposes no way to ask what the cubin modules it loaded cost. The whole
``cuModule*`` surface returns functions and globals, never bytes, and
``cuModuleGetLoadingMode`` reports only lazy-vs-eager; measured on sm_89 /
cu13, forcing cuBLAS/cuDNN/cuFFT/sort/SVD modules to load moved
they are OUTSIDE every allocator we own, and nothing enumerates them.

So the term is measured the only way it can be: bracketed once at a
configuration, then held to. What that buys is the point — while it was
computed as the bracketed residual it EQUALLED whatever the phases grew, so
an unregistered pool, a creeping NCCL buffer or a raw ``cudaMalloc`` was
absorbed into it by construction and ``driver.residual`` stayed 0. Held
against a baseline, that growth has nowhere to go but the residual line.

The cache identity is :func:`_driver_modules_cache_key`: the graph-pool key
(model, TP degree, KV backend, shape buckets, GPU arch, ``torch_version``,
``turbo_attn_version``) WIDENED by :func:`_driver_modules_identity`, which
adds the build revision, the exllamav3 fork revision and the resolved EXL3
pin geometry set.

The graph-pool key alone was not enough, and the gap was not theoretical.
It carries dependency VERSIONS and no identity for our own code, so every
change we made to which kernels load — the EXL3 shape pin going from inert
to resolving 409 linears, a drafter adding its own geometries — reused a
baseline measured before the change. Held against it, the extra residency
had nowhere to go but ``driver.residual``, which fires the freeze alarm.
``turbo_attn_version`` is a plain dist version (``0.49.0``) and does not
move when that library is rebuilt with a different tile ladder either.

Establish this configuration's ``driver.modules_loaded``. SEED-ONCE.

Returns whether a baseline was written.

Deliberately NOT monotonic and deliberately not refreshed: the caching
overhang next to it takes ``max(new, cached)`` because it is a BUDGET BOUND
and erring high is safe, but this value is a GATE. Re-writing it from every
boot lets a leak walk the expectation up a little at a time and the guard
reports green the whole way — the ratchet is the failure mode, not the
safeguard.

So a configuration whose module residency legitimately changes must reach a
new KEY rather than a refreshed value, and the key has to actually carry
the change. :func:`_driver_modules_cache_key` covers the build revision,
the exllamav3 revision and the EXL3 pin geometry set alongside the
graph-pool inputs. What it still does NOT cover is a kernel library rebuilt
at an unchanged version string with different contents: that re-seeds only
if the build revision moved with it, which for an image it does.

Driver-resident physical SERVING has been observed to add here, or ``None``.

This covers bytes that appear in no pool at all — the far side of the
allocator boundary.

WHAT THE QUANTITY IS. ``driver.modules_loaded`` is bracketed over the BOOT
phases and then held (:func:`predict_driver_modules_bytes`), which is what
it names. But CUDA loads a module image on FIRST DISPATCH, and the loading
mode is lazy unless an operator overrides it, so a kernel a served step
reaches and no boot phase launched materialises its cubin *after* the KV
layout is frozen. The same window carries a kernel-stack pool the driver
re-grows for a deeper frame and any graph instantiated post-boot. Every one
of those is driver-resident physical that no allocator of ours can see,
taken out of the free VRAM the serving floor holds.

THE REPORT FOR IT ALREADY EXISTED; THE RESERVE DID NOT. ``driver.
modules_serving`` is the first of those three terms, split out by
:func:`~arbi_serve.engine.memory_budget.driver_residency.
attribute_driver_residency` over the bracket ``seal_boot`` opens, and the
other two are named beside it. But nothing SAMPLED that bracket on a
schedule — the only caller re-measured on an operator's request — so the
rows read 0 for every process nobody opened the memory panel on, and no
term of ``serving_floor_for_grow`` reserved a byte of any of them. This
cache is that reserve's input, and it is the SUPERSET of the three: for a
floor the distinction between a cubin, a kernel-stack re-grow and a
post-boot graph exec does not matter, because all three are free VRAM the
grow gave to KV and the driver then took.

It is not derivable. CUDA exposes no way to ask what a module costs, or
which modules a workload will reach; the ``cuModule*`` surface returns
functions and globals, never bytes. So it is measured where it happens —
:class:`~arbi_serve.engine.memory_budget.driver_residency.
DriverResidencyMeter` opens a serving bracket at ``seal_boot`` and the
metrics export samples it — and read back here on the next boot.

``None`` when this configuration has never been observed serving. The floor
then holds NOTHING for it and says so on the row, which is the honest
answer: a stand-in here would be the fudge margin this term exists to
replace.

Keyed by :func:`_driver_modules_cache_key`, the same widened identity the
modules baseline uses, because it is the same question — WHICH CUBINS LOAD
— asked over a different window. The graph-pool key alone carries no
identity for our own build, and our own build is what decides which kernels
a step dispatches.

The per-configuration serving-observation record, or ``None``.

ONE file holds every reading a served process leaves for the next boot at
this configuration, and one reader parses it, so a second observation
cannot arrive with its own copy of the key derivation and its own idea of
what a stale version means.

Put a rising serving driver-growth reading on record. Returns the high-water.

MONOTONIC, unlike :func:`persist_driver_modules_baseline` beside it, and
the difference is what each number is FOR. That one is a GATE: refreshing
it would let a leak walk the expectation up a step at a time while the
guard reported green, so it is written once and a legitimate change must
reach a new key. This one is a BUDGET BOUND on a quantity that only ever
grows within a process — a cubin, once resident, stays resident — so the
largest reading is the right one and erring high costs KV rather than a
serving OOM. The two live in the same directory under the same key and
must not be confused: one file gates, the other reserves.

MEASURED capture-pool bytes PER cuMem tag, from the same cache entry.

``{}`` when the configuration has never been measured with a per-tag
counter available — a cold cache, or a boot on a card with no cuMem driver,
where ``total_cudagraph_bytes`` can only report an undifferentiated
estimate. The caller must then split the total itself and say how.

The keys are the pool names the budget uses as rows
(``capture.cudagraphs``, ``capture.io_buffers``), so a row can take its own
measured value instead of the sum standing in for both.

MEASURED ``cudaGraphInstantiate`` driver bytes per exec, by capture family.

``{}`` on a cold cache — "never measured at this configuration", which the
caller answers with an explicitly-labelled SEED, never with a fabricated
point estimate. The special key ``"*"`` carries the all-family mean and is
the fallback rate for a family this configuration has not yet exercised.

Same measure-once-and-persist contract, same cache identity and same file as
``graph_pool_bytes`` — because it is the same kind of quantity: a driver
cost with no trustworthy closed form, cheap to measure exactly once at the
configuration that will pay it. See
:class:`arbi_serve.runtime.capture.graph_exec.InstantiateMeter` for why the
refuted on this driver.

Persist measured graph_pool size to the budget cache.

Called AFTER the actual cudagraph capture sweep finishes. The
caller measures ``eng.total_cudagraph_bytes()`` and hands the
int to this function; we write the JSON entry with atomic rename
(``.tmp`` → final) so a partial / crashed write does not poison
the cache. Per-pool sub-component bytes (decode / drafter /
piecewise) are persisted alongside for diagnostic value but are
not consumed on cache hit.

Args:
    monotonic: when True, NEVER shrink an already-cached budget — keep
        ``max(new, existing)``. It guards ONE case: a value that is an
        ESTIMATE of the pool rather than a reading of it. On the non-cuMem
        path ``total_cudagraph_bytes()`` falls back to a per-graph
        ``memory_reserved`` delta which under-counts the true private-pool
        reserve, so persisting it would make the warm boot under-reserve KV
        and OOM.
    bytes_by_tag: the SAME reading split by cuMem tag
        (``capture.cudagraphs`` / ``capture.io_buffers``). The budget keeps
        a row per pool, so persisting only the sum forces the consumer to
        either apportion it by guess or feed one number into both rows. An
        empty / absent split is carried forward from the cached entry rather
        than clobbered, exactly like ``instantiate_bytes_per_graph``.
    authoritative: the caller read the pool from the cuMem per-tag counter
        (``capture.cudagraphs`` + ``capture.io_buffers``), which is EXACT
        for the whole sweep. An authoritative reading REPLACES the cached
        value at this key in BOTH directions, ``monotonic`` notwithstanding.

Why authority beats the ratchet. A ratchet over a value that is exactly
measurable turns a one-off high reading into a permanent one: the cache key
already forks on every input that moves the pool (ladder shape counts,
kv-page / LoRA buckets, ``prefill_capture``, ``split_attn``, the draft
geometry), so two readings under ONE key describe the SAME sweep and the
fresh one is strictly better information. Keeping the older, larger one
republishes a figure this boot did not measure and no counter on the machine
agrees with — and that figure is what the operator card and the boot ledger
print.

Truncation is NOT what monotonic protects. A capture that fell to eager is
refused a persist outright, before this function is reached
(:func:`~arbi_serve.engine.build_graph_pool.persist_graph_pool_budget`
returns early on
:func:`~arbi_serve.engine.boot_degradation.capture_pool_truncations`), so a
truncated reading never becomes a cache entry to ratchet against.

Returns the persisted ``graph_pool_bytes`` (caller-friendly: chain
``budget_bytes = measure_and_persist_graph_pool(...)``).

Measure-once-and-persist; never re-measure on every boot.

Persist the CUMULATIVE capture-sweep residual after a real grow.

``residual_bytes`` is the total correction the next boot should apply, i.e.
the caller's ``cached_residual + (predicted_now - realized_now)``. Negative
inputs clamp to 0 (a prediction that came in UNDER realized needs no
correction; it is already conservative).

Returns the persisted value.

Budget assembly — sum the per-pool predictions into a KV page budget.

:func:`compute_kv_budget` folds every ``predict_*_bytes`` contribution
into ``non_kv_total_bytes`` and derives the KV page count, raising
:class:`BudgetExceeded` with the per-pool breakdown when the sum exceeds
the usable VRAM. Also holds the paged-KV per-page sizing, the pool-member
KV caps, the deferred-KV sizing, and the post-capture KV-grow helpers.

The one line a refusal on a SHARED card needs, or "" on an unshared one.

A budget that cannot fit is read as a configuration problem, and every
remedy the message offers is a knob on this engine. None of them applies
when the bytes belong to a process this engine does not own, so the
breakdown has to say whose they are and what the actual remedy is.

Pure budget arithmetic — no CUDA, no engine state.

Sums every ``predict_*_bytes(...)`` contribution into
``non_kv_total_bytes``, then subtracts from
``total_bytes * gpu_memory_utilization`` to derive ``kv_bytes``.

Args:
    total_bytes: total device VRAM (binary GiB × 1024**3).
    weights_bytes: bytes used by loaded model weights (post-flat-
        load; ``torch.cuda.memory_allocated``).
    activation_reserve_bytes: the single activation reserve line
        item — always ``profiled_peak * safety_factor`` (or the
        cold-boot closed-form upper bound). Whether the runtime
        :class:`ActivationArena` standing buffer is allocated or
        not, this is the same number; the activation traffic is
        either standing (arena on) or transient (arena off), but
        the budget reserves the same bytes either way. See
        :func:`predict_activation_reserve_bytes`.
    per_page_bytes: per-page byte cost of the KV pool — read from
        the active backend's ``bytes_per_token``, never derived
        from bit widths in this module.
    gpu_memory_utilization: fraction of total VRAM the engine is
        allowed to claim.
    extra_safety_bytes: optional additional reserve (maps to
        ``--activation-headroom-gb``). Default 0.
    external_reserved_bytes: physical THIS PROCESS holds that its own torch
        caching-allocator reservation does not account for — the CUDA
        primary context, the driver's cubin/kernel-stack residency, and the
        regions we cuMem-map directly. Booked as ``driver.cuda_context``.
    foreign_process_bytes: VRAM OTHER PROCESSES hold on this device. Booked
        as its own ``driver.foreign_process`` row, NOT folded into
        ``driver.cuda_context``: it is charged (``gpu_memory_utilization``
        is a ceiling on TOTAL device usage, so a co-tenant's bytes are
        bytes this engine may not claim) but it is not ours, it does not
        scale with TP the way a per-rank context does, and an operator
        reading the breakdown has to be able to tell the two apart.
    per_pool_predicted_bytes: dict mapping pool name → predicted
        bytes (from ``predict_*_bytes`` calls). Summed into
        ``non_kv_total_bytes`` along with ``weights_bytes`` and
        the activation reserve.
    serving_activation_floor_bytes: the irreducible SERVING-phase
        activation floor (max non-prefill serving shape: decode /
        chunked-mixed / stochastic-verify). Booked as the
        ``transient.serving_step`` row, so the pool this call sizes leaves
        it FREE on top of every resident pool — a serving step's fresh
        allocations come out of driver-free VRAM, not out of the arena's
        mapped physical. Also derives ``serving_kv_ceiling_pages``, the
        post-capture grow's VA ceiling, by substituting the full prefill
        ``activation_reserve`` (a one-time boot transient) for the
        serving-resident terms. Zero (the default) reserves nothing and
        makes the serving ceiling equal ``num_pages``.
    serving_cumem_invisible_reserve_bytes: bytes resident during serving
        that are invisible to both torch's caching allocator and
        ``CuMemPoolAllocator.mapped_bytes`` (the extra TP per-rank CUDA
        contexts + NCCL workspace — the ``_staging`` gather buffers are
        now a visible cuMem ``nccl_staging_pool`` and are NOT in this
        reserve; see :func:`cumem_invisible_serving_reserve_bytes`).
        Held OUT of the
        serving KV ceiling so the post-capture grow does not plan into
        VRAM the driver re-occupies after the prefill arena frees. Zero
        (the default / TP1) leaves the legacy ceiling unchanged.

Returns:
    :class:`KvBudget` with every intermediate exposed.

Raises:
    :class:`BudgetExceeded` when Σ predictions > usable budget.
    :class:`MemoryBudgetError` when ``kv_bytes < per_page_bytes``
        (can't fit even one KV page).

Read ``per_page_bytes`` off the active PAGED_KV backend.

Sums per-layer ``bytes_per_token`` (already per-rank when
``num_kv_heads_per_rank = num_kv_heads // tp_size``) over every
PAGED_KV layer, then multiplies by ``block_size``. Hetero-geometry
models (Gemma 4 sliding vs full layers) get the correct sum
because each layer is queried with its own (num_kv_heads, head_dim).

The TKV / bf16 variants implement ``bytes_per_token`` directly —
we never hardcode bit widths here. TKV's per-page codebook tile
(when present) is included in the backend's reported
``bytes_per_token`` because the backend's ``kv_cache_shape``
already returns the slot byte budget that includes any per-slot
codebook tile.

A page is not only its KV slot. A sparse-attention layer's index-key
stream and the pooled per-page rows derived from it are allocated by
:class:`~arbi_serve.cache.paged_kv_pool.PagedKVStatePool` out of the same
page budget, and the backend cannot report them — it describes the KV
slot, not the layer's indexer. They are added here from the same
:func:`~arbi_serve.cache.paged_kv_pool.index_stream_page_bytes` the pool
charges itself, so the number this returns IS what a page costs. It has
to be: every KV sizing path (the profiler ceiling, the deferred resize,
the post-capture grow, the wake re-provision) divides free VRAM by this
to decide how many pages to allocate, and the allocation that follows is
the pool's, not the backend's.

True when a stable-VA pool spans more than one distinct ``model.path``.

In a VRAM on every swap (``drop_parked_member_weights``), so only ONE
member's weights are ever resident — each member serves as the SOLE VRAM
resident. A SAME-model config-variant pool (every member reuses the boot
path, ``PoolMemberConfig.path=None``) donor-shares one weight set and can
co-reside, so it is NOT distinct.

Reads the boot ``model.path`` plus each ``model.path``).
Pure config inspection — no engine state.

The KV-page CEILING for a multi-model residency-pool member, or
``None`` when the boot may grow KV into the full VRAM budget.

A POOL MEMBER that genuinely CO-RESIDES with other members must not size
KV beyond what its own config can ever address — ``max_batch`` sequences
× full ``max_context`` each (+1 partial page per seq, +1 null-sentinel
page), plus the non-servable head of a growable capture-backed slab. Pages
past that are pure waste that steals the pool-mates' budget and wake-remap
headroom. Mirrors ``active.py``'s ``batch_floor`` arithmetic and
:func:`kv_pages_floor_for_context`'s capture-prefix accounting.

The cap applies ONLY when there is real co-residency to protect:

  * ``cfg.pool_members`` is non-empty — the operator declared sibling
    models that will co-reside, OR
  * ``member_build_active`` — this build IS a runtime member-add
    (``build_member_into_engine``), which sizes against the live
    co-resident state.

When ``distinct_model_pool`` is set the pool
spans >1 DROPS the non-active members' weights on
every swap and each member serves as the SOLE VRAM resident — there is no
co-resident pool-mate whose budget the cap protects. Capping such a member
at its tiny config-addressable count refuses the KV a big-KV member can
hold on the whole card while the others' weights are dropped (the phantom
max-weights + max-KV cross-product that never co-occurs). Size it against
its OWN single-resident footprint instead (uncapped, like a solo model / a
swap-in that fully evicts the outgoing member). The cap is KEPT for
SAME-model config-variant pools, which donor-share one weight set and can
co-reside.

A SOLO boot — residency on (it is ON by default for the runtime-reconfig
surface) but NO declared members and NOT a member build — is right to
grow KV into ALL the budgeted VRAM (more pages = more cache /
concurrency), exactly like a non-residency model. Capping it at the tiny
config-addressable count was a footgun: on a low-concurrency config
(``max_batch`` small, large ``max_context``) it starved KV to a fraction
of a GiB and left GiB of VRAM idle, then the post-capture grow could not
expand past the cap. When a variant member is later added at
runtime, that member's build re-profiles against measured free VRAM and
sizes itself within the co-resident budget — the primary is not forced to
pre-shrink at boot for a sibling that may never arrive.

Slack that a map round-up at the KV/holdout boundary can consume.

``cuMemMap`` rounds sizes and offsets up to the device allocation
granularity, so the slab's end and the holdout's start can each shift by at
most one granule. Two granules is therefore the bound, and it is a
MEASUREMENT of the driver rather than a chosen cushion — ask the allocator
for its granularity and multiply.

``None`` / non-positive means no driver answered, and only then does
:data:`DEFERRED_KV_MARGIN_BYTES` stand in.

Rows the verify SamplingParams chain widens to fp32 at once.

Reads the engage threshold and the chunk step from the sampler module that
owns them, so the reserve and the loop it prices cannot drift into two
answers about which regime a given slate takes. Mirrors
``_apply_sampling_params_uniform_block``'s own gate, including its clamp of
the per-slot step to the slate it is walking.

The chunked branch there needs a persistent destination, which exists on
every boot this reserve is non-zero for (the destination and the vocab
width come from the same ``VerifyBuffers``), so the gate is the row count
alone.

Analytic peak bytes of the stochastic MTP verify tail at max_batch.

The stochastic rejection sampler runs an EAGER tail after the verify
forward where several vocab-scale fp32 tensors are co-resident. This is
NOT decode-class — it scales with concurrency (B) AND vocab (V), and the
forward profiler never runs the sampler. Used as a LOWER BOUND for the
serving floor when the measured profile is unavailable, and by the boot
gate that refuses a config that can't serve stochastic verify at B=max_batch.

TWO terms, on two different row counts, because the tail holds two kinds of
tensor and only one of them is slate-scale:

  * the RESIDENT set (:data:`MTP_VERIFY_STOCHASTIC_VOCAB_TENSORS`) — the
    ``VerifyBuffers`` vocab scratch plus the drafter-q carry — over the
    whole ``(K+1) * B`` slate;
  * the SamplingParams chain's fresh intermediates
    (:data:`MTP_VERIFY_SAMPLING_CHAIN_VOCAB_TENSORS`) over the rows that
    chain actually widens at once, which is one chunk on a slate wide
    enough to chunk and the WHOLE slate below that threshold
    (:func:`_verify_sampling_chain_rows`).

Folding the second into the first as a flat multiple of the slate is what a
single constant does, and it cannot be right in both regimes: a slate too
narrow to chunk holds the chain's blocks at full slate width, so a multiple
fitted to the chunked regime under-reserves there — reachable on an ordinary
``max_batch`` with a shallow draft depth, and the OOM direction.

``draft_depth`` is the CHAIN depth every caller has to hand; the row
count those tensors actually carry is resolved from it by
:func:`~arbi_serve.spec_decode.tree_spec.verify_slate_k` — the same
function ``VerifyBuffers`` sizes its vocab-scale scratch with. A tree
proposes one draft per NODE, so booking the chain's rows against a
tree's slate reserves a fraction of what the step holds.

``dense_draft_q`` is the caller's answer to
:func:`~arbi_serve.spec_decode._mtp_driver_ops.dense_draft_q_possible`.
``False`` drops :data:`MTP_VERIFY_DENSE_Q_VOCAB_TENSORS` from the resident
term, matching a pool that skipped ``sc_draft_probs``. It defaults to the
booking (``True``) so a caller that cannot resolve it over-reserves.

No CUDA and no engine state; the slate width is config (the tree flag).

A conservative serving-phase activation floor at PROFILE time.

``profile_and_size_kv_pool`` runs BEFORE the per-shape activation profile
(``eng.activation_profile``) exists, so it cannot read the measured
decode / chunked-mixed peaks the post-capture grow uses
(:func:`serving_activation_floor_bytes` via ``serving_floor_for_grow``).
This derives a closed-form UPPER bound on the serving floor from what IS
known at profile time, used ONLY to set the growable-region VA ceiling
(:attr:`KvBudget.serving_kv_ceiling_pages`) — never the physical sizing.

The serving forward never runs the FULL ``max_batched_tokens`` prefill:
chunked prefill caps a live prefill step at ``chunk_prefill`` tokens
(default 2048 < 8192), so the serving activation is bounded by the chunk
fraction of the measured full-prefill peak (activation scales ~linearly in
tokens for the dominant QKV / MLP GEMM scratch). ``arena_footprint_bytes``
is then subtracted from that fraction, because the serving prefill runs
inside ``scratch.forward_arena`` and the budget books that pool separately:
what belongs on THIS row is only what a step needs ON TOP of the arena. The
stochastic MTP verify tail is folded in as a max (vocab-scale, not
token-scale). Floored at the allocator granularity.

Being an UPPER bound is intentional and safe: the VA ceiling only RESERVES
address space (~0 physical); the actual mapped pages are bound at grow time
by the MEASURED ``serving_floor_for_grow`` AND the gmu-on-total cap. A
too-generous ceiling just lets the grow fill more of the genuinely-free
VRAM; a too-tight one would strand it (the bug we are killing). Pure
arithmetic — no CUDA, no engine state.

Free VRAM the post-capture KV grow must LEAVE for serving — the real
serving-time activation peak, NOT the one-time capture/prefill peak.

The post-capture KV grow reclaims (final_free − floor). The floor is the
safety boundary: it must cover the live serving step's transients (the
per-step decode / verify forward, cuBLAS workspace that expands on first
decode, sampler scratch, RoPE-cache growth) so a live decode can never OOM
the card after the Phase-2 freeze locks the layout.

Critically it must EXCLUDE the capture-time forward-activation peak — the
full-prefill-width arena (``activation_reserve``, profiled at
``max_batched_tokens``) is a one-time boot transient that tears down after
the capture sweep; holding it free would strand multiple GiB. The serving
floor is derived from the MAX over the measured non-PREFILL serving shapes
(``SHAPE_DECODE``, ``SHAPE_MIXED``, ``SHAPE_VERIFY_STOCHASTIC``). The
stochastic MTP verify tail is NOT decode-class: at concurrency with
temperature>0 the rejection sampler holds several vocab-scale ``(K+1, B, V)``
fp32 tensors co-resident, scaling with both concurrency and vocab.
``decode_shape_peak_bytes`` is that MAX (the caller
already folds in scratch + the profiler's safety margin).

Floored at ``allocator_floor_bytes`` (the granularity / fragmentation floor
the post-capture gate verifies) so we never reclaim below the allocator's
own minimum even when the decode peak measures tiny. ``margin_bytes`` is the
round-up slack so a boundary allocation can't nudge KV past the holdout.

``stochastic_verify_peak_bytes`` is an analytic LOWER BOUND on the stochastic
verify tail (:func:`mtp_verify_stochastic_peak_bytes`) folded into the max as
a belt-and-suspenders guard: it keeps the floor above the vocab-scale tail
even if the measured profile under-reported it (the measured value normally
dominates via ``decode_shape_peak_bytes``). Zero on greedy-only / no-MTP.

Pure arithmetic — no CUDA, no engine state.

Target KV mapped-page count for a reclaim grow: claim (free − floor).

Maps as many extra pages as fit in ``free_bytes − serving_floor_bytes``,
capped at the reserved-VA ceiling (``profiled_max_pages``) and never below
the currently-mapped count. Pure arithmetic — the single source of truth
for BOTH the post-capture grow and the pre-freeze second reclaim pass so
the two stages size identically. No CUDA, no engine state.

``mapped_bytes_for_pages`` PRICES a page target in the physical the grow
actually maps
(:meth:`~arbi_serve.cache._growable_kv.GrowableKvMixin.kv_mapped_bytes_for_pages`).
It is not ``pages × per_page_bytes``: the slab maps one granularity-rounded
``GrowableRegion`` prefix PER LAYER, so the physical is a step function of
the page count and the two counts differ by up to one granule per layer in
either direction. Sizing against ``free`` with the capacity number is wrong
in both: too high and the grow overdraws the serving floor; too low and the
leftover reads as a late release to the ``kv_sizing_late_growth`` detector,
which then claims a page whose real cost is a whole granule per layer.
Because the cost is monotone in the target this bisects it. Left ``None``
(a pool that cannot answer) the capacity arithmetic stands, so a pool with
no per-layer regions is unaffected.

Re-cut a KV page ceiling by the bytes a plan row turned out to be short.

:func:`compute_kv_budget` hands KV ``usable - sum(per_pool_predicted)``, so
a row priced BELOW what the boot then makes resident is, page for page,
KV the plan granted against bytes the card does not have. When the boot
measures that row directly, ``measured - planned`` is exactly how much of
the ceiling was cut too high, and this takes it back.

A CEILING re-cut, NOT a floor row, and that distinction is the whole
correctness argument. The bytes are already RESIDENT by the time a sizing
pass reads free VRAM, so they are outside that reading; subtracting them
from the free side as well would hold the same physical twice and shrink KV
on every boot. The ceiling is the other side of the plan's arithmetic —
where the term was booked as a forecast — and that is the side to correct.

Rounds the page count UP: the pages come out of a footprint that is a
bound, so a partial page is held rather than handed back.

Pure arithmetic — no CUDA, no engine state. Identity whenever the
measurement does not exceed the plan, so a row the plan already covers
leaves the ceiling byte-for-byte where it was.

Re-cut the KV page ceiling against the overhang bound this boot MEASURED.

``ceiling_pages`` was cut by :func:`compute_kv_budget`, whose activation
footprint is ``max(activation_reserve, serving_torch_caching_overhang)`` —
one physical footprint counted once. That plan runs BEFORE the capture
sweep grows most of the overhang, so its overhang term is a previous boot's
reading or the cold seed. When the boot then measures MORE
(:func:`~arbi_serve.engine.memory_budget.reserves.
serving_torch_caching_overhang_grow_bound_bytes`), the plan promised KV
exactly ``bound - planned`` bytes that the caching allocator is holding, so
the ceiling was cut that much too high. Take the pages back.

A CEILING re-cut, NOT a floor row. These bytes are already RESIDENT when
the grow reads free VRAM, so they are outside that reading already;
subtracting them from the free side as well would hold the same physical
twice and shrink KV on every boot. The ceiling is the other side of the
plan's arithmetic, where the term was booked as a forecast, and that is the
side that has to be corrected.

Rounds the page count UP: the pages come out of a footprint that is a
bound, so a partial page is held rather than handed back.

Pure arithmetic — no CUDA, no engine state. Identity when the bound does
not exceed the plan, which is every boot whose budget cache carries the
measurement.

The re-cut itself is :func:`kv_ceiling_pages_for_measured_excess`, which
every plan row measured against its own realization shares.

Reconcile the DRIVER free against the accounted residual at freeze time.

The Phase-2 freeze reports the driver free (``mem_get_info`` — ground
truth). The accounted free is ``physical − pooled − growable-KV − captured-
graphs``: weights live in the cuMem ``weights_pool`` (in ``pooled``), KV in
``growable-KV``. When all three are summed the accounted free should track
the driver free within ``tolerance_bytes``.

``cudagraph_mapped_bytes`` must be ONLY the captured bytes that are NOT
already in ``pooled_mapped_bytes``. On the cuMem path (the default) that is
ZERO: ``capture.cudagraphs`` / ``capture.io_buffers`` are cuMem-backed
``NamedMemPool``s, so their physical carries their tags and is already
inside ``mapped_bytes(None)``. Passing the full captured figure there
subtracts the same physical twice and drives accounted-free below the truth
by the whole capture footprint. The parameter is non-zero only on the
``--no-cumem-pools`` / CPU fallback, where captures do go through PyTorch's
graph-private cudaMalloc and are genuinely invisible to ``mapped_bytes``.

A gap beyond tolerance means resident physical the accounting omits (a
leaked pool) — a defect to flag, never a serving risk. Returns the accounted
free, the signed gap (driver − accounted), and whether it reconciled. Pure
arithmetic — no CUDA, no engine state.

Evaluate the Phase-2 memory identity against one live device reading.

The freeze states two things about the card and then stops looking: the
serving floor is held out of the driver free it measured, and that free
reconciles against the accounted residual (:func:`reconcile_serving_free`)
up to a gap the CUDA context and default-heap scratch always leave. Both
statements are re-read here from the same terms, relative to the freeze's
own reading, so physical that appears during serving under no accounted
name shows up as a signed figure rather than as the next step's OOM.

  * ``free_minus_floor_bytes`` — driver free minus the line at which the
    floor is spent (``free_at_freeze − floor``): what the floor's consumers
    have left. It is the serving-time reading of the consumption
    ``assert_post_freeze_floor_intact`` measures once at boot; negative
    means they have drawn past what the KV grow held out for them.
  * ``reconcile_gap_bytes`` — the live reconcile gap minus the gap the
    freeze recorded: bytes mapped since the freeze that neither the cuMem
    pools nor the KV regions name. Zero while every byte the driver lost
    is a byte a pool mapped.

``violated`` is either bound broken: the floor overdrawn, or the gap moved
by more than ``tolerance_bytes``. Pure arithmetic — no CUDA, no engine
state.

Bytes the deferred KV resize must HOLD OUT of free VRAM before sizing.

The deferred resize (build Phase 4c) runs FINAL pool (captured graphs bake
the KV slab data_ptrs) — so when a capture sweep is still pending
against a fully-physical pool, the resize must leave the sweep's whole
working set free or capture OOMs the card and the engine "serves" at
~0 B free (first eager forward then dies at ``cublasCreate``):

  * ``graph_pool_reserve_bytes`` — the ``graph_pool`` budget the
    profiler already computed; never invented here);
  * ``capture_transient_bytes`` — the capture-time forward-activation
    transient (the profiler's ``activation_reserve`` plan line: the
    warmup + capture forwards allocate it from free VRAM because the
    resize's ``empty_cache`` returned every cached segment);
  * ``post_capture_floor_bytes`` — the runtime free-VRAM floor the
    post-capture gate verifies;
  * ``margin_bytes`` — allocator-granularity slack.

When NO capture will run against the resized pool (``capture_pending``
False) — or the pool is growable prefix-backed (``kv_prefix_backed``;
the capture transient lives in the unmapped tail by construction) —
only ``baseline_reserve_bytes`` is held out, preserving the legacy
sizing on those paths.

``member_graph_floor_bytes`` (residency-pool member builds: driver-side
``cudaGraphInstantiate`` headroom per resident record) is additive on
every path.

``instantiate_reserve_bytes`` (``n_captured_graphs × per-graph
``cudaGraphInstantiate`` executable-graph reserve, ``build._PER_GRAPH_
INSTANTIATE_BYTES``) is the driver memory each captured stream's
instantiation allocates on the DEFAULT device heap — OUTSIDE the cuMem
``graph_pool`` tag, so the map-time cap does NOT contain it and it draws
straight from free VRAM during the sweep. The pre-capture feasibility gate
(:func:`build.assert_capture_sweep_fits`) already counts it against the
headroom it demands. Holding it out HERE makes the grow and the gate
agree: KV stops eating the instantiate reserve and the sweep's driver
allocations fit. Additive on every path (like the member floor); the
caller passes 0 when no capture runs.

Pure arithmetic — no CUDA, no engine state.

Pages the deferred resize may allocate: ``(free − reserve) // per_page``
capped at the profiler's ceiling (the operator's
``gpu_memory_utilization`` plan). When measured free matches the plan,
the ceiling binds and the KV is allocated; when boot-time
reality ate more than the plan's slack (compile artifacts,
fragmentation), KV shrinks by exactly the shortfall — never the capture
sweep's holdout. Pure arithmetic.

Bytes a residency-pool MEMBER build may hand its KV pool.

Derived from MEASURED free VRAM at build time, never from the boot-style
``total × gmu`` plan: whatever the other resident members still hold
mapped (weights, parked pools, graph-internal driver memory) is already
inside ``total − free``. Two binds, the smaller wins:

  * ``free − reserve`` — the capture sweep + runtime floor holdout must
    survive the KV allocation;
  * the operator's ``gpu_memory_utilization`` as a ceiling on TOTAL
    device usage: KV may only grow until ``used + kv == total × gmu``.

Pure arithmetic — no CUDA, no engine state.

The closed KV re-provision ledger a wake / reinstate sizes against.

The wake counterpart of the boot ``KvBudget`` / deferred-resize ledger. Every
byte of measured free VRAM is accounted: it is either the KV re-provision
``budget_bytes``, a held-out reserve term, the ``gpu_memory_utilization``
headroom the operator asked to leave, or the ``unmodelled_residual_bytes``
(0 when every term is named; negative when the reserve exceeds free, i.e. the
wake is starved). ``coresidency_residual_bytes`` is the outgoing member's
OUT_OF_ARENA default-allocator residual — already excluded from
``free_bytes``, carried here so the ledger names WHY free is below a solo card.

Closed-ledger KV re-provision budget for a wake / reinstate.

The wake counterpart to the boot deferred resize
(:func:`deferred_kv_reserve_bytes` + :func:`member_kv_available_bytes`): it
holds the SAME named terms out of measured free and derives the page budget
through the SAME free∧gmu bind, so a wake can never silently miss a reserve
the boot ledger accounts for. It differs from the deferred resize ONLY in the
terms a wake does not incur — and names each as zero rather than dropping it:

  * capture pool + capture transient + torch-caching overhang — ZERO on wake.
    No capture sweep runs (captured graphs are RE-INSTANTIATED, not
    re-recorded) and the cuMem capture pool is already remapped at its stable
    VA by ``wake_all`` BEFORE this budget is read — so those bytes are already
    resident and already netted from ``free_bytes``, never held out again.

The held-out terms:

  * ``serving_floor_bytes`` — the serving-time activation peak the woken
    member must keep free (``serving_floor_for_grow``); already folds in the
    post-capture runtime floor and the residency member-graph floor.
  * ``instantiate_reserve_bytes`` — the ``cudaGraphInstantiate`` executable-
    graph reserve ``_reinstantiate_graph_execs`` allocates AFTER the KV
    re-provision, on the DEFAULT device heap OUTSIDE the cuMem arena (the same
    term the boot deferred resize + capture gate hold out). Zero when
    ``cuda_graphs`` is off / no captured graphs.
  * ``state_arena_bytes`` — the sentinel-alias recurrent arenas (GDN /
    Mamba slabs) the wake re-maps AFTER the KV re-provision. They own their
    own VA reservations, so no allocator counter sees them and the free
    reading this ledger divides up still holds the bytes they are about to
    take.
  * ``per_tensor_restore_bytes`` — the registered per-tensor sleep set,
    re-mapped last of all.
  * ``margin_bytes`` — allocator-granularity slack.

Those two middle terms are the whole of what a wake maps after the grow and
the grow does not: everything else it remaps (the cuMem named pools, a
cross-model member's weights) is already back and already netted from
``free_bytes`` by the time this is called. Both are read off the PARKED
member — the region sizes its own wake will re-bind
(``state_arena_parked_wake_bytes`` /
``SleepableTensorPool.parked_wake_bytes``) — so they are measurements of
this wake rather than a forecast or a previous member's reading.

``coresidency_residual_bytes`` is the measured OUT_OF_ARENA sibling residual
(pytorch#145168): already excluded from ``free_bytes``, so it is carried for
the ledger — NOT re-subtracted (that would double-count it).

Pure arithmetic — no CUDA, no engine state.

The smallest KV pool that is honestly servable, derived from config.

A KV pool must hold at least ONE sequence at the full configured
``max_context`` to serve anything beyond a tiny truncated stream. That is
``ceil(max_context / block_size)`` data pages plus the null-sentinel page
at index 0. NOT a hardcoded number — it is exactly the per-request page
window the scheduler reserves for one full-context sequence (mirrors
``active.py``'s per-seq ``per_seq_pages + 1`` floor unit). Below this the
pool cannot even fit one request's worst-case context: the boot is broken,
not merely tight.

THE servable-KV floor: pages the pool must hold to serve ``max_context``.

Single source of truth for "how many KV pages does this config actually
need". Every consumer that has to answer that question — the pre-capture
viability gate, the post-capture viability gate, the live-swap capacity
gate, and the prefill-capture ladder's KV reserve — MUST call this, because
the whole class of bug it exists to kill is consumers disagreeing:

  * a boot refusal computed from a floor of one window while the gate
    that actually decides the boot uses ``capture_prefix + one window``
    would name a ``--max-context`` cap that is arithmetically guaranteed
    to fail on re-launch, and
  * a prefill-capture ladder budget that reserves only
    ``max(MIN_VIABLE_KV_PAGES, max_batch)`` pages of KV leaves the sweep
    free to capture rungs that eat the very KV the ``max_context``
    window requires, so the boot then refuses (or OOMs) on a shortfall
    the ladder created itself.

Three additive/maxed terms, all config-derived, none hardcoded to a dtype:

  * ``capture_prefix_pages`` — non-servable pages held out IN ADDITION to
    the window's own null sentinel. For a live serving pool this is ``0``:
    the ONLY page the pool withholds is page 0 (the null sentinel,
    :data:`SERVING_NULL_SENTINEL_PAGES`), and ``min_useful_kv_pages`` (the
    window term below) ALREADY counts it — the serving free-list is
    ``[1, num_pages)`` and ``grow_to_pages`` hands out every mapped page, so
    a single request reaches pages ``[1, num_pages)``. Passing the
    ``max_batch + 3`` capture-mapping prefix
    (:func:`arbi_serve.engine.inprocess_capture.capture_valid_prefix_pages`)
    here is WRONG — that quantity is a capture-time PHYSICAL-mapping margin,
    never withheld from serving, and double-counting page 0 on top of it.
    Kept as a parameter (rather than dropped) so the signature is stable and
    the term is explicit; live callers pass ``0``.
  * one full ``max_context`` window (:func:`min_useful_kv_pages`) — the
    per-request page window the scheduler reserves for one full-context
    sequence, INCLUDING the shared page-0 null sentinel. The KV bytes/token
    of the live codec is NOT an input here: this is a PAGE count, and the
    codec is priced once, in ``per_page_bytes``, wherever pages are
    converted to bytes.
  * ``MIN_VIABLE_KV_PAGES`` / ``max_batch`` — the legacy count floors,
    kept so this is a strict tightening of every floor it replaces.

Pure arithmetic — no CUDA, no engine state.

THE servable-KV floor when the served context is an OUTPUT, not a pin.

:func:`kv_pages_floor_for_context` answers "how many pages does THIS
``max_context`` need". That question only has an answer when
``max_context`` is a REQUIREMENT — an operator ``--max-context N``, which
:func:`~arbi_serve.engine.build_memory_sizing.assert_kv_pages_floor`
REFUSES below. Under ``--max-context auto`` (or a recipe default) the
served context is the engine's OUTPUT: the same gate NARROWS it to whatever
the realized pool proves and boots. There is then no window to protect, and
a pre-capture consumer that reserves one full window of the PROVISIONAL
sizing bound is reserving for a context the boot has already said it may
not serve.

That is not a conservatism, it is an impossibility. The provisional bound
is the model ceiling clamped by reserved VA — on any card where one such
window outsizes free VRAM (measured: 27B/TP1 exl3-4.0bpw + DFlash2 on a
24 GiB 4090, one 262144-token window = 1025 pages = 4356 MiB against
3122 MiB free at the gate) the reserve can never be satisfied whatever the
consumer does, so the consumer is disabled by construction rather than
making a trade-off.

What an engine-flexed boot must still not lose is the CONCURRENCY the
operator DID pin. ``max_batch`` is a pin at every one of these seams (the
auto-narrow's own words: "concurrency is the pinned knob, the context is
the flexible one"), and the widest step the scheduler may issue for that
batch is ``step_tokens`` (``max_batched_tokens`` / ``chunk_prefill``). A
pool that cannot hold ``max_batch`` sequences of one such step cannot run
the configured batch through a single chunked prefill, and no narrowing
fixes that — it is the "4 KV pages" footgun with the count floor removed.
So the floor is that, never below the legacy count floors:

  * ``max_batch x pages_per_seq(step_tokens)`` + the null sentinel,
  * ``MIN_VIABLE_KV_PAGES`` / ``max_batch``.

Pure arithmetic — no CUDA, no engine state. Every term is config-derived
and moves with the config; none of them is a constant chosen to make a
particular card pass.

The largest ``--max-context`` a KV budget of ``kv_pages_available`` pages
can ACTUALLY serve, plus the slack pages held back to get there.

The exact inverse of :func:`kv_pages_floor_for_context` — and it PROVES it
before returning (the loop below re-evaluates the real floor and steps down
until it fits), so the number a refusal prints is the number the boot's own
gate will accept. That self-check is the point: both KV refusals name a
concrete ``--max-context <= N`` for the operator to re-launch with, and that
N is a PROMISE. The inversion accounts for:

  * the capture-sentinel prefix — a fixed ``max_batch + 3`` pages of the
    growable slab that can never serve a request, and
  * slack: the REALIZED pool only ever lands at or MEASURED free VRAM, which
    the capture pool's prediction error and cuMem/allocator granularity
    round DOWN), so one full window must not consume the entire budget
    with zero pages spare.

So this holds back ``margin_bytes`` (default :data:`DEFERRED_KV_MARGIN_BYTES`,
the SAME allocator-granularity slack the deferred KV resize keeps below its
reserve) worth of pages before inverting. The default is deliberately
CONSERVATIVE and both refusals say so: an under-promise that boots is a
correct answer, an over-promise that OOMs is not. Callers narrowing against
an already-MEASURED, mapped-and-locked pool (no prediction error, no
allocator growth) pass ``margin_bytes=0`` to serve ``pool − prefix``.

It is conservative in a second way that matters on hybrid/long-context
configs: several competing pools (the cudagraph block tables, the DFlash
drafter slabs, ``PiecewiseBuffers``) themselves SCALE with ``max_context``,
so the KV ceiling measured at the REFUSED (too-large) context is smaller
than the ceiling the quoted cap would actually get. Inverting against the
smaller ceiling can only under-promise.

Returns ``(max_context_tokens, margin_pages)``. ``margin_pages`` is 0 when
``per_page_bytes`` is unknown (CPU stub / non-paged KV) — the prefix
subtraction still applies, so the quote is never LESS safe than before.

Pure arithmetic — no CUDA, no engine state.

Decide what the post-capture KV grow should do — pure, no CUDA.

Returns one of:

  * ``"grow"``     — ``target_pages > cur_mapped_pages``: expand the pool.
  * ``"broken"``   — the pool stayed at/below ``cur_mapped_pages`` AND is
    below ``min_useful_pages`` AND there is enough FREE VRAM (past the
    runtime floor + margin) to have mapped at least ``min_useful_pages``.
    This is the silent-degradation footgun: substantial VRAM idle yet KV
    starved — the caller must FAIL LOUD, never serve.
  * ``"constrained"`` — the pool stayed at ``cur_mapped_pages`` because
    the card is genuinely too full to grow it past the floor. A legitimate
    degraded-but-servable state; the caller WARNS (naming the degradation)
    and serves. The page-count VIABILITY gate
    (:func:`arbi_serve.engine.build.assert_kv_pages_floor`) runs AFTER the
    grow and REFUSES when this degraded pool is below the servable floor,
    so a "constrained" verdict here is not the last word.

The ``"broken"`` vs ``"constrained"`` split is what distinguishes a sizing
bug (KV starved while VRAM sits free) from honest scarcity (no VRAM to
grow into). Thresholds are config-derived: ``min_useful_pages`` is one
full-context window (:func:`min_useful_kv_pages`); the free-VRAM test is
whether the headroom above the runtime floor+margin could itself have held
``min_useful_pages``.

Per-pool predicted byte costs — one ``predict_<pool>_bytes`` per named pool.

Each predictor takes ONLY config / shape inputs (no CUDA, no engine
state) and returns an ``int``; the docstrings show the formulas. Covers
the weights, recurrent, MTP-snapshot, DFlash-slots, short-conv, MLA,
TKV-scratch, scratch-GDN, RoPE-cache and LoRA pools plus the DFlash /
tkv-bypass / TKV-prefill transient peaks.

Resolve the per-element byte cost for a given quant method.

``quant_method`` (``str | None``) is the canonical name from the
safetensors ``quantization_config.quant_method`` field
(``"awq"`` / ``"awq_marlin"`` / ``"exl3"`` / ``"fp8"`` /
``"compressed-tensors"`` / ``"nvfp4"`` / ``"gptq"`` / ``"int8"`` /
``None``). Unknown / ``None`` falls back to the dtype byte cost
(bf16 baseline = 2.0).

See ``arbi_serve/weight_quant/awq/__init__.py`` for the canonical
detection; ``arbi_serve/weight_quant/loader.py`` is the runtime
branch this predictor mirrors.

Predicted bytes for :data:`weights_pool`.

Quant-aware. The dense Linear projections (q/k/v/o + gate/up/down)
are sized at ``_bytes_per_elem(quant_method, dtype)`` — bf16 = 2.0,
AWQ-INT4 ≈ 0.5625, FP8 ≈ 1.02, EXL3/GPTQ-INT4 ≈ 0.530, INT8 = 1.0
(each value INCLUDES the per-group scale + zero-point overhead).
LM-head and embeddings size at their own quant method (defaulting
to ``quant_method``); pass the no-quant marker (``None`` /
``"bf16"``) explicitly when the body is INT4 but the embed / lm_head
is bf16 (Q36 case — AWQ leaves the embeddings bf16). RMSNorm
parameters always size at ``dtype`` (LN scales never quantize).

Formula (per rank, after TP sharding):

    bytes = (
          embed_elem * vocab * hidden                  # input embed (bf16 even on AWQ)
        + lm_head_elem * vocab * hidden                # untied lm_head (only when not tied)
        + per_layer_attn_elem * Σ_layer (q + k + v + o param counts)
        + per_layer_mlp_elem * Σ_layer (gate + up + down param counts)
        + dtype_elem * 2 * num_layers * hidden         # rmsnorms
    ) / tp_size

For mixed-quant models (LoRA adapters bf16 on top of INT4 base,
etc.) the caller can sum two predict_weights_pool_bytes calls — one
for each quant slice — to land at a true per-layer total.

See ``arbi_serve/weight_quant/loader.py`` for the runtime path this
mirrors.

Bytes the sentinel-row alias removes from a predicted slab total.

The slabs still SPAN row 0; the granules wholly inside it resolve to one
shared page, so the physical is not there to budget for. Returns 0
whenever no arena would be built — the alias off, or no cuMem driver on
this host — so the prediction always matches what the boot allocates.
See :mod:`arbi_serve.runtime.sentinel_alias_arena`.

Predicted bytes for :data:`recurrent_pool` — base recurrent slabs.

Formula (N = ``max_num_seqs``; slab dim 0 = ``N + 1`` to cover the
zero-sentinel row at index 0 — see :class:`RecurrentStatePool`
module docstring):

    per-layer base = (N + 1) * (num_v_heads_local * head_k_dim * head_v_dim
                                + conv_dim_local * conv_kernel) * elem
    total = Σ_recurrent-layer base

The base slabs are zeroed at admission and persist for the
request's lifetime. The MTP verify snapshot buffers attached on top
of these slabs are a SEPARATE line item —
:func:`predict_mtp_snapshot_pool_bytes`.

Returns 0 when ``layer_specs`` has no recurrent (GDN / Mamba) layers.

``(base_elems, recompute_per_token, replay_per_token)`` for one
Mamba-2 layer, TP-local.

``base_elems`` is the 1-deep ``snap_{ssm,conv}_state`` pair per slab
row. The per-token terms are the retained recurrence INPUTS:
recompute keeps the pre-conv ``[x ‖ B ‖ C]`` stream + raw ``dt``;
masked replay keeps that stream, the POST-conv stream the SSM step
splits, and raw ``dt``. Mamba-2 has no gating pair, so both terms are
strictly smaller than GDN's.

Predicted bytes for :data:`mtp_snapshot_pool` — the per-token MTP
verify rollback buffers on the recurrent layer views.

Formula (T = :func:`~arbi_serve.spec_decode.tree_spec.mtp_snapshot_rows`
of ``mtp_n_draft``; N = ``max_num_seqs``; slab dim 0 = ``N + 1``, the
zero-sentinel row). Under masked replay a GDN layer reserves only the
per-token inputs:

    per-layer = T * (N + 1) * replay_per_token * elem

Mirrors ``attach_mtp_snapshot_buffers`` (``cache/recurrent_pool.py``),
which allocates into the ``recurrent_pool`` MemPool AFTER the budget
sized the pools — which is exactly why it must be a budget line item.
Same geometry source as :func:`predict_recurrent_pool_bytes`.

``T`` is resolved by the SAME function every attach route passes to
that method, so a tree boot cannot book a chain's rows against a
tree's allocation: under ``ARBI_MTP_TREE`` the staging holds one row
per NODE, and a K-derived count under-books it. The tree has no K
axis, so threading the chain depth deeper is not the fix — reading the
one definition is.

MAMBA layers are priced by :func:`_mamba_snapshot_elems`; both slabs
are at the activation dtype (no fp32 recurrent-state opt-in on the
Mamba path), and the Mamba-2 step has no gating pair, so a Mamba
layer's per-token term is smaller than the GDN layer's.

Returns 0 when MTP is off or ``layer_specs`` has no recurrent layers.

Mode-aware (``gdn_mtp_rollback_mode``). ``replay`` prices the
in-kernel-saved raw per-token post-conv ``k``/``v``/``b``/``a`` + pre-conv
``x_conv`` and NO base frame on a GDN layer: its verify forward commits
no state, so the slab itself is the rollback base. ``recompute`` prices
the base frame plus the retained per-token inputs, and a MAMBA layer
keeps its base frame in both modes. The SAME resolution runs at pool
attach
(``attach_mtp_snapshot_buffers``), so the budget line and the
allocation always agree — this is the accounting chain that lets the
per-position-ladder cut GROW the KV budget instead of being
phantom-freed.

Predicted bytes for :data:`drafter_pool` (``model.drafter``) — the
separate spec-decode draft model's persistent WEIGHTS.

Formula (over the draft checkpoint's safetensors shard HEADERS — no
tensor is materialized, no CUDA is touched):

    per_tensor = numel * elem(engine dtype)   if the on-disk tag is a
                                              dense float (F64/F32/F16/BF16)
               = on-disk nbytes               otherwise (EXL3 trellis,
                                              packed int, fp8)
    fused_kv   = Σ over the per-layer k_proj / v_proj / k_norm weights
                 (the second copy ``build_fused_kv_buffers`` stacks)
    bytes      = Σ per_tensor + fused_kv, per-layer projections // tp
                 when ``tp_shard``

Mirrors ``DFlashDraftModel.from_checkpoint``: a tensor a QuantBackend
claims binds at its STORED width (``backend.bind``), a tensor no backend
claims loads dense at the engine dtype (``DFlashQuantLoader``'s
``get_tensor(..., dtype=self.dtype)``). The dense/packed split is
``dflash_quant_load._DENSE_FLOAT_TAGS`` — imported, not restated, so a
wire tag added there cannot silently re-price this pool. Then the
migration (``build_helpers.move_module_tensors_into_pool``) parks all of
it in the ``model.drafter`` cuMem pool. The same sum covers the external
draft-model drafter (``build_external_drafter``), which migrates into the
same pool.

``spec.drafter_slots`` is the SEPARATE per-request draft-KV line item
(:func:`predict_dflash_slots_pool_bytes`); it lands in this same MemPool
but is priced on its own axis, so this predictor must not restate it.

Returns 0 when MTP is off or no draft checkpoint is configured.

Predicted bytes for :data:`dflash_slots_pool` — the DFlash
drafter's persistent :class:`DraftKVSlots` per-layer K/V buffers.

Per-layer-type sizing (k AND v, never TP-sharded — the draft model
and its KV state live whole on rank 0):

    eff      = dflash_effective_context(max_context, num_pages,
                                        page_size, max_num_seqs)
    eff_full = min(eff, full_context_cap)  if the cap bites else eff
    full_len = eff_full + block_size                  # full-attn layer
    ring_len = min(eff, sliding_window) + block_size  # sliding layers
    per_slot = num_kv_heads * head_dim * (
                   n_full * full_len + n_sliding * ring_len)
    bytes    = 2 * (max_num_seqs + 1) * per_slot * elem

The slab's leading dim is ``max_slots + 1``: ``DraftKVSlots.__init__``
appends a TRASH row the device append redirects masked-out writes to.

Each ``full_attention`` draft layer holds the whole effective context
(+ block guard); each ``sliding_attention`` layer holds only its
trailing ``sliding_window`` (a RING) since the block queries can never
attend context older than the window.

``full_context_cap`` (ARBI_DFLASH_DRAFT_CTX_CAP) mirrors
``DraftKVSlots.__init__``: when set and the effective context EXCEEDS it,
the full-attention layers ring to a trailing ``cap`` window too, so
``full_len = min(eff, cap) + block``. The budget MUST apply the same cap as
the real allocation or the reserved bytes drift from the slab.

``max_context`` is the model RoPE ceiling (the raw config gate);
``num_pages`` / ``page_size`` (when given) bound it down to the
EFFECTIVE per-request context the realized paged-KV pool can serve at
``max_num_seqs`` concurrency — exactly what ``build_dflash_drafter``
passes the drafter. Geometry comes from the draft checkpoint's
``config.json`` (:meth:`DFlashConfig.from_dir` — no weights load),
matching ``DraftKVSlots.__init__``. Returns 0 when MTP is off or
``dflash_draft_path`` is unset.

See ``spec_decode/dflash_driver.py`` (``DraftKVSlots.__init__``).

Resolved draft-KV geometry of ONE :class:`DraftKVSlots` slot row.

Mirrors ``DraftKVSlots.__init__``: a ``full_attention`` layer holds the
whole effective context (+ the ``block`` guard band) unless
``full_context_cap`` rings it, and a ``sliding_attention`` layer holds only
its trailing window as a RING. The predictors below read the layer split
off this record so the persistent-pool bound and the per-step transient
bound can never resolve it differently.

Resolve one ``DraftKVSlots`` row's geometry from the draft checkpoint.

``n_draft`` supplies the runtime block width for checkpoints that declare
no ``block_size`` — the slab carries a ``block``-wide guard band, so the
budget must resolve it exactly as the drafter build does.

Bytes of the DFlash draft-context DESTINATION buffers.

Both context assemblies fill one ``(max_num_seqs, n_kv, layer_len_i,
head_dim)`` K and V pair plus the ``(max_num_seqs, layer_len_i)`` int64 key
positions per draft layer — the persistent
:class:`~arbi_serve.spec_decode.dflash_kv_slots.DraftKVSlots` geometry
re-materialised at ``max_num_seqs`` rows, and the dominant draft-step term.

These are PER-STEP tensors: they die with the caller's references when
``DFlashDrafter.draft`` returns. That lifetime is what lets the serving
floor fold :func:`dflash_draft_transient_peak_bytes` against the verify
tail with a MAX on the synchronous path instead of adding it. Retaining
them across steps would make them co-resident with every other serving
transient and the reserve would have to become additive by exactly these
bytes — so this term is resolved here, counted once inside the transient
peak, and pinned by ``tests/test_dflash_assemble_coresidency.py``.
Returns 0 when DFlash is off.

Analytic peak bytes of the DFlash *draft-step* transient working set.

Unlike ``spec.drafter_slots`` (the PERSISTENT per-slot draft KV — a budgeted
:func:`predict_dflash_slots_pool_bytes` pool) none of these tensors appear in
the serving pool table, yet at ``max_batch > 1`` they must fit in the
post-freeze free VRAM or the draft OOMs at RUNTIME. This is the analytic
bound the serving floor reserves for them so the KV pool is sized AROUND the
draft step (deterministic KV shrink) instead of the grow claiming VRAM a live
draft step needs.

EVERY tensor counted here is allocated FRESH on a draft step and freed when
:meth:`DFlashDrafter.draft` returns — the whole set is a step peak, never a
residency. The serving floor relies on that: on the synchronous path the
grow folds this bound against the verify tail with a MAX (the draft runs
after the verify sync), which under-reserves the moment any term here
becomes permanently resident.

``graph_assemble_covered`` is the caller's PROOF that the per-step
destination is unreachable — ``DFlashDrafter.freeze_graph_assemble``
returned True: ``ARBI_DFLASH_GRAPH_ASSEMBLE`` is pinned on for the life of
this member, capture is armed, and the sealed draft-graph pool covers
every ``(B, caps)`` a served step can present, so every draft step
assembles STRAIGHT into a captured graph's own input buffers. Those
buffers are the draft-graph pool's, already priced in ``capture.io_buffers``
by the SAME :func:`dflash_draft_context_buffer_bytes` this term uses, and
already resident before ``grow_kv_after_capture`` takes its single
``mem_get_info`` read. Under the proof, reserving the destination here
reserves that one geometry TWICE — once where it lives and once for a
route that cannot run — so the term is dropped. Default False keeps the
destination reserved, which is what an unproved boot needs: the per-step
destination is reachable there (a shape the pool does not cover, a
disarmed capture, a slate with no active row) and the bound must cover it.

The step allocates in three SEQUENTIAL stages on top of one buffer that
lives across all of them, so the peak is ``resident + max(stages)``:

RESIDENT — the context-assemble DESTINATION, and zero under
``graph_assemble_covered``. ``_assemble_fixed_context``
(host) / ``_assemble_fixed_context_device`` fills padded
``(B, nkv, Ccap_i, hd)`` k/v buffers plus the ``(B, Ccap_i)`` int64 key
positions from each active row's per-layer draft KV. Its geometry is
IDENTICAL to the persistent ``DraftKVSlots`` at B = ``max_num_seqs`` rows —
one whole slots-pool worth, re-materialised. ``DFlashDrafter.draft`` holds
it through the forward AND the lm_head leg, so every later stage stacks on
top, and drops it on return.

Stage A — the assembly's own scratch. Both assemblies blit each row
STRAIGHT from the persistent slab into its destination slice, so no
gathered copy of the context is ever co-resident with the buffers it
feeds. What remains is the key-position bookkeeping of ONE layer at a
time: the device path derives its ``(B, Ccap_i)`` positions from two
int64 candidate vectors and a bool validity mask, and the host path
writes them in place. Plus the O(B) int64 row/slot index vectors the
device path carries across the whole assembly.

Stage B — the forward's per-layer K/V concat. ``forward_block_fixed`` runs
``torch.cat([ctx_k[i], k_n], dim=2)`` while ``ctx_k[i]`` is still live, so
one layer's context is momentarily doubled at ``Ccap_i + block`` width,
alongside that layer's additive attention bias, its ``(B, block, keys)``
bool allow mask, the concatenated key positions, and the
``(B, nh, block, keys)`` mask ``gqa_folded_attend`` materialises when it
broadcasts that bias across the folded query-head groups. ONE layer's pair:
``forward_block_fixed`` drops each layer's concat as soon as the attention
has read it, so the next layer's never allocates beside it.

Stage C — the lm_head leg. The flattened denoise hidden ``(rows, H)`` and
the draft logits ``(rows, V)``.

``embed(blk_ids)`` ``(B, block, H)`` and ``dhid_all`` ``(B, block-1, H)``
predate the assembly and outlive it, so they join the resident term — as
does the TP>1 vocab-parallel all_gather staging (a per-rank shard buffer
``rows*V/tp`` plus the full-vocab gather buffer ``rows*V``), which is sized
LAZILY at the FIRST draft step into the PERSISTENT staging arenas and then
stays mapped, so a later step's gather stacks on top of it.

``rows = B * (block-1)`` is the flattened draft-token count. ``vocab_size``
is the width the DRAFTER projects — every term it feeds (the logit row and
the TP all_gather staging that carries it) is a draft-side tensor, so under
``--draft-vocab-prefix`` it is the sliced ``N`` and not the verify
vocabulary (:func:`~arbi_serve.spec_decode.drafter.resolve_draft_vocab_size`).
``device_slots`` selects which assembly runs (``ARBI_DFLASH_DEVICE_SLOTS``);
the default covers the device path, whose gather copies every layer. Returns
0 when DFlash is off. Pure arithmetic — no CUDA, no weights load (drafter
geometry from the checkpoint ``config.json`` via
:meth:`DFlashConfig.from_dir`).

Rows ONE :meth:`~arbi_serve.spec_decode.dflash.DFlashDraftModel.
project_context_stacked` pass may take, so its peak stops scaling with
``max_batched_tokens``.

The projection has NO cross-row reduction — ``fc``, the hidden norm, the
K/V projections, the K-RMSNorm (which reduces over ``head_dim``, the last
axis) and RoPE are all per-row — so splitting the row axis is numerically
EXACT, not a reassociation. Tiling therefore costs nothing but kernel
launches, and it is the only lever that reaches the largest term: a
quantized ``fc`` stages a converted copy of the whole tapped block, five
times the hidden state, and tiling shrinks that copy along with everything
else without needing the quant kernel to consume a view in place.

WHERE THE TILE COMES FROM, and why it is not a number someone liked:

  * ``target_peak_bytes`` — the reserve that EXISTS ANYWAY. The serving
    floor's draft row is a MAX over this projection and the draft
    FORWARD's step peak (:func:`dflash_draft_transient_peak_bytes`), so
    tiling below the draft forward's peak buys nothing: those bytes are
    held either way. Tiling to it makes the projection stop being the
    binding term BY CONSTRUCTION, which is the whole objective. Going
    smaller trades launches for VRAM nobody gets back.
  * ``per_row_bytes`` — :func:`dflash_context_projection_peak_bytes` at one
    row, so the tile follows the drafter's geometry and dtype rather than
    assuming this checkpoint's.
  * ``min_rows`` — the floor the tile may not cross. On an EXL3 drafter the
    GEMM picks its leg on the row count
    (``_auto_reconstruct_threshold``): a tile that fell below it would
    change which kernel runs, which changes both speed and numerics. The
    tile is a memory decision and must not become a dispatch decision.
  * ``row_bound`` — no tiling below it. A step that already fits takes the
    single-shot path unchanged, which is what keeps the fixed-shape,
    loop-free contract the captured draft graphs rely on.

The arithmetic itself is :func:`~arbi_serve.runtime.row_tiling.
rows_per_tile`; what lives here is the POLICY — which budget, which per-row
size, which floor — because those are properties of this drafter and this
reserve, while the division is not. Returns ``row_bound`` (no tiling) when
any input is missing, so an unpriceable tile degrades to today's behaviour
rather than to a guess.

Peak fresh bytes ONE ``lm_head`` call over ``rows`` post-norm rows holds.

Formula::

    bytes = rows * vocab * out_elem
          + (rows * vocab + rows * hidden) * kernel_elem   [when the
                                                            kernel computes
                                                            in its own dtype]
          + rows * gather_vocab * out_elem                 [when the vocab
                                                            axis is sharded
                                                            and the head
                                                            gathers]

An MTP verify step runs the head over the WHOLE slate — ``max_batch x
(K + 1)`` rows, one per verify position — through
:meth:`~arbi_serve.models._layer_stack_model_mixin.LayerStackModelMixin.
logits_from_hidden`, which passes ``batch=None`` and therefore binds no
``logits_out`` destination. Every byte below is a fresh heap allocation on
the served step, outside the captured pool and outside every activation
probe: the decode and mixed shapes run the head at ONE row per sequence,
so nothing in the profile has ever seen the slate width.

``kernel_elem_bytes`` is the whole difference between a head that hands
back what it computed and one that converts. A backend whose GEMM writes
the model dtype directly (AWQ / Marlin / fp8 / dense bf16) passes 0 and
the peak is the one output block. A backend that computes in its OWN dtype
— the EXL3 trellis kernel is fp16 (:attr:`~arbi_serve.models.linear.
LinearBase.kernel_compute_dtype`) — casts the input in and casts the result
back out, so the converted input, the kernel's block and the returned block
are all live across the conversion.

The vocab passed here is the vocab the KERNEL writes, not the tokenizer's:
a trellis head Hadamard-aligns its output axis and allocates the padded
width before the caller's slice narrows it. ``gather_vocab_width`` is the
SECOND width a sharded head holds — its rank's block from the GEMM and the
full-vocab block the all-gather lands in are both live — and is 0 on every
head that returns what it computed.

Returns 0 when any dimension is missing, which is a head that cannot be
priced from geometry — never a head that costs nothing.

Predicted bytes for :data:`short_conv_pool`.

Formula (per layer; slab dim 0 = ``max_num_seqs + 1`` to cover
the zero-sentinel row at index 0):

    bytes = (max_num_seqs + 1) * (conv_dim // tp_size) * (conv_kernel - 1) * elem

Sums across every ``StateKind.SHORT_CONV`` layer. Returns 0 for
models with no short-conv layers (Q36 has none).

See ``cache/short_conv_pool.py:127``.

Predicted bytes for :data:`dsv4_pool` at full residency.

Per DSV4_SPARSE layer, per slot row (``max_num_seqs + 1`` to cover the
reserved sentinel at row 0)::

    window * slot_bytes                                  the ring
  + ceil(max_context / ratio) * slot_bytes               the stream
  + ceil(max_context / ratio) * index_slot_bytes         the indexer's
  + 2 * pool_span * pool_width * 4                       accumulator pairs

The two stores span whole GROUPS of slots
(:class:`~arbi_serve.models._deepseek_v4_kv.DSv4SlabAddress`), so their
slot count rounds up to the group boundary the pool lays them out on —
the pool's own :func:`~arbi_serve.cache.dsv4_pool.group_slots_for` off
the same pinned granularity, so this equals what the pool allocates.

This is the RESERVED extent, which is what a budget has to be sized
against even though the pool physically backs only what live requests
have reached — the reservation is virtual, but a budget that assumed
the mapped figure would admit contexts the pool cannot then grow into.

Returns 0 for models with no DSv4 layers.

Predicted bytes for the DSv4 partial-accept rollback buffers.

Mirrors :meth:`~arbi_serve.cache.dsv4_pool.DSv4StatePool.attach_mtp_snapshot_buffers`,
which runs at drafter attach — AFTER the budget sized the pools — so
it has to be a budget line of its own or a DSpark boot reserves less
than it allocates.

Per DSV4_SPARSE layer, per slot row (``max_num_seqs + 1``, covering
the reserved sentinel), with ``T`` =
:func:`~arbi_serve.spec_decode.tree_spec.mtp_snapshot_rows` of
``mtp_n_draft`` — the SAME count the attach fan-out hands this pool,
so a tree boot's node-wide staging is booked as a tree's::

    T * slot_bytes          the ring rows a verify block overwrites
  + T * 8                   those positions, int64

and, on a layer that carries a compressor::

  + 2 * pool_span * pool_width * 4          the accumulator frames
  + 2 * index_pool_span * index_pool_width * 4
  + T * hidden_size * elem                  the hidden the replay pools

Returns 0 when MTP is off or the model has no DSv4 layers.

Predicted bytes for :data:`mla_pool`.

Formula (one stacked latent slab over the MLA layers only):

    slot_bytes = elem * (kv_lora_rank + qk_rope_head_dim)
    bytes = num_mla_layers * num_pages * page_size * slot_bytes

Returns 0 when no MLA layers (Q36 case).

Lockstep with ``MLAStatePool.__init__``'s slab shape.

Predicted bytes for :data:`tkv_scratch_pool` (the ``scratch.attn_codec`` pool).

Covers BOTH paged codec modes' pool residents: TKV's ``TQBufferPool``
(the formula below) and tkv-bypass's q-prescale pair + CSR triplet
(see the ``has_tkv_bypass`` branch — lockstep with
``TkvBypassMetadataBuilder.__init__``). Returns 0 when neither is
active. For TKV, per-rank H =
``num_q_heads // tp_size``,
D = ``head_size``, Dpad = ``next_pow2(D)``, fp32 for most fields,
bf16 for ``q_rot_bf16``:

    bytes = 8 * B * M                              # block_table + indices (i32 each, ×2)
          + 8 * N                                  # safe_slots i64
          + 8 * (B + 1)                            # indptr + prefill_cu_seqlens_k (i32 each)
          + 4 * (B + 2 * M + 2 * B + 1)            # bs_scalar / col_range / nb / lpl_rem (small)

where N = ``max_num_tokens``, B = ``max_num_seqs``, M = ``max_blk``. The
pool carries no q/out staging tiles: they had no reader and were removed,
so nothing here scales with the head count.

Only PAGED_KV layers contribute. This is one
``TQBufferPool`` per active TKV backend; there is at most
one TKV backend per engine (``backends/tkv_backend.py:139``).

WHAT THIS DOES NOT COVER. The same pool holds the per-layer codec residents
the core build allocates into it — the calibration tables and the rotation
state (:func:`arbi_serve.engine.tkv_install.ensure_tkv_cores_built`). Those
are not priced here and cannot be: under an OSCAR rotation the state is
per-KV-head and which matrices exist depends on the rotation file's sides,
the o_proj fold and which dtype casts a path reached. The boot MEASURES
that residency at the install instead
(:func:`arbi_serve.engine.tkv_install.record_attn_codec_basis_bytes`) and
holds it out of the KV ceilings sized after it, so this predictor stays the
closed form for the part that has one.

Analytic peak bytes of the ``TkvBypassBackend`` per-STEP split-K scratch.

The tkv-bypass-SPECIFIC sibling of the codec-independent verify-churn headroom
(:func:`mtp_verify_stochastic_peak_bytes` × :data:`_MTP_VERIFY_SERVING_CHURN_FACTOR`,
reserved by the serving floor for every stochastic-MTP config). This term
covers ONLY the extra attention-side workspace the raw-bf16 turbo split-K
kernel allocates fresh each decode / MTP-verify step
(``tq_splitk_batch_decode`` / ``tq_splitk_batch_mtp_decode``) with no
persistent pool:

  * ``partial_o`` fp32 split-K reduction workspace
    (``_turbo_attn_simt.cu``): ``units × H_q × D × 4``.
  * ``partial_lse`` fp32: ``units × H_q × 2 × 4``.
  * KV-scatter temps
    (``2 × B × M × H_kv × D × 2 + 2 × B × M × 8``).

``units`` is ``batch × num_splits × block_m`` maximised over the two arms
(see the body): they cannot co-peak, because a step is either decode-shaped
or verify-shaped.

``M`` is the verify step's ``block_m`` — the SERVED slate width plus the
committed token, resolved by
:func:`~arbi_serve.spec_decode.tree_spec.verify_slate_k` rather than
derived from ``mtp_n_draft`` here. A tree's verify block is one row per
NODE, and the chain depth it was built from does not describe it.

The ``q_scaled`` fp32+bf16 prescale pair is NOT in this reserve — it
lives in the builder-owned ``scratch.attn_codec`` pool residents
(``TkvBypassMetadataBuilder`` — sized by
``predict_tkv_scratch_pool_bytes``), and the kernel writes the
engine's forward output buffer directly via ``out=`` (no bf16 result
staging alloc + copy).

``num_sms`` is what separates the verify arm's per-DEVICE peak from a
per-batch one; pass it. Omitted (``0``), the verify arm falls back to the
ceiling at every batch, which is sound but reserves the batch factor the
arm's own occupancy formula rules out. ``H_q`` / ``H_kv`` are per-rank.
Returns 0 when no ``TkvBypassBackend`` is active (tkv / MLA paths) or
the model has no PAGED_KV layers. Pure arithmetic — no CUDA, no engine state.

See ``backends/tkv_bypass_backend.py`` (``_decode`` / ``_mtp_verify``) and
``tkv/kernels/_turbo_attn_simt.cu`` (``partial_o`` / ``partial_lse``).

Bytes the Turbo prefill padded gather stages per KV token of accumulated context.

The per-token slope of :func:`tkv_prefill_staging_peak_bytes` — the same
``index_select`` rows, K/V packed + norm repacks, dummy-KV scratch and
int64 slot temps, priced for ONE token so a caller that knows a row's
accumulated context can price that row. Admission uses it to charge a
prefill row's gather while it builds the slate
(:mod:`arbi_serve.scheduler.activation_budget`); the reserve above uses it
at the widest admissible context. Returns 0 when no ``TkvBackend`` is
active or the model has no PAGED_KV layers. Pure arithmetic.

Analytic peak bytes of the tkv Turbo chunked-prefill staging transients.

The tkv-SPECIFIC sibling of :func:`tkv_bypass_scratch_peak_bytes`, for the
PREFILL side: the B==1 eager Turbo prefill paged short-circuit
(``tkv/runtime/attend/turbo_prefill_paged_loader.py::turbo_prefill_from_paged_cache``)
materializes the whole accumulated context FRESH per (chunk, layer) —
``index_select`` of ``S_kv`` cache rows, K/V packed + norm repacks, an
``S_kv``-scaled dummy-KV scratch and int64 slot-index temps — none of it
in a budgeted pool, and none of it visible to the boot activation probe
(which runs at chunk-scale ``S_kv``).

WHAT THIS TERM DOES AND DOES NOT COVER. tkv takes the padded gather
whenever the step carries exactly one prefill row (``B == 1``, eager, not
route-pinned) and that row's accumulated KV is at or below
``PREFILL_PAGED_GATHER_MAX_TOKENS``; there is no ``S_kv <= total_q``
condition, so a continuation chunk of a long chat gathers its whole
accumulated context, not its query width. That context-scaled term is
charged where it can be BOUNDED rather than guessed — per row, at admission
(:mod:`arbi_serve.scheduler.activation_budget`, sharing this module's
:func:`tkv_prefill_gather_bytes_per_kv_token` slope), against the same step
budget the profiled activation floor reserves. What stays here is the
query-width residue every prefill step stages regardless of route, sized at
``Sg = min(max_context, max_batched_tokens)``:

  * ``rows`` gather: ``Sg x row_bytes`` (packed K/V + 16-aligned norm
    sections, the paged-cache per-token row).
  * K/V packed + norm repack copies: ``Sg x Sg x H_kv x 2``.
  * dummy K/V (bf16, K=V aliased): ``(Sg + 256) x H_kv x D x 2``
    (256 = max Turbo prefill tile_n).
  * int64 slot/index temps (tok / pages / slot_mapping / slot_idx):
    Sg x 8``.

Staging is reused across layers within a step (stream-ordered), so ONE
layer's set at ``Sg`` is the peak of this residue. Returns 0 when no
``TkvBackend`` is active or the model has no PAGED_KV layers. Pure
arithmetic — no CUDA, no engine state. Keep in lockstep with
``turbo_prefill_paged_loader.py``.

Predicted bytes for the pinned ``scratch.gdn`` workspace pool.

Only the monolithic capture path pins this pool: under ``split_attn``
(the production default) the GDN scan runs eager between the captured
per-layer graphs and its FLA chunk-fwd workspace routes through the
default allocator, covered by the profiled ``scratch.forward_arena``
reserve (never-pin contract — see ``cudagraph_admin.py`` /
``profile_peak.py``). Under ``split_attn=OFF`` with cudagraphs on, the
capture pin sites route ``chunk_gated_delta_rule``'s O(num_tokens)
intermediates into the pinned pool so they stay OUT of the captured
mempool — one pool's worth reused across buckets and layers. That pool
is engine-lifetime resident and must be held out of the KV budget.

Formula — fp32 upper bound over the chunk-fwd intermediates of ONE
call at the widest tokens ``T`` (``H_v`` = per-rank value heads, ``K``
/ ``V`` = key/value head dims, ``BT`` = :data:`_FLA_CHUNK_LEN`):

    per_token_elems = H_v * (2*K            # GVA-expanded q/k copies
                             + 1            # g cumsum
                             + K            # w  (WY repr)
                             + V            # u  (WY repr)
                             + BT           # A  (intra-chunk)
                             + V            # v_new
                             + V)           # o
                      + H_v * K * V / BT    # h-state blocks
    bytes = per_token_elems * 4 * T * _GDN_WORKSPACE_LIVE_SETS

taken as the MAX over GDN layers (the pool is reused per-layer, never
summed). Returns 0 under ``split_attn``, with cudagraphs disabled
(the pin sites never run), or with no GDN layers.

Per-tenant breakdown of the ``capture.io_buffers`` prediction.

One field per owner so the boot log can name which tenant moved when the
pool's measured row moves, instead of reporting a single number that only
says the total is wrong.

Closed-form byte cost of every ``capture.io_buffers`` tenant.

The pool holds the persistent kernel I/O a captured graph bakes by
``data_ptr`` — deliberately a different MemPool from ``capture.cudagraphs``
so address reuse cannot corrupt a replay. It is freeze-capped at its live
mapped size (it is not in ``FREEZE_CAP_BUDGETED_POOLS``), so every byte a
serving forward can need is resident before the post-capture KV grow and
comes straight out of KV: it needs a budget line, not a residual.

Terms (``T`` = ``max_num_tokens``, ``B`` = ``max_num_seqs``, ``H`` =
``hidden_size``, ``V`` = ``vocab_size``, ``P`` = ``block_table_pages``,
``e`` = the model dtype's element size):

``piecewise_batch``
    ``16*T + 17*B + 8*(B+1) + 4*B*P`` — ``PiecewiseBuffers``' int32 ids +
    positions, int64 slot_mapping, the per-row seq_lens / cu_seqlens pair,
    the block table and the three recurrent per-row buffers. Zero without
    cudagraphs (the ``cuda_graphs=False`` twin is built outside the pool).
``verify_batch``
    ``n_flat*16 + 8*(B+1) + 4*B + 4*B*Pv + 8*K*B`` at
    ``n_flat = B*(K+1)`` — the same shape family at the verify slate width.
``verify_sampler``
    ``(K+1)*B*V*4`` (``sc_p_target``) ``+ (K+1)*B*V*e_logits``
    (``sc_verify_logits``) ``+ (K+1)*B*8``, plus ``2*K*B*V*4`` on a boot
    that can carry a dense drafter ``q`` — the densified proposal the
    sampler reads (``sc_draft_probs``) and the drafter's own destination
    it is densified FROM (``sc_drafter_q``), which are separate members
    because the overlap path has them live at once — and
    ``2*(K+1)*B*V*4`` where the sampler materialises its Gumbel noise
    instead of drawing it in-kernel (Philox off).
``decode_capture``
    per captured ``(B_i, S_i)`` and per kv-page x LoRA bucket: the
    ``S_i > 1`` shapes take prefix SLICES of ``VerifyBuffers`` and allocate
    only ``B_i*S_i*H_verify*e`` of hidden output; the ``S_i == 1`` shapes
    allocate their own ``N*16 + 4*B_i + 8*(B_i+1) + 4*B_i*P`` inputs.
    Recurrent models add ``13*B_i`` of per-row state indices.
``shared_logits``
    ``max(B_i)*V*e`` — one buffer for the whole ladder.
``residual_bufs``
    ``sum(n*H*e)`` over the pre-warmed / captured rung set: one permanent
    slab per distinct ``num_tokens`` the cross-layer fusion runs at.
``dflash_tap``
    ``T*H*n_tap*e + 8*T + 4`` — the tap slab plus its token-count and
    flat-position validity buffers.
``dflash_capture`` / ``exl3_reconstruct``
    Resolved by the caller (:func:`dflash_draft_context_buffer_bytes` at
    the capture sweep width; the widest bound EXL3 linear's fp16
    reconstruct), because both read a checkpoint rather than the config.
``exl3_kernel_rows``
    ``sum(2 * W)`` over the DISTINCT EXL3 kernel input widths ``W`` bound
    on this member — the ``(1, W)`` fp16 row each ``BC_LinearEXL3``
    descriptor Hadamard-transforms a single input row into, which the
    bsz==1 decode graphs bake by ``data_ptr``. One row per width, not per
    linear: exllamav3 caches them on the shape. Checkpoint geometry, so the
    caller resolves the widths (see
    :func:`~arbi_serve.weight_quant.exl3.kernel_scratch.kernel_row_widths`).

NOT covered, and deliberately: the sub-MiB per-capture odds and ends (the
DFlash draft graph's ``noise``/``out`` pair, the drafter-chain and
prefill-replay staging tensors), and the allocator's segment rounding —
the measured row is cuMem MAPPED bytes, which rounds every allocation up
to a granule. Both land in the reconcile's delta, which is where rounding
belongs; a tenant does not.

Predicted bytes for ``capture.io_buffers`` — the sum of every tenant.

Takes the same keyword arguments as :func:`capture_io_buffer_terms`, which
documents the formula per tenant. Split so the boot log and the tests can
read the breakdown while the budget consumes one int.

Segment-level residency for the named pools — what is RELEASABLE, not just free.

The freeze ledger reports each pool as ``(live / held-free)``, where held-free is
``reserved − allocated``. That number is the right alarm but the wrong unit for
deciding what can be handed back to KV, because the caching allocator returns
segments that each still carry one live tensor can release NOTHING; a pool

So this module splits held-free into the only two buckets that matter:

  * ``fully_free_segment_bytes`` — segments with zero ``active_allocated``
    blocks. The upper bound on what any release mechanism could return.
  * ``partial_free_bytes`` — free blocks inside segments that are still pinned
    by a live allocation. Unreachable without defragmentation, i.e. not
    reclaimable at all on this allocator.

WHY ``partial_free_bytes`` CANNOT BE RELOCATED AWAY. Emptying a partially-free
segment means moving every live block out of it, and no mechanism available
here can do that:

  * The allocator picks the destination, not the caller. A replacement
    ``torch.empty`` is served best-fit from the pool's free blocks — i.e. from
    the very holes being vacated — so a copy inside the same pool moves the
    tenant into another segment's hole and leaves an equal hole behind.
  * A ``NamedMemPool`` destination re-creates the hole and makes it permanent.
    A private ``MemPool`` is this same caching allocator over a pluggable
    backing malloc, so it applies the same segment sizing and round-up: the
    relocated tenant costs the same reserved bytes there, and the round-up tail
    is then never returnable at all (pytorch#145168) instead of merely idle.
  * The permanent residents are pinned. A resident that survives to serving is
    either referenced by a captured cudagraph — its address is baked into the
    replayed kernel arguments — or held by a library in C++ with no Python
    referrer a ``Tensor.set_`` rebind can reach.

So ``partial_free_bytes`` is serving working reserve, not slack. It is held out
of the KV ceiling by
:meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.torch_caching_reserved_overhang_bytes`
for exactly that reason. The one defragmenting mechanism torch offers —
``expandable_segments`` — is mutually exclusive with ``torch.cuda.MemPool``
(pytorch#147851) and is therefore forced off for the whole process (see
:mod:`arbi_serve.__init__`), so there is no allocator here that could return it.

A pool holding CAPTURED CUDAGRAPHS admits no such split from torch's view at
all: capture drops the intermediates' Python refs, so their blocks read free
while the graph's launches still address them. That pool's row is recomputed
against its cuMem tag-mapped bytes (:func:`_correct_cumem_mapped`) — nothing in
it is reclaimable while a graph is alive, and the unit the one release path
works on is the whole pool, never a segment.

MEASURED CONSTRAINT (torch 2.12.1+cu130, sm_89): ``torch.cuda.empty_cache()``
does NOT visit the block pools of a private ``torch.cuda.MemPool``
(pytorch#145168), so even a wholly-free segment inside a named pool stays
resident. Both reclaim passes in the capture handoff call ``empty_cache`` and
per-pool held-free on a boot whose ``driver.residual`` is 0.00: the bytes are
accounted, they are simply not returnable while their pool object is alive.

The one mechanism that DOES return them is destroying the pool
(``NamedPoolRegistry.free_all``'s per-pool release → drop → ``empty_cache``),
which is only sound for a pool with no live allocations.

Read-only and best-effort: every helper here degrades to zeros rather than
raising, so a diagnostic can never wedge a boot or a serving step.

The ``(id0, id1)`` PRIVATE ``MemPool`` id a snapshot segment belongs to.

``None`` means the segment belongs to no private pool, i.e. it is a plain
default-caching-allocator segment — which is what makes it the one kind
``torch.cuda.empty_cache()`` can hand back (pytorch#145168 immobilises only
private ``MemPool`` block pools).

Both the ``(0, 0)`` sentinel and an absent key read as "no private pool".
Treating ``(0, 0)`` as a real id is not a cosmetic slip: it makes
:func:`default_pool_residency`'s filter match nothing, so its
``fully_free_segment_bytes`` reads 0 on every boot and the KV ceiling holds
out reclaimable bytes it was written to subtract; and it merges the default
pool into ``unpooled.unregistered_pool`` in :func:`residency_by_pool`, where its
held-free reads as a reclaim opportunity that ``empty_cache`` then returns
nothing for.

Residency split of the DEFAULT caching pool — segments in no NamedMemPool,
i.e. those carrying the ``(0, 0)`` :data:`_DEFAULT_POOL_ID` sentinel.

The default pool is the ONE place ``torch.cuda.empty_cache()`` can return
fully-free segments (pytorch#145168 immobilises only private ``MemPool``
block pools, not the default pool). So its ``fully_free_segment_bytes`` are
reclaimable and only its ``partial_free_bytes`` are the genuinely-pinned,
fragmentation-held residual that must stay resident. All-zeros when CUDA or
the snapshot is unavailable.

Sorted ``(start, end)`` address ranges of LIVE blocks in the default pool.

Block addresses are accumulated from each segment's base rather than read
off the block, so this works on snapshots that omit the per-block
``address`` field (mirrors ``NamedPoolRegistry._debug_dump_pool_leak_owner``).

Sorted ``(start, end, total_size)`` of every DEFAULT-pool segment.

The handle a reader has on an allocation no Python object names: which
segment holds it, and how large that segment is — i.e. how much reserved
memory this one tenant pins, which is the number that decides whether
moving it is worth anything.

What an allocation IS when no Python object names it.

dtype, shape, device and the segment holding it. That is recognisable on
sight and says what moving this tenant would buy; a placeholder says
neither. ``no_referrer`` distinguishes "nothing gc-visible holds this" — a
captured graph's baked buffer, an extension's own cache, an executing
frame's locals — from "held, but the walk could not name the holder",
because the two want different fixes.

Attribute the ``unpooled.torch_default_pool`` LIVE bytes to their tensors.

That row is the one resident term with no predictor and no owner — the
budget reconcile books it "UNBUDGETED". Naming its tensors is what makes it
readable as "real, in use, and this is what it is" rather than a blank.

Reads the allocator's segment snapshot and the Python heap, and NOTHING
else. In particular it does not need ``_record_memory_history``, which
:func:`arbi_serve.engine.build.build` arms only on a re-entered build (the
cold boot's compile sweep OOM-kills the recorder on host) — so a snapshot's
per-allocation frames are absent on exactly the boot that sizes KV, while
this attribution is not.

Rows are per STORAGE, not per tensor. Many live tensors routinely address
ONE allocation — a width-max capture buffer handed to each batch rung as a
``narrow(0, 0, B)`` view is the standard case — and
``untyped_storage().nbytes()`` reports the WHOLE allocation for every one of
them. Summing per tensor multiplies a single resident allocation by its view
count and prints bytes that are not on the card, under per-view SHAPES that
read as separate allocations. Each storage is therefore counted once, under
the shape of the widest tensor addressing it, and ``views`` says how many
live tensors do.

A row names EVERY storage it aggregates, with no sampling. Two different
tensors can share a shape and a dtype exactly — an untied ``(vocab,
hidden)`` input embedding and its ``lm_head`` are the same bytes in the
same dtype — so a row that named only a sample would answer the question it
was asked backwards. Attribution is
:func:`arbi_serve.engine.arena_tenants.owner_paths`, the same walk the
forward-arena reclaim refusal uses, which resolves a named module root, an
instance attribute or a closure-held cache; a tensor it cannot name is
described by what it IS instead — dtype, shape, device, and the segment it
pins.

Returns ``[{bytes, count, views, shape, dtype, owner}]`` sorted by bytes
desc, where ``count`` is distinct storages, ``views`` is live tensors over
them, and ``owner`` names every distinct holder the row aggregates. Empty
on any failure — a diagnostic must never wedge a boot.

Split the walk into its two halves: FIND the tensors, then NAME them.

They are separate problems with separate fixes — the find is one
``gc.get_objects`` pass over the heap, the name is a graph search — and a
single total cannot say which one to attack. Off unless ``ARBI_BOOT_PROFILE``.

Publish an already-walked attribution as this build's cached answer.

The boot walks the heap once, at the serving freeze, while nothing is being
served. Recording the result here is what keeps :func:`resolve_unpooled_owners`
a dict read on the request path: the walk holds the GIL for on the order of a
second, and the admin memory snapshot runs on the engine loop, so a walk
triggered by a request is a stall for every sequence in flight.

The unpooled-row owner attribution published for ``eng``'s current build.

A dict read, and ONLY a dict read. The invariant this function exists to
hold is that :func:`unpooled_live_owners` — a full ``gc.get_objects`` pass
plus a module-namespace BFS, measured at 1.62s on a warm 27B boot, holding
the GIL throughout — never runs on the engine loop, which is where the admin
memory snapshot executes. "Never" is checkable by reading this function;
"only on a cold-cache path" would stop being true the first time someone
added a caller.

So a miss returns ``()`` and says so at WARNING. The boot's freeze publishes
these via :func:`stash_unpooled_owners`; if it did not, that is a bug worth
reporting, and an admin panel missing an owners row is a diagnostic gap
while a 1.6s stall of every in-flight request is a production incident.

Keyed on the build generation, so a hot-swap reload (which re-enters
``build`` and bumps it) reports a miss rather than serving the retired
build's owners.

Log :func:`unpooled_live_owners` as a table; returns rows logged.

Always emits the table, the zero-row case included. A silent return is
indistinguishable from a diagnostic that never ran, which is the confusion
this dump exists to end — so an empty walk prints WHY it is empty, loudly.

The table states its own trust: taken after a re-entered build it can carry
a previous engine generation's allocations, and a reader must not have to
infer that from the boot log.

It also states its own COST. The walk names every storage rather than a
sample, which is what makes a row actionable, and the price is heap scans
proportional to how many tenants nothing names — a number only the boot it
ran on knows.

Book a CUDAGRAPH pool's tag-mapped bytes as live, not fully-free.

``torch.cuda.graph(pool=...)`` drops the capture-time tensors' Python refs,
so torch's block view calls a captured graph's memory ``inactive`` even
though the graph replays out of it and its device addresses are baked into
the launches. EVERY field derived from that view is wrong for such a pool —
the byte split and the segment counters alike — so the whole row is
recomputed against the pool's cuMem tag-mapped bytes: no field may survive
naming a wholly-free segment the corrected bytes deny.

Graph-holding pools ONLY — the rule
:func:`arbi_serve.engine.phase2_freeze.measure_pool_free_bytes` applies.
Tag-mapped bytes count a pool's mapped-but-FREE segments too, so flooring an
ordinary cuMem pool at them reports 0 held-free for it and erases the idle
overhang this table exists to surface.

``{pool name: residency dict}`` from ONE snapshot walk.

One walk for every pool (the per-pool accessors in ``NamedMemPool`` each take
their own snapshot, which is fine at metrics cadence but wasteful — and
inconsistent — when sampling all pools together).

Segments outside every registered pool split by the ledger's own rule:

  * no PRIVATE ``MemPool`` id (the ``(0, 0)`` sentinel) →
    ``unpooled.torch_default_pool``, the caching allocator's working reserve
    re-taken by every forward — resident and unbudgeted, but NOT a leak.
    Its fully-free segments are the only ones ``empty_cache`` returns.
  * a ``MemPool`` id this boot DESTROYED
    (:meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.release_empty_pool`,
    e.g. the forward-arena overhang reclaim) → ``address_space.released_pool_va``.
    Its cuMem physical went back to the driver through
    ``cuMemUnmap``/``cuMemRelease`` while torch kept counting the VA, so the
    row reads reserved AND fully-free and is neither: the physical is
    already KV. Named apart from the residual because it is a KNOWN
    category, not an unexplained one.
  * a ``MemPool`` id that matches no registered and no released pool →
    ``unpooled.unregistered_pool``: an unregistered pool, and the only one of the
    three that is a leak signal — the bytes have an owner and nobody
    declared it.

``unpooled.unregistered_pool`` here is the CACHING-ALLOCATOR quantity. The boot
ledger's ``driver.residual`` row (:mod:`arbi_serve.engine.vram_ledger`) is a
different number on a different surface: the whole-card identity's
remainder, read from ``mem_get_info``, which legitimately absorbs
driver-resident cubin growth and is NOT a leak signal.

Log the releasable-vs-pinned split of every pool's held-free bytes.

The line the freeze ledger cannot give: of the held-free bytes on each pool,
how many sit in wholly-unused segments (the ceiling on any reclaim) and how
many are stranded behind a live tensor (never reclaimable).

The table itself reads the allocator's segment snapshot. Naming the unpooled
row's tensors does not: :func:`unpooled_live_owners` walks the whole Python
heap and then BFSes the module namespaces, which costs seconds. Under
``attribute_unpooled=False`` the NOTE states the row's bytes and says where
the names are served from instead of paying that walk inline — for the boot
path, where the walk lands on the critical path of every start and the names
are read by an admin request that may never come.

Returns the :func:`unpooled_live_owners` rows it named — empty when the
attribution was deferred.

cuMem bytes torch still counts as reserved whose physical is released.

``tracked − mapped`` over the pluggable allocator — the same term
:func:`arbi_serve.runtime.cumem_allocator.real_reserved_bytes` subtracts to
keep ``memory_reserved`` honest. Read here so the residency table can say
which part of ``unpooled.unregistered_pool`` is a ghost rather than a reclaim
opportunity. 0 when no cuMem allocator exists.

Background sampler for the per-pool live/reserved HIGH-WATER MARK.

``NamedMemPool.peak_allocated_bytes`` only advances when ``snapshot()`` is
called — i.e. at the metrics tick — so a prefill-width transient that lives
for tens of milliseconds is invisible to it. Deciding whether a pool's
held-free headroom is ever TOUCHED at serving time needs a real peak, so this
samples on its own thread at a configured rate and keeps the max.

The document is ONE schema throughout: ``peak``, every ``marks`` entry and
``final`` all carry :func:`residency_by_pool` rows, so a peak row and the
final row are read with the same field names and subtract directly. A peak
row is a per-FIELD envelope over the run, not a copy of any one sample: each
field is its own maximum, so two fields may come from different moments and
the per-sample identity ``held_free = reserved - live`` does not survive the
reduction. Only extensive fields are reduced (:data:`_PEAK_FIELDS`).

Document fields:

  * ``samples`` — residency reads that returned a table, ``read_errors``
    those that raised (a read that raises is skipped, never folded in as a
    sample of zero). A write whose own read raised carries ``final: {}``
    rather than dropping the document: the peak is what a run leaves behind.
  * ``interval_s`` — the REQUESTED sleep between samples, not the achieved
    period: each read walks the allocator snapshot, so the true period is
    this plus the walk. ``elapsed_s`` and ``observed_interval_s`` report what
    actually happened, so the sample rate is a measurement rather than a
    restatement of the request.
  * ``peak`` — the per-field maximum over every read, including the reads
    taken for ``marks`` and ``final``.
  * ``marks`` / ``final`` — instantaneous tables, at a labelled checkpoint
    and at write time respectively.

Strictly opt-in (``ARBI_DEBUG_POOL_HWM``), off by every default path, and a
daemon thread so it can never hold up shutdown. The JSON is rewritten
atomically so a reader always sees a complete document.

Start the per-pool HWM sampler when ``ARBI_DEBUG_POOL_HWM`` is set.

Writes ``/tmp/pool_hwm_dev<N>.json``, mirroring how
``ARBI_DEBUG_FREEZE_MEMDUMP`` places its allocator snapshot — one fixed,
device-scoped path per diagnostic, so there is no path to mistype and no
second flag to keep in sync.

RESIDENT bytes the named pool holds right now, or ``None`` when unreadable.

Read with the SAME meter the boot VRAM ledger's row for that pool uses
(:func:`arbi_serve.engine.phase2_freeze._all_pool_mapped_bytes`): the cuMem
allocator's per-tag mapped bytes when the driver path is live, and the
pool's torch-reserved bytes when it is not. A plan line and the ledger row
it is compared against have to come from one instrument, or the difference
between them is the instruments rather than the prediction.

``None`` — not 0 — when there is no such pool or no reading can be taken,
so a caller can tell "this pool holds nothing" from "nothing measured it".

Rewrite the json atomically, folding this table into the peak first.

Every table this class reads — sampled, marked, or read for ``final`` —
goes through :meth:`_observe`, so a written document always carries a
peak covering every read taken, from the first write onward.

RoPE / LoRA pool predictors and the serving-reserve byte estimators.

The RoPE-cache and LoRA pool predictors, the per-rank NCCL workspace
estimate, the cuMem-invisible serving reserve, the torch caching-overhang
reserve, and the single activation-reserve line consumed by
:func:`arbi_serve.engine.memory_budget.compute_kv_budget`.

Predicted bytes for :data:`rope_cache_pool`.

Formula (per unique rope cache key):

    bytes = 2 * rope_seq_len * int(head_dim * rope_partial_factor) * elem

where ``rope_seq_len`` is ``model_dims.effective_rope_cache_seq_len()``
— ``min(rope_cache_seq_len, max_position_embeddings)`` when the
engine has stamped the serving gate, else the full
``max_position_embeddings``. Sums over the unique ``(head_dim,
rope_partial_factor)`` keys across the model's layers.

Only the leading ``int(head_dim * rope_partial_factor)`` channels
rotate, so the cos/sin tables are that wide — a full-``head_dim``
table over-predicts partial-rope models (Qwen 3.5 / 3.6:


See ``models/layers.py``.

Predicted bytes for ``scratch.persistent_fold``.

The slab that backs the folded scratch sub-tags is ONE allocation of their
summed budgets, and the caching allocator serves it from a segment rounded
up by its own sizing rule. That rounding is all this pool holds of its own —
each sub-tag's bytes stay predicted on its own row — so::

    bytes = segment_bytes_for(sum(subtag_bytes)) - sum(subtag_bytes)

Zero when nothing is folded (fewer than two non-empty sub-tags), which is
also when the fold is skipped.

See :mod:`arbi_serve.engine.persistent_fold`.

Predicted bytes for ``scratch.penalty_accum``.

Formula::

    bytes = max_batch * vocab * (elem + 4) + vocab * (elem + 4)

Two ``(max_batch, vocab)`` buffers plus one scratch ROW of each:

* the accumulator itself, in the LOGITS dtype (``elem``);
* the int32 occurrence counts (4 B) that ``frequency_penalty`` reads
  to recompute ``freq * count`` exactly and ``repetition_penalty``
  reads as its seen set;
* one spare row of each, used to rotate rows when the slate reorders.

Reserved WHOLE at boot, including the counts half. The alternative —
allocating counts lazily on the first request that needs them — would
take those bytes out from under a KV pool that was sized as if they
were free, at ``gpu_memory_utilization`` 0.99. A reserve that is
honest about the worst case is cheaper than an OOM under traffic.

Scales with ``max_batch x vocab`` and NOT with context: the bound is
the widest batch the SCHEDULER can admit, not the widest CAPTURED
shape (eager decode past the capture ladder still needs a row).

Predicted bytes for ``scratch.logprobs``.

Formula::

    bytes = tile * vocab * (elem + 4 + 1)

One ``(tile, vocab)`` lm_head output in the LOGITS dtype (``elem``),
the ``(tile, vocab)`` fp32 log-softmax it is reduced through, and the
``(tile, vocab)`` bool mask the rank comparison produces.

``logprobs_cfg is None`` — the surface was not armed — returns 0 and
the pool stays empty.

Scales with ``tile x vocab`` and NOT with prompt length: the prompt-
logprobs path walks a prompt one tile at a time and reduces each tile
to the requested top-k before the next, so the bound is the tile.

Predicted bytes for :data:`lora_pool`.

``lora_cfg`` is a duck-typed object with attributes:

  - ``max_loras`` (int) — 0 disables LoRA → return 0
  - ``max_rank`` (int)
  - ``max_batched_tokens`` (int)
  - ``target_shapes`` (dict[str, tuple[int, int]]) — module → (in, out)

Or ``None`` → return 0.

Formula (slots = ``max_loras + 1``; ``T = max_batched_tokens``):

    bytes = 4 * T                                 # lora_token_indices i32
          + 4 * slots                             # scalings f32
          + 2 * elem * slots * R * Σ_target (in_f + out_f)  # packed_a + packed_b

Q36 (``lora_max_loras=0``): returns 0.

See ``lora/capture.py:111-122``.

Resolve the per-rank NCCL workspace bytes, honouring the env override.

Returns :data:`_NCCL_WORKSPACE_PER_RANK_BYTES` unless
``ARBI_NCCL_WORKSPACE_BYTES_PER_RANK`` is set to a non-negative integer,
in which case that value is returned and a one-line LOUD log is emitted
the first time an override (or an unparseable value) is seen. A bad value
falls back to the measured default — LOUD, never silent.

Predicted per-rank bytes for the (out-of-torch) NCCL communicator
workspace, held out of the KV ceiling so the pool does not claim VRAM the
libnccl-internal channel buffers will occupy.

  - tp_size == 1: 0 (no communicator)
  - tp_size > 1:  the per-rank workspace reserve — the measured 2-GPU
    no-P2P default (``_NCCL_WORKSPACE_PER_RANK_BYTES``) unless the
    ``ARBI_NCCL_WORKSPACE_BYTES_PER_RANK`` env override is set (see
    :func:`_resolve_nccl_workspace_per_rank_bytes`)

``cudagraphs_enabled`` does not change the figure: the cudagraph-capture
path does NOT mint extra per-channel symmetric buffers on this topology.
The argument is kept for call-site compatibility.

NCCL's libnccl-internal allocations are tracked by the driver but are
invisible to torch's caching allocator AND to the cuMem accounting — we
still subtract them from ``usable_bytes`` so torch's KV pool doesn't claim
driver memory the NCCL workspace will need.

Bytes resident in the serving phase that are INVISIBLE to torch's
caching allocator AND to ``CuMemPoolAllocator.mapped_bytes``, scaled by
the TP degree — held OUT of the serving-phase KV grow ceiling so the pool
does not plan to grow into VRAM the driver will keep occupied.

At TP1 the only invisible residual is the single CUDA primary context,
already folded into ``external_reserved`` and inside the pre-capture
budget; this returns 0 there so the legacy (correct) TP1 sizing is
unchanged.

At TP>1 two driver/library-level residuals appear that the pre-capture
budget under-holds because the post-capture grow assumes the freed
prefill arena is fully reclaimable to KV — it is not, the driver
re-occupies it:

  * The sibling rank's CUDA primary context + cuBLAS/cuMem driver
    scratch. ``external_reserved`` is profiled as a single per-rank
    ``own_used − own_reserved`` snapshot (``profile.py``) taken before the
    sibling rank's context is fully resident, so it captures roughly ONE
    context — and only OURS: a co-tenant's VRAM is booked separately as
    ``driver.foreign_process``, so it is never re-scaled by (tp−1) here. The serving footprint is ``tp_size`` contexts —
    the extra ``(tp_size − 1)`` contexts are the under-held bytes,
    approximated by scaling the measured ``external_reserved`` by that
    factor.

  * The NCCL communicator workspace (``nccl_workspace_bytes``), one per
    rank, libnccl-internal and never cuMem-tracked.

The ``GroupCoordinator`` reduce/gather ``_staging`` buffers are NOT
counted here — they are routed through the cuMem-backed
``nccl_staging_pool`` NamedMemPool (``parallel_state.attach_staging_pool``
in engine build). As a named cuMem tag their mapped bytes are VISIBLE to
``CuMemPoolAllocator.mapped_bytes`` and already inside the post-capture
pool accounting / KV ceiling — counting them here too would double-count
the same VRAM, over-holding the grow ceiling.

Pure arithmetic — no CUDA, no engine state.

Serving-resident torch-caching scratch overhang held out of the ceiling.

``measured_bytes`` — the value the Phase-2 freeze measured at THIS
configuration on a previous boot
(``NamedPoolRegistry.torch_caching_reserved_overhang_bytes``, persisted by
:func:`arbi_serve.engine.memory_budget.graph_pool.
measure_and_persist_serving_caching_overhang`). When present it is used
verbatim: a direct measurement of a physical quantity always beats a
multiple of a different physical quantity.

``None`` (cold cache) falls back to ``profiled_activation_peak_bytes x
_SERVING_TORCH_CACHING_OVERHANG_FACTOR`` — see that constant for why the
factor is a seed rather than a model, and why it is 1.0.

Pure arithmetic — no CUDA, no engine state. Returns 0 when neither input is
available (the cold-boot path supplies the closed-form activation reserve
instead).

The overhang the POST-CAPTURE KV sizing is entitled to size against.

``planned_bytes`` — what the pre-capture plan held out for the activation
footprint's overhang form (``boot_state.
serving_torch_caching_overhang_bytes``): a previous boot's persisted
reading, or the cold seed. ``measured_bytes`` — what this boot's settled
card actually holds
(:meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.
torch_caching_reserved_overhang_bytes`).

MONOTONE, ``max`` of the two, and the same discipline
:func:`~arbi_serve.engine.memory_budget.graph_pool.
measure_and_persist_serving_caching_overhang` keeps ACROSS boots, applied
WITHIN one. The quantity is a BOUND, not a point estimate, and its two
error directions do not cost the same: a shortfall costs the whole boot,
an over-hold costs some KV. It is also not a modelling error to fit —
overhang and KV are SUBSTITUTES (see
:data:`_SERVING_TORCH_CACHING_OVERHANG_FACTOR`), so a reading that comes in
UNDER the plan is the allocator having taken less because KV took more,
which is not evidence the plan over-held.

Pure arithmetic — no CUDA, no engine state.

Bytes the serving floor reserves for a speculative step's TRANSIENTS.

Two per-step transients sit above the base floor on a DFlash spec-decode
boot: the stochastic verify tail — the rejection sampler's vocab-scale
``(K+1, B, V)`` fp32 draws, allocated EAGER and fresh on every temp>0 step,
outside the captured pool — and the DFlash draft transient
(:func:`~arbi_serve.engine.memory_budget.pool_predictors.dflash_draft_transient_peak_bytes`).
Neither appears in the serving pool table, so the post-capture KV grow will
claim their VRAM unless the floor holds it back.

WHICH FOLD APPLIES IS A CONCURRENCY QUESTION, and ``device_slots`` answers
it. On the synchronous path the draft runs AFTER the verify's host sync, so
the two peaks are strictly sequential and the step's high-water mark is the
LARGER of them. Under ``ARBI_DFLASH_DEVICE_SLOTS`` the draft is eligible for
the async verify path and overlaps the verify tail instead of waiting on it,
so both can be live at once and the floor must hold their SUM.

The device-slots premium is therefore exactly That is the KV the flag costs.

``verify_logits_bytes`` is the THIRD term and it is not a third moment: the
``lm_head`` call it prices produces the very block the rejection sampler
reduces, so the step holds it across the sampler's own peak and the two are
charged together as one verify moment against the draft
(:func:`~arbi_serve.engine.memory_budget.pool_predictors.
verify_logits_epilogue_peak_bytes`).

Collapses to the plain verify tail whenever DFlash is off
(``draft_transient_bytes == 0``) — every non-DFlash model — so both folds
are byte-identical there and this cannot perturb them.

Predicted bytes for the engine activation reserve.

Whether the runtime :class:`ActivationArena` standing buffer is
allocated or not, the budget always reserves the same number of
bytes for activation memory:

  * When a profile is available (``profiled_peak_bytes > 0``):
    ``int(profiled_peak_bytes * safety_factor)``. The serving path
    passes ``safety_factor=1.0`` — the raw profiled peak, no pad.
  * When no profile is available (cold-boot fail-fast path):
    a closed-form upper bound
    ``8 * max_num_tokens * hidden_size * elem`` (8× covers
    attention QKV temporaries + MLP up/gate/down + a 2× margin
    for ping-pong buffers).

The arena does NOT shrink this line. It is a standing buffer sized
at ONE LAYER's attention output — it rewinds at every layer
boundary — while this reserve covers a whole step's transients,
which stay on the heap, arena or no arena. Reserving the same bytes
either way keeps a single activation budget line and never
under-reserves the arena-on case.

See :mod:`arbi_serve.runtime.activation_arena`.

The serving-step reserve as ``((row, bytes, note, provenance), ...)``.

Same arithmetic as :func:`~arbi_serve.engine.inprocess_capture.
serving_floor_for_grow`, expressed as rows: the sum over ``bytes`` is
BYTE-IDENTICAL to the floor that function returns, which is what lets a
reader see what holds the free VRAM instead of one number to take on faith.

Two shapes of term, and they compose differently:

  * ADDITIVE — the base activation peak, the forward arena's serving
    re-growth, and the perception / speech / logprobs reserves. Each is
    charged in full because a step can hold them at once.
  * SINGLE-MOMENT — the speculative step reserve, the tkv-bypass decode
    scratch and the tkv prefill staging peak at DIFFERENT moments of a
    step, so only the LARGEST is charged. The other two get a row of 0
    carrying what they would have cost, because a term that vanishes is
    indistinguishable from a term nobody priced.

The speculative reserve splits into its own two rows exactly as
:func:`spec_step_reserve_bytes` folds them: both charged under device-slots
(the draft overlaps the verify sampler), otherwise the larger only.

``base_arena_excluded_bytes`` is the arena share the caller took OUT of the
base term because ``transient.serving_step.arena_regrow`` already holds it.
The two rows are what a reader compares, so the base row has to say that its
number is a peak minus a transfer rather than a smaller measurement.

``base_note`` lets the caller say the base term is not the per-shape max at
all — a boot with no profile prices it from a probe of a DIFFERENT, wider
step (:func:`~arbi_serve.engine.inprocess_capture.
unprofiled_serving_step_bytes`). That row is still MEASURED, so provenance
alone cannot carry the difference and the note has to.

``arena_regrow_provenance`` / ``arena_regrow_note`` come from
:func:`~arbi_serve.engine.inprocess_capture.forward_arena_regrow_plan`,
which picks that row's bytes from three candidates of different standing —
a bound admission enforces, a serving reading carried across boots, and
this boot's own boot-workload mark. The row cannot be labelled from the
number alone, so the caller says which one won and against what.

``gmu_floor_bytes`` / ``residency_gate_bytes`` are the two floors that can
bind ABOVE the step's own needs (the operator's ``gpu_memory_utilization``
ceiling; the stable-VA residency registration gate). Only the excess over
the step reserve is charged, and to an ``unclaimed.*`` row — those bytes are
held free for a policy, not for a step. Because they are MAX-ed rather than
added, either can be the whole floor while every measured step term sits
invisibly underneath it, so both rows carry the arithmetic that put them
there: ``residency_gate_note`` says what the gate's number is a reading OF
and what has already been paid against it, and ``residency_gate_provenance``
says whether that reading is a measurement. A row that prints only its
bytes hides exactly the question an operator is asking when the floor is
larger than everything under it.

``driver_growth_bytes`` is the ADDITIVE row for the far side of the
allocator boundary: physical the DRIVER takes after the KV layout is
frozen, which no allocator counter can see and which the card reports but
never reserved (see :func:`~arbi_serve.engine.memory_budget.graph_pool.
predict_serving_driver_growth_bytes`). ``driver_growth_observed=False``
means this configuration has never been watched serving, so the row holds
0 and carries ``UNBUDGETED`` rather than a measurement it does not have —
the bytes are known to be non-zero and nothing on this boot predicts them.

``verify_logits_bytes`` rides WITH the verify tail rather than against it:
the ``lm_head`` call it prices produces the block the sampler reduces, so
the step holds it across the sampler's own peak. The two therefore win or
lose the single-moment comparison together, and the pair is what the draft
transient is compared against.

Each row carries its own :data:`~arbi_serve.runtime.pool_taxonomy.PROVENANCE`
so a reader can see which of these numbers were observed and which were
computed. ``verify_tail_measured=False`` says the verify tail fell back to
its config-dims analytic bound, which downgrades that row from ``MEASURED``
to ``DERIVED`` — the row must not keep claiming a measurement that did not
happen.

Turn a memory FACT into a memory ACTION.

An OOM is a statement about the card, not about the code that was running when
it landed: "this allocation did not fit" -- UNLESS an allocator cap refused it,
which is a statement about a pool and reaches this module wearing the card's
words. WHICH ONE IT IS COMES FIRST, because the two facts have disjoint
reactions and the message alone cannot tell them apart (see
:func:`_classify_refusal`). The serving path had no way to say either. Every
path that met one either charged it to a defect counter whose threshold ends
the process, or finished the slate with an error and moved on --
so an OOM produced a restart or a silent stream of 500s, and in neither case
anything that made the NEXT step more likely to fit.

This module is the missing half. One OOM, wherever it is caught, reaches
:func:`note_oom`, which runs the whole reaction:

``count`` -- separately and visibly, in its own window. A cold-pathed drafter
draw and a broken drafter are different facts and cannot share a counter: the
first is the correct response to memory pressure, the second is a defect whose
persistence should end the process.

``narrow`` -- the step budget admission enforces is RE-MEASURED from the card
and every scheduler is re-armed with it
(:func:`~arbi_serve.engine.activation_admission.rearm_step_budget`). The
scheduler then narrows the prefill chunk and defers rows on its own, which is
the reaction that actually makes the next step fit. The budget is not adjusted
by a margin -- it is the same measurement the boot took, taken again at the
moment the card contradicted it, so nothing here is a cushion and nothing needs
a threshold flag.

``back-pressure`` -- when the re-measured budget can no longer cover the
model's OWN minimum step (its constant plus one row, no prefill tokens), there
is nothing left for admission to narrow: every slate it can build already
exceeds what the card has. The front door then refuses fresh work
(``memory_pressure``, HTTP 429) so the server stops admitting into a card that
cannot serve, while the rows already in flight keep running -- they are the only
thing that can give memory back. That refusal is armed only WHILE there is
something to drain: the latch is cleared by a step succeeding and steps come
from the rows in the scheduler, so a closed door on an empty scheduler would be
a deadlock rather than back-pressure. On an idle engine the door opens, the
latch stays set, and the next OOM is the escalating one -- which is right,
because the card drained completely and still could not run a row.

A CAP DENIAL IS NOT PRESSURE AND TAKES NONE OF THOSE RUNGS. When an ENFORCED
per-tag cap refuses a map, the allocator shim returns ``nullptr`` BEFORE the
driver is asked (``runtime/_cumem_shim.py``), and torch formats that refusal
with the same text -- and the same card-wide "N GiB is free" figure -- as a
genuine driver OOM. Every rung above is a lever on the CARD: re-measure what is
free, build a smaller slate, stop admitting until rows drain. None of them acts
on a cap, which bounds a pool's MAPPED bytes and is unmoved by the size of the
step; a narrower slate still needs one new block in a size class the pool has
no room to map. So the ladder would have run, reported a remedy that cannot
apply, and left the actual fact -- WHICH pool refused, and by how much --
sitting unread in the allocator, which is exactly what it did. The refusal is
therefore CLASSIFIED first
(:func:`~arbi_serve.engine.pool_caps.drain_cap_denials`), and a cap denial is
counted, named and reported instead of narrowed. The classification is
determinate in both directions: an armed gate that refused nothing is a null
control that says the driver refused, and a disarmed gate could not have
refused at all.

``escalate`` -- and if an OOM lands while that door is ALREADY shut and no step
has succeeded since it shut, the reaction has been applied and has not worked.
That is the line: not a count of failures and not a timer, but the state in
which narrowing has provably run out. It latches the sticky memory-exhausted
fault (:func:`~arbi_serve.engine.infra_health.note_memory_exhausted`), so
``/health/ready`` is 503 and the loop tears down for a restart. A card that is
genuinely too small fails loud instead of serving degraded forever.

WHY THE RE-MEASURE IS SINGLE-RANK ONLY. It reads ``mem_get_info`` through the
boot's own budget reader, which collapses the reading to the all-rank minimum --
a collective. An OOM at TP>1 is asymmetric by nature (one rank's allocation did
not fit), so a collective issued on that edge is issued by one rank and waited
on by none: the NCCL wedge this repo refuses. At multi-rank an OOM is already a
loud fault by the engine loop's own contract (the ranks are desynced and cannot
be reconciled in-band), so the narrowing ladder has nothing to add there and
this module only counts and reports.

WHY THE BUDGET DOES NOT SPRING BACK ON A TIMER. It follows the card's latest
reading and nothing else. A reading is taken on the OOM edge and again on the
first successful step after one -- the recovery edge, where the allocator is in
an ordinary state and the reading is the more trustworthy of the two. Between
those edges nothing is measured and nothing moves, so no step pays for this and
no decay constant has to be chosen.

Account one OOM and apply the reaction. ``True`` when it ESCALATED.

``source`` names where the allocation failed ("step", "drafter") for the
logs and the degraded registry; it does not change the ladder, because the
fact does not depend on which allocation met the wall.

A ``True`` return means the sticky memory-exhausted fault is already
latched: the caller's only remaining job is to propagate loudly so the run
loop tears down. ``False`` means the step's own path continues -- the slate
was cold-pathed or error-finished by the caller, and the next slate is built
against a smaller budget.

Never raises. An accounting or re-measure failure must not replace an OOM
with a second exception in the same handler.

A step succeeded: clear the pressure state and re-read the card.

The engine's post-step success seam calls this on every served step, so the
common path must cost one attribute read and return -- which it does:
``_oom_since_success`` is zero on a process that has not OOM'd, and zero
again the moment a step succeeds after one.

On the recovery edge (the FIRST success after an OOM) it re-measures the
step budget once. That reading is the more trustworthy of the two this
module takes: the OOM-edge reading is taken with the allocator holding the
fragments of a failed allocation, while this one is taken on a card that
just ran a step. So the budget follows the card back up when the pressure
genuinely passed, without any decay constant deciding when that is.

Whether the narrowing line should be written for this OOM.

Always for the FIRST of a burst — the one that says pressure started — then
at most once per
:data:`~arbi_serve.engine.infra_health.INFRA_LOG_INTERVAL_S` while it
persists. The two lines that matter operationally (admission closing, the
escalation) are outside this gate and always written: they are edges, not a
standing condition.

Ask the armed caps whether one of THEM refused this allocation.

``None`` only when the question could not be put — the classification is
then UNKNOWN and no caller may name a cause on its behalf. Every other
answer is determinate in both directions; see
:class:`~arbi_serve.engine.pool_caps.CapDenialDrain`.

Never raises: this runs inside an OOM handler, where a second exception
replaces the fact being diagnosed with the diagnosis's own failure.

The classification, as a clause for the line that reports the reaction.

A reaction line that does not carry it is the defect this exists to
remove: the card levers below are correct ONLY for a driver refusal, and a
reader has to be able to see that the alternative was ruled out rather than
never asked. When it could not be asked, it says so — an unknown is not a
licence to name the cause the reaction assumes.

Put the pressure on the ``serving_degraded`` registry.

That registry is what the ``serving_degraded`` gauge and the
``x-arbi-degraded`` response header read, so a server that is cold-pathing
its speculation or narrowing its chunks under memory pressure SAYS so on
every response instead of quietly serving slow -- which is the degradation
the drafter breaker was built to prevent and the reason simply stopping the
breaker from tripping would not be a fix.

A cap denial gets its OWN sentence here for the same reason it skips the
ladder: "the step budget follows the card and admission narrows to it" is
a claim about a reaction that, on that fact, does not run.

Copy per-layer centroids + per-channel scales onto the MLA ops.

Only the K side is read: MLA folds K and V into one quantiser and
:meth:`tkv.codec.TkvCodec.compress_shared_kv` drives it from the
K-side tensors.

K from disk.

Returns:
    - ``None`` on miss (no cache file, parse error, key mismatch, or
      explicit-override gate). The caller falls through to
      ``cfg.mtp.n_draft`` / head defaults.
    - The calibrated ``best_k`` (``int``, possibly 0) on hit.
      ``best_k == 0`` is meaningful — the workload's optimum was
      MTP-off — and the caller must honor it instead of falling
      through.

Whether THIS boot should build the checkpoint's bundled MTP head.

``False`` when spec-decode is off, OR when an explicit drafter source
(DFlash block-diffusion head / external draft-model) overrides it. The
bundled head is a full MoE decoder block plus a full attention-layer KV
slab per token; when :func:`build_mtp_driver` routes drafting to DFlash or
an external model, that head is loaded and then never touched. Threading
this into ``from_safetensors(mtp_enabled=…)`` means the head — its weights
AND its KV slab — is never built, so the freed VRAM sizes into the KV pool
(see :func:`arbi_serve.models._qwen3_5_loading.maybe_append_mtp_spec`).

Kept in lockstep with the source dispatch in :func:`build_mtp_driver`: any
source that wins over the bundled head there must return ``False`` here.
The assistant (Gemma-4 EAGLE3) path builds its own head separately and is
left untouched.

Resolve the speculation load valve's mode and announce it.

The valve is tri-state — ``off`` (always speculate), ``on`` (stop at a
pinned decode width), ``auto`` (decide from live measurement). The
verifier's FFN topology is only a FALLBACK: it seeds ``auto``'s
starting arm before anything has been measured, and supplies the width
a bare ``on`` pins. It is never presented as a measurement.

Logged at boot so the resolved policy is never silent, and at WARNING
when the valve is pinned to stop speculating at a width the operator
did not measure.

The single gate point for a boot that ends with a drafter attached.

INVARIANT: when control leaves :func:`build_mtp_driver` or
:func:`attach_mtp_driver`, the engine's drafter slot is either empty —
the plain K=1 decode path, nothing to refuse — or holds a drafter that
has passed every refusal named here. Both functions reach this as their
single exit, so a refusal added here is reachable from every drafter
source by construction and none of them carries its own copy.

The refusals belong to the ATTACH, not to any one route, because what
they read is not route-local: the verify side takes its speculation
geometry from process-global flags whatever the drafter is, and the
async/sync verdict is a whole-engine property. A route that builds its
own verify scaffolding rather than going through
:func:`attach_mtp_driver` (DFlash, external draft model) is subject to
both for exactly the same reasons as one that does not.

Scope, deliberately: an empty slot is a no-op. A boot that attaches no
drafter serves the K=1 decode path, which none of these refusals
describes — they each name a condition for serving SPECULATION
correctly. Detach (``attach_mtp_driver(eng, None)``) is the same case.

Every refusal here must stay a pure predicate — read the engine,
raise or return — because a source that attaches through
:func:`attach_mtp_driver` passes this on that function's exit and
again on :func:`build_mtp_driver`'s, and both must reach the same
verdict.

Build the MTP driver from the active model's bundled draft head.

Two modes:

  * **External draft model** (``cfg.mtp.draft_model_path`` set) —
    load a separate (smaller) draft model into its own pool and
    wire :class:`arbi_serve.spec_decode.external_drafter.ExternalModelDrafter`.
    See :func:`arbi_serve.engine.build_external_drafter.build_external_drafter`
    for the 8 boot asserts and the per-StateKind metadata + attn-op
    wiring.

  * **Bundled MTP head** (default) — pick the head off
    ``eng.model.mtp_head`` (constructed inside the model's
    ``__init__`` when the safetensors carry ``mtp.*`` keys) and
    wrap it in :class:`MtpDriver`.

Single exit: the source dispatch lives in
:func:`_route_drafter_source` and :func:`_seal_drafter_attach` is the
last statement here, so every source — including the ones that attach
their own drafter without going through :func:`attach_mtp_driver` —
passes the boot refusals on the way out.

Build and attach THIS engine's speculation source, or none.

One branch per source, each ending with a drafter attached (or with
``eng.drafter`` left ``None`` for the K=1 decode path). Carries no boot
refusal of its own: :func:`build_mtp_driver` seals every branch through
:func:`_seal_drafter_attach`.

Build + load the Gemma-4 EAGLE3 assistant drafter head.

Constructs :class:`Gemma4MtpHead` bound to the verifier's scaled
embedding, with its draft marker layer indices matching the engine's
appended marker specs (so ``pool.layer_view(idx)`` aliases the verifier
donor layers and ``attn_ops[idx]`` are the skip-write ops), and loads
the head's own weights from the separate assistant checkpoint via
``Gemma4MtpHead.weight_map``.

Resolve — once per boot — the draft depth this engine runs.

"Draft depth" is how many speculative tokens the driver proposes per
step. It is not ``head.n_draft`` (the blocks the checkpoint ships):
:meth:`MtpDriver.draft` chains a single trained block autoregressively,
so the depth run may exceed the block count. See
:mod:`arbi_serve.engine.mtp_depth` for the vocabulary distinguishing
the two.

Resolution order (first hit wins):

  1. already resolved this boot → the latched value (idempotent, so
     the KV sizer and the driver build can never disagree about K);
  2. MTP off → ``0``;
  3. external / DFlash drafter → the explicit ``cfg.mtp.n_draft``
     (the CLI already refuses those without one — a plain decode
     model / block-diffusion head carries no block count to chain);
  4. a cached K calibration → its ``best_k`` (``0`` means the measured
     optimum was MTP-off and must be honored);
  5. an explicit ``--mtp-n-draft K`` → ``K``;
  6. ``auto`` + cache miss → the chained sweep ceiling, priced against
     VRAM when ``budget`` is supplied, and the boot-sweep latch set.

``budget`` is the capture-ladder pricing input; it is only available
from ``profile_and_size_kv_pool`` (which is where the post-weights
VRAM picture exists). Called without it — the CPU/no-GPU path and the
direct-``build_mtp_driver`` path — the ceiling is taken unpriced, and
the authoritative downstream gates (``compute_kv_budget``,
``assert_capture_sweep_fits``) still refuse an over-budget config.

Pins ``cfg.mtp.n_draft`` to the verdict so every K-derived boot
quantity (KV sizing, the verify ``(B, K+1)`` and drafter ``(B, K)``
capture ladders, the driver's ``max_k``) reads the same number.

Install an :class:`MtpDriver` for opt-in speculative decoding.

The slot is hot-swappable: passing ``None`` detaches. Once
attached, requests with ``SamplingParams.mtp_k > 0`` route through
the driver; ``mtp_k == 0`` requests still use the K=1 path in the
same step.

The driver's head must reuse this engine's main-model
``embed_tokens`` and ``lm_head``; the bundled-head identity check
runs once at construction (when an engine is known) and once again
here for the test path that constructs the driver without an engine
first.

Single exit: :func:`_seal_drafter_attach` is the last statement, so a
hot-swapped driver passes the same boot refusals a booted one does.

Bind a drafter to the engine and build its verify scaffolding.

Carries no boot refusal of its own — :func:`attach_mtp_driver` seals
the attach. The ``RuntimeError`` here is not an exit that needs
sealing: it aborts before anything is bound, so no drafter is attached.

Run the auto-pick K-sweep against the live, built engine and pin the
winner. Returns the chosen ``best_k`` (or ``None`` if no sweep ran).

Called from the rank-0 serve path after ``run_forever`` is active —
the only point where the engine can serve requests, and (under TP>1)
where a rank-0 ``asubmit`` drives all ranks via the per-step plan
broadcast. The driver was already built at the sweep ceiling, so every
swept K hits a captured graph. Pins ``cfg.mtp.n_draft`` to the winner
(the per-request default) and writes the cache so a 2nd boot loads it
without re-sweeping. ``best_k=0`` detaches the driver (MTP off).

Bundled-MTP-head draft-depth vocabulary + the VRAM price of depth.

Two distinct quantities are easily conflated under the name ``n_draft``:

  * **blocks shipped** — :func:`blocks_shipped` (``head.n_draft``): how
    many trained MTP decoder blocks the checkpoint carries. Qwen 3.5 /
    3.6 ship exactly one (``mtp_num_hidden_layers = 1``); DeepSeek-V3
    ships two. It is a property of the weights. Config never raises it.
  * **draft depth run** — :func:`chained_depth_ceiling` /
    ``cfg.mtp.n_draft``: how many speculative tokens the driver actually
    proposes per step. For a depth-1 head this is the length of the
    autoregressive chain — :meth:`MtpDriver.draft` feeds the block its
    own ``(token, hidden)`` output back in, so the depth run may exceed
    the block count. It is a property of the deployment.

Depth is what pays: the spec-decode speedup is ``1 + a + a² + … + a^K``
(``a`` = per-token accept rate), so a K=1 head has to clear break-even on
a single extra token and usually cannot. Chaining is the only way a
one-block EAGLE head reaches the depth where the geometric series wins.

The price of depth is capture, not compute: each extra K adds a verify
``(B, K+1)`` cudagraph rung and a drafter ``(B, K)`` rung, and every
captured graph retains its working set in the shared private graph pool
for the whole boot. On a tight card that pool is taken straight out of
KV. :func:`price_depth_ladder` prices each candidate depth before
anything is allocated and :func:`deepest_affordable` picks the deepest
one that still leaves a servable KV pool — the same "price it, then
refuse with numbers" shape as
:func:`arbi_serve.weight_quant.head_quant._assert_vocab_head_quant_fits`,
so an over-ambitious ceiling is skipped cleanly instead of OOMing
mid-capture.

True iff this head can be run autoregressively past its block count.

The chain in :meth:`MtpDriver._draft_live` feeds step ``s``'s output
back as step ``s+1``'s input, so the head's ``forward`` must accept a
``prev_hidden`` carry and be able to hand back the post-norm hidden
(``return_hidden=True``). A head without that contract can only ever
emit as many drafts as it has blocks — chaining it would silently
feed garbage, so the ceiling clamps to :func:`blocks_shipped`.

Deepest draft depth this deployment is willing to run, pre-VRAM.

``max(blocks shipped, max_chain_depth)`` — a checkpoint that ships
more blocks than the configured chain ceiling still runs all of them
natively — clamped to :func:`blocks_shipped` when the head cannot
chain, and to ``kernel_max_k`` (the tkv served verify ``block_m - 1``)
when the codec imposes one. Never below 1: a bundled head always
drafts at least one token.

Price every candidate depth against the post-weights VRAM budget.

``non_kv_bytes_at(k)`` returns the non-KV bytes the boot would hold
at depth ``k`` (captured graph pool + the depth-scaling pools +
verify tail). ``budget_bytes`` is what is left of the gmu-capped card
after weights / activation reserve / external context — i.e. what the
non-KV pools and the KV pool must share.

Pure and CUDA-free so the ladder is unit-testable on CPU.

Deepest priced depth that still leaves a servable KV pool.

``floor_depth`` is the depth this boot would run absent any pricing
(the head's block count). It is returned unchanged when nothing
fits, so pricing can only ever refuse the extra depth beyond the
floor — it can never make a boot that fits at ``floor_depth`` fail.
The authoritative gates downstream (``compute_kv_budget``,
``assert_capture_sweep_fits``, the post-capture headroom check) still
fail loud on a genuinely over-budget config; this is the chooser that
keeps them from ever having to.

Boot-time reservation of the TP ``all_gather`` staging arenas.

WHY THIS EXISTS. The live-eager ``all_gather`` staging strand is the one
collective buffer that is VOCAB-scale: ``TiedLMHead.forward`` gathers a
``(rows, vocab/tp)`` logit slab, so one staging entry is ``rows x
vocab_local`` elements. ``rows`` is the live ``B`` on plain decode and ``B x
(K+1)`` on the MTP verify step, so it takes a different value on almost every
step of a continuously-arriving server.

Before :class:`~arbi_serve.distributed.parallel_state.GroupCoordinator` grew
its gather arenas, each distinct value minted its own permanent buffer. That
made the staging strand grow, silently, on the SERVING path — after the KV
sizer had already handed the remaining VRAM to the KV pool. Measured
vocab 248320 (``vocab_local = 124160``):

    rows seen  = {1..16} (decode) u {6, 12, ..., 96} (verify)

i.e. the cache cannot fit. The ``torch.empty`` that overran it raised
``torch.OutOfMemoryError`` on ONE rank *inside* the collective region — the
peer had already enqueued the paired ``all_gather_into_tensor`` — and the TP
group then sat in NCCL for the full watchdog timeout. See
``tests/test_nccl_gather_arena.py`` for the arithmetic this module pins.

Reserving the arena HERE — during pre-capture warmup, before Phase 4c reads
``torch.cuda.mem_get_info`` to size the KV pool — makes the bytes ACCOUNTED by
construction: the KV pool simply sizes itself around a reservation that is
already resident, instead of being handed headroom the serving path then
takes back.

Rows the widest live-eager ``lm_head`` all-gather can carry.

``max_batch`` rows on plain decode (one logit row per sequence) and
``max_batch x (K + 1)`` on the MTP verify step, whose flat ``B x S``
layout is the widest shape the engine ever gathers. Pure arithmetic.

``(vocab_local, weight)`` of the tied lm_head shard, or ``None``.

Resolved off the model's ``lm_head`` — the ONLY caller of the vocab-scale
all-gather. Returns ``None`` for an untied / absent head (nothing to
reserve) rather than guessing a width.

The text tower's hidden width, or 0.

Resolved off the EMBEDDING MODULE first (``VocabParallelEmbedding`` stores
``hidden_size`` explicitly), because that is the module whose output the
reduce actually stages. Falls back to the weight's trailing dim and then to
the HF config — the latter two only as a safety net, since a quantized head
may carry packed storage whose trailing dim is not the logical hidden size.

Elements the widest live-eager ``all_reduce`` can carry.

The live-eager reduce strand is ACTIVATION-scale: every staged reduce is a
``(tokens, hidden)`` activation, and the widest ``tokens`` a single forward
can present is the per-step token budget ``max_batched_tokens``. Pure
arithmetic; returns 0 when either quantity is unresolvable.

Pre-allocate the TP reduce arena for this engine's widest shape.

Same contract as :func:`reserve_gather_staging` — returns the bytes
reserved (0 when there is nothing to reserve) and NEVER raises, because a
failed reservation is a lost optimisation rather than a boot failure and
the arena still grows on demand (loudly).

WHY. ``VocabParallelEmbedding.forward`` all-reduces a ``(tokens, hidden)``
activation from OUTSIDE any compiled region, so before the arena landed
every distinct token count minted a permanent staging buffer. Under chunked
prefill that set is unbounded, and it grew on the SERVING path — after the
KV sizer had already handed the remaining VRAM to the KV pool. Measured
the ``torch.empty`` that overran it raised ``torch.OutOfMemoryError`` on
ONE rank *inside* the collective region — the peer had already enqueued the
paired ``all_reduce`` — wedging the TP group until the watchdog fired
(arbi-serve#1302). Reserving here bounds the strand at ``max_batched_tokens

Resolve this TP group's peer access, once, at boot. No consumer required.

Its own boot step rather than a side effect of arming a transport. The
copy-engine tier declines before it reaches any topology question on every
shipped deployment (``min_rows`` is 0), so a resolution that lived inside
that arm would never run — and then the fact nothing else can be routed on
would not exist. The probe is the seam #2172 is named for: the engine has
never asked what the interconnect is, and three separate sites carry "this
box has no NVLink" as a comment premise that nothing checks.

Peer access decides which of three collective tiers a TP group can use, and
they want opposite answers from it, so the fact is resolved here and each
tier reads it:

==================  =========================================
peer access         tier
==================  =========================================
available           async-TP (fused all-gather / reduce-scatter)
absent              the copy-engine host-staged exchange
unresolved          NCCL, unchanged
==================  =========================================

Returns the verdict, or ``None`` when the ranks disagreed or the group is
trivial. Costs two ``all_gather_object`` calls at boot and nothing after.

Arm async-TP for this boot, on the shapes this checkpoint actually issues.

The row-parallel widths are read off the MODEL rather than configured: the
predicate is about the weights, and a list an operator maintains by hand
drifts from the checkpoint the moment either moves. Every
``RowParallelLinear`` in the model contributes its per-rank contraction
width, and each distinct width is timed once against the exposed
all-reduce it would replace.

Arm this engine's copy-engine all_reduce tier. Returns bytes reserved.

The tier's landing buffers must be resident before the deferred KV resize
reads ``mem_get_info``, for the same reason the two staging arenas beside
it must: bytes that appear after the sizer has handed the rest of VRAM to
the KV pool are bytes the pool already promised away. So the arm sits here,
in the boot phase that reserves the TP collectives' memory, and nothing on
the serving path can reach it.

Default OFF (``ARBI_TP_COPY_ENGINE_MIN_ROWS=0``). The arm itself resolves
every remaining term — group size, the prefill capture mode, and the pair's
peer capability — and refuses loudly rather than silently falling back, so
an operator who asked for this transport is never left believing the copy
engines are carrying a reduce that NCCL is.

NEVER raises, and the collectives inside the arm are why that matters more
here than for the arenas: the arm issues an ``all_gather`` and a
``broadcast``, so a rank that raised past them while its peer did not would
wedge the group in NCCL rather than lose an optimisation. Every term this
function reads to decide whether to CALL the arm is rank-symmetric — the
flag comes from the shared process environment, the group and the config
are the same object on both ranks — so the two ranks enter and leave the
collective region together.

Pre-allocate the TP gather arenas for this engine's widest shape.

Returns the bytes reserved (0 when there is nothing to reserve: TP=1, no
tied lm_head, CUDA unavailable, or the arena disabled). NEVER raises — a
failed reservation is a lost optimisation, not a boot failure, and the
arena still grows on demand (loudly) if it was under-sized.

Incremental ``<TOOLCALL>[...]</TOOLCALL>`` detector for NemotronVoiceChat's
SEPARATE function-logits channel.

Why this is a new, minimal class rather than reusing
:class:`arbi_serve.engine.tool_parsing.StreamingToolCallExtractor` (the
existing, fully-built SSE-path tool-call extractor): that class's whole
design is shaped around ONE interleaved text stream, where content and
``<tool_call>...</tool_call>`` markers share the same buffer and the
extractor must separate them so the SSE generator can stream clean
``content`` deltas alongside framed tool-call deltas. NemotronVoiceChat
doesn't have that problem — the function channel carries NOTHING but
(potential) tool-call markup; there is no interleaved "content" to
separate out, so the fast-path/careful-path split that class exists for
has no work to do here. Two further shape mismatches make reuse actively
wrong rather than merely unnecessary:

  * the marker is spelled ``<TOOLCALL>``/``</TOOLCALL>`` (uppercase, no
    underscore) — confirmed against the real reference template
    (``nim_extract/s2s/prompt_template.jinja``) and the reference's own
    parser (``offline_voicechat.run_fc_offline_inference``), NOT the
    Qwen3-style lowercase ``<tool_call>``/``</tool_call>`` that
    :mod:`arbi_serve.engine.tool_parsing` is hardcoded to look for;
  * the body is a JSON ARRAY of ``{"name": ..., "arguments": ...}``
    entries (``<TOOLCALL>[{...}, {...}]</TOOLCALL>``, one block can name
    several calls), not one ``<tool_call>{...}</tool_call>`` block per
    call the way Qwen3's template (and therefore
    :func:`arbi_serve.engine.tool_parsing.extract_tool_calls`) emits.

So this module ports just the "detect a complete marker, parse its JSON
body" shape, adapted to array-of-calls, over a caller-supplied plain text
stream (the function channel, already sampled + detokenized by
:class:`arbi_serve.runtime.nemotron_voicechat_stt_step.EngineSttStep` —
this module does no sampling/detokenization itself, so it is pure
Python/CPU-testable, mirroring how :mod:`tool_parsing` itself only ever
sees already-decoded text).

One parsed call from a ``<TOOLCALL>[...]</TOOLCALL>`` block.

``arguments`` is always a JSON STRING (OpenAI tool-call shape,
matching :func:`arbi_serve.engine.tool_parsing.extract_tool_calls`'s
own convention) even though the reference's raw JSON may carry
``arguments`` as a nested object — see :meth:`NemotronFunctionChannelExtractor.feed`.

Stateful, incremental ``<TOOLCALL>[...]</TOOLCALL>`` detector.

One instance per turn (owned by
:func:`arbi_serve.realtime.nemotron_voicechat_turn.stream_nemotron_voicechat_turn`).
:meth:`feed` is called once per STT step with that step's DECODED
function-channel text fragment (may be ``""`` for a step whose
function token has no visible text); it accumulates into an internal
buffer and returns ``None`` until a complete
``<TOOLCALL>...</TOOLCALL>`` block has arrived, at which point it
returns the parsed calls (an empty list on a malformed/unparseable
body — treated as "no calls", not an error, matching
:func:`arbi_serve.engine.tool_parsing.extract_tool_calls`'s own
"skip what doesn't parse" convention) and resets to scan for a
POSSIBLE FURTHER round (the reference's own chat template supports
calling additional tools after a ``<TOOL_RESPONSE>`` is injected —
see the template's own trailing sentence).

One duplex connection's function-calling state — the object
:attr:`arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession.fc_state` holds (design doc §7.9.16).

It bundles the things a full-duplex FC cycle needs to survive across
ticks, because they share a lifetime and a reset:

  * :attr:`extractor` — the connection-persistent incremental
    detector. Turn-based mode builds one per turn
    (``stream_nemotron_voicechat_turn``); a duplex connection has no
    turn boundary, so it builds exactly one, for the whole connection.
  * :attr:`in_progress` — the reference's ``_fc_in_progress`` gate.
    While ``True``, :func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step
    .finish_duplex_tick` overrides the tick's sampled TEXT token to the
    checkpoint's text PAD id before the frame-lockstep feedback, the
    Site-3 BOS/EOS switch and the TTS subword all read it — so the
    agent goes silent for the duration of the call without the 80 ms
    frame grid or the recurrence itself pausing.
  * :attr:`ack_token_ids` — the ``ack_messages`` "on hold" phrase, as
    forced text tokens (design doc §7.9.18). One id per 80 ms frame,
    spoken INSTEAD of the PAD the gate would otherwise force.

Lives in this module rather than on the session dataclass so the
runtime-layer tick (:mod:`arbi_serve.runtime.nemotron_voicechat_duplex_step`)
and the realtime-layer WS session can both name it without either
importing the other. Dropping the whole object (``fc_state = None``,
which :meth:`~arbi_serve.realtime.nemotron_voicechat_turn.ToolResultBridge
.wait_for_result`'s timeout escape hatch does) resets all of it at
once: the gate reopens, the detector starts scanning clean and any
unspoken ack is abandoned, which is exactly the right state after an
abandoned call.

Drop an opening marker that is never going to close.

On a duplex connection this extractor is CONNECTION-lifetime
(``DuplexFcState.extractor``), and the buffer only ever resets
when a block closes. So a turn cut short mid-block — by a
barge-in or a forced EOS, which is exactly what an over-eager
turn gate produces — leaves ``<TOOLCALL>[{"name": ...`` in the
buffer permanently. That is not merely a leak: :meth:`feed`
matches the FIRST ``_OPEN``, so every later, complete block is
parsed as `stale-prefix + intervening speech + real body`, fails
``json.loads``, and is silently discarded as "no calls". One
interrupted attempt therefore disables tool calling for the rest
of the connection, and it looks from outside exactly like a model
that simply never calls tools.

The reference bounds the same scan with ``MAX_TOOL_TOKENS`` and
per-state frame timeouts (design doc §7.52.7's divergence 6);
this is the same idea at character granularity. Once an unclosed
marker has been open for longer than any real block could be, the
marker and everything before it is discarded and scanning
continues from after it — so the NEXT block still works.

The OPERATOR's view of the live card: a few named groups, and warnings.

:mod:`arbi_serve.engine.live_card` closes the device total as several dozen
rows named in engine terms. That row set is the contract and is not touched
here — it stays the stable API, and every row of it remains in the payload.
What it is not is something a person running the server can read: the question
"is there enough KV, and what is eating the card" is not answered by forty
lines of allocator bookkeeping, most of them zero.

This folds those rows into the handful of categories an operator can act on —
KV cache, model weights, the drafter, CUDA graphs, working memory, allocations
no pool claims, driver overhead, the step reserve, free VRAM a named rule holds
back, and the unused remainder — and lifts the rows that mean something is
WRONG into an explicit warning list, so a healthy card renders no warnings at
all and an unhealthy one cannot be read as healthy.

EVERY GROUP OPENS INTO ITS OWN NAMED CLAIMS (``items``). A bar answers "how
much"; a bar that folds a residency gate, an operator ceiling and a rounding
cushion into one figure cannot answer "what for", and that is the question
behind every "why is my context small". ``items`` is the group's own card rows,
in the pool's own words, each stating what KIND of number it is — and, where a
peak was measured, stating the reserve BESIDE the peak so the slack between
what is held and what was ever seen is a glance rather than a subtraction. It
is not a second grouping: the items of a group are exactly the rows its
``rows`` field names.

ONE MODEL AT TWO ZOOM LEVELS, not two schemes. Each group is exactly the rows
of one or more taxonomy CATEGORIES (:data:`_CATEGORY_GROUP` is the whole rule),
so zooming from a group into its rows is opening it, never switching vocabulary.
Groups are then sorted into TENURE sections — what KIND of number each is, from
:data:`~arbi_serve.runtime.pool_taxonomy.CATEGORY_TENURE`. Without that axis a
weight slab, a driver context, a reservation nothing has taken and a span of
address space with no pages behind it render in one column, where they read as
four allocations. Two of them exist, and only one of those is ours to place.

Two distinctions are deliberately NOT groups, because making them groups would
double-book bytes:

  * "resident but idle" is a COLUMN, not a bucket. A pool's reserved-minus-live
    is the same physical its own row already counts, so it travels as
    ``idle_bytes`` on the group that holds it. A separate idle bucket would
    count those bytes twice and make each pool's row read smaller than it is.
  * ghost address space is neither resident nor free. It takes part in no sum,
    so it is reported BESIDE the groups (``ghost_bytes``) rather than inside
    one — which is also why it needs its own taxonomy category rather than a
    flag on a ``driver.*`` row.

Server-side by rule. The grouping is arithmetic over the card, and a console
that sums rows itself silently books every row the payload omitted into
whatever it computed, under a label it chose. The console renders ``groups``,
``tenures`` and ``warnings`` and adds nothing.

The group's TEXT is the DISPLAY layer: operator language for identifiers the
payload keeps in engine language. It comes in two registers, and the split is
the whole readability contract. ``blurb`` is the ONE short plain sentence a
reader sees under the bar; ``detail``, ``growth``, ``lever``, ``cost`` and
``split_text`` are everything else the group has to say, which the console
carries on hover. A page where every row states its lever, its price, its
in-use split and its tenure inline is correct and unreadable — the reader
cannot find the shape of the answer, because every row is arguing its case at
full length. Nothing is dropped; it is one gesture away. ``growth`` also
answers the only question a group name has to survive — "if this grows, what do
I do about it?" — and a group whose answer is the same as its neighbour's is
one group, not two. ``rows`` on each group names the card rows folded into it,
so the diagnostic view and this view are provably the same bytes.

TWO ROWS NAME A THING, NOT JUST A QUANTITY. "Model weights — 13.56 GiB" and
"External drafter — 0.50 GiB" are unreadable without knowing WHICH model at
what bit width and WHICH drafter, and the KV row's byte count says nothing
about what is inside a byte. Those identities arrive in ``evidence``
(:mod:`arbi_serve.engine.card_identity` reads them from the checkpoint's own
index, the loaded weights' geometry and ``cfg.mtp`` — never from a served name,
which spells out a quantisation for one checkpoint and not the next), and
:func:`_identity_text` puts them on the line. A term the caller could not read
is not stated: the static sentence stands, rather than a placeholder.

THREE ZOOM LEVELS, one model. The taxonomy was already two of them; what it
never had was a top.

  * Level 1 — WHERE THE CARD WENT. ``groups``, in CLAIM ORDER: the rows that
    take their share first, then the KV cache, which is last because it is the
    residual — everything above is claimed and KV grows into what is left.
    Ordering it anywhere else invites the reading that KV is a peer cost you
    could trim, and it is not: "why is my context small" is answered by reading
    UPWARD from the last bar. Each group carries its ``lever`` (the flag or
    knob that moves it) beside its ``cost`` (what moving it takes away), and
    states its bytes as ``in_use`` / ``idle`` / ``promised`` rather than as one
    figure, because a group that took 1128 MiB and is sitting on 974 of them is
    not comparable with a group that took nothing and merely promised 237, and
    rendering both as a percentage of the card invites exactly that comparison.
  * Level 2 — WHY THAT NUMBER IS WHAT IT IS. ``row_drivers``: per CARD ROW, the
    measured inputs its number is a function of (captured graph count and the
    metered bytes per graph, the module baseline and the boot it was measured
    on, the kernel-stack limit and the thread count it multiplies). A group's
    ``drivers`` is its own rows' drivers concatenated — one mechanism, so
    opening a group can never show a line the row does not, and the static half
    of the same answer lives on the row's taxonomy entry as ``function_of``.
    Every entry is a value the engine already measured, passed in as
    ``evidence``; this module invents no constant and models nothing.
  * Level 3 — the card's own rows, unchanged, behind the diagnostics
    disclosure.

ONE CATEGORISATION, WHICH IS THE POOLS'. Every group here is exactly the rows
of one or more taxonomy CATEGORIES, so a name at the top of the page is a name
at the bottom of it. This module adds no vocabulary of its own — no bucket that
is not a pool category, no sectioning by a property the taxonomy does not
carry. What it adds is ORDER (claim order), TEXT (a lever and its price) and
ARITHMETIC the console must not do. A second arrangement of the same bytes, if
one ever earns its keep, ships as an explicit toggle between named views, never
as a silent second scheme — which is the failure this whole surface exists to
prevent, one zoom level up.

CROSS-CUTS ARE VIEWS, NEVER GROUPS (:data:`OPERATOR_VIEWS`). The true cost of
CUDA-graph capture is spread over two groups — the capture pools are ours, the
instantiated exec allocations and the cubins the capture phases loaded are the
driver's — and an operator deciding whether to keep the capture ladder needs
the sum. Summing it into a group would double-book the driver's share, and the
card's whole contract is that every byte is on exactly one row. So a cross-cut
is carried BESIDE the groups, states the group it overlaps and by how much, and
takes part in no total. A view that did not name its overlap would be a second
partition wearing the first one's clothes.

Pure arithmetic + naming: no CUDA, no engine state, fully CPU-testable. The
measured inputs Level 2 states arrive as the ``evidence`` argument, already
read by the caller (``proc.admin_registry.memory_live_snapshot``), so this
module stays a function of its arguments.

One bar of the operator view.

``key`` is stable (a console keys colour off it); ``label`` is what a human
reads.

THE DEFAULT VIEW IS ``label`` + BYTES + PERCENT + ``blurb``, AND NOTHING
ELSE. So ``blurb`` is one short, flat, plain sentence — what this holds,
said the way a person would say it, with no flag names, no allocator
vocabulary and no second clause. Everything the group has to say beyond
that lives in ``detail`` (what it holds, at length), ``growth`` (what a
reader DOES when it grows), ``lever`` (the knob) and ``cost`` (its price),
all of which the console carries on hover. A page where every row argues
its case at full length has no shape a reader can find; the argument is
kept, one gesture away.

``growth`` is also the test for whether a name earns its place, and two
groups with the same answer fail it.

``tenure`` is what kind of number the group's bytes are. Declared here
rather than derived because :data:`OTHER_GROUP` has no category to derive
from; that it equals the tenure of every category folded into the group is
asserted by ``tests/test_operator_card.py``, which is the seam that keeps
the declaration honest.

One CROSS-CUTTING view: a cost that is spread over more than one group.

Not a group, and the difference is load-bearing. The groups PARTITION the
card — every byte on exactly one row, the rows summing to the device — and
that property is the only reason any figure on this page can be trusted. A
cost that straddles the partition (CUDA-graph capture pays the allocator,
the driver's exec heap and the driver's cubin residency) cannot be made a
group without booking the driver's share twice.

So it is carried BESIDE the groups: it takes part in no total, it names
every row it claims, and it states which group each of those rows is
already counted in and by how much. ``overlap_bytes`` is that statement as
a number, and a view that could not produce it would be a second partition
pretending to be the first.

``rows`` are card rows the view claims WHOLE. ``evidence_parts`` names the
keys in the ``evidence`` block that supply the parts no single row carries
— a share of a row, measured by something other than the allocator.

The named pieces of one cross-cut, each with the group it is already in.

A row the card does not carry contributes nothing rather than a zero: a
view listing rows that do not exist on this boot reads as a breakdown with
holes in it.

The cross-cutting views, each stating the overlap it has with the groups.

``overlaps`` is per GROUP, not per row, because that is the sentence a
reader needs: "1.8 GiB, of which 476 MiB is also inside Driver overhead".
Without it the two figures on one page add to more than the card and the
reader has no way to know which one is wrong — the answer being neither.

Return the operator group ``name`` belongs to.

``category`` is the row's own category as the card reports it; the row name
wins where :data:`_ROW_GROUP` overrides it. A name whose category is not a
known engine category lands in :data:`OTHER_GROUP` rather than being
dropped — an unnamed byte is the failure this whole surface exists to
prevent.

Bytes as a reader would say them, at a precision the number can carry.

Formatted HERE, like every other figure on this view, because the console
is not allowed to compute: a client that divides is a client that can
divide by the wrong thing, and the unit is part of the claim. The scale
steps down as the quantity does — a 114 KiB owner list rendered as "0 MiB"
is a row telling a reader it holds nothing, which is the exact defect this
line exists to fix.

A card row's name as a reader would say it, and nothing more.

Mechanically the pool name with its category prefix dropped and its
underscores opened out, so this adds no vocabulary: the word under the bar
is the word on the row in Diagnostics and in the taxonomy. A hand-written
label per row would be a second name for the same thing, and the second
name is the one that goes stale.

The ONE string a console prints beside an item, formatted here.

A reservation and an allocation are different kinds of number and the text
says which, because the byte count cannot: "384 MiB" leaves a reader with
the question, and "384 MiB reserved — no step measured against it" answers
it. Where a peak exists it is printed BESIDE the reserve rather than
instead of it, so the slack between what is held and what was ever seen is
one glance rather than a subtraction.

A group's own card rows, named, in descending size — the group OPENED.

The bar answers "how much"; a bar that folds a residency gate, an operator
ceiling and a rounding cushion into one number cannot answer "what for",
and that is the question an operator asking where their context went is
actually asking. So every row the group holds gets its own named line, and
each line says what KIND of number it is.

Rows holding no bytes are left to Diagnostics, which keeps them all: a term
that is 0 because a surface is off is worth finding when you are hunting
one and is noise under a bar you are scanning.

The two rows that name a thing, not just a quantity.

``Model weights`` and ``External drafter`` are the only groups whose bytes
belong to a NAMED artefact, and the static blurb can only say "the model"
and "the drafter". Where the caller read this boot's checkpoint
(:mod:`arbi_serve.engine.card_identity`), the line says which model, at what
bit width, and — for the drafter — which family and checkpoint, or that the
draft head ships inside the main one and so weighs nothing here.

A term the caller could not read is simply not stated. The fallback is the
static sentence, never a placeholder: "unknown quantisation" beside a real
byte count is worse than not raising the subject.

The KV row's ONE line: how many bits a key and a value are stored in,
and whether that is the same on every layer.

The row's byte count says nothing about what is inside a byte, and "is my
context small" is a different question at 4 bits than at 8. The SHAPE of the
precision is the second half: one width everywhere, or a width per layer.

Level 2, per CARD ROW: the measured inputs each row's number is made of.

Per ROW rather than per group, and a group's ``drivers`` is then its own
rows' lines concatenated. One mechanism, so a group can never show a line
its rows do not, and the two zoom levels answer "why is it this size" with
the same words. The STATIC half of that answer — the term list the value is
a function of, independent of any boot — lives on the taxonomy entry as
``function_of``; this is only the half that needs a measurement.

A row with nothing measured gets NO entry. That is a correct answer: the
alternative is a page where every row has an explanation and half of them
are invented, which is worse than a page with gaps a reader can see.

The rows that mean something is wrong, as operator-readable warnings.

Every entry is derived from a verdict the SERVER already reaches — the card
failing to close, a ghost row the allocator does not back, a leak-signal row
with bytes on it, a residual past the alarm threshold, a row whose number is
a guess rather than a measurement. Nothing here is a threshold invented for
the view, so a healthy boot produces an empty list.

The card's two totals, named apart, with the gap between them stated.

``total`` is what this process can address —
``torch.cuda.mem_get_info``'s device total, the total the live card closes
against, and the denominator of every ``pct`` on this view.
``physical_total_bytes`` is the whole card, which NVML reports and which is
LARGER: the driver holds a carve-out (page tables, ECC, its own
bookkeeping) that no allocation can reach.

The carve-out is not folded into a group. The groups partition the
addressable total and that partition is why any figure here can be trusted;
a row for memory no pool can ever claim would be a claim on a card it is not
on. It is stated beside them instead, so a reader can follow the physical
card down to the addressable total and then through the bars with no
unexplained remainder anywhere.

``physical_total_bytes`` of ``None`` means NVML could not be asked. That
renders as unavailable — never as equal to the addressable total, which
would be the page asserting a carve-out of zero that nothing measured.

Fold a closed live card into operator groups, tenure sections, warnings.

Args:
    card: the ``card`` block of ``GET /v1/admin/memory/live``
        (:func:`arbi_serve.engine.live_card.build_live_card`). Ghost rows
        are excluded from every group, exactly as they are from the card's
        own sums, and reported apart as ``ghost_bytes``.
    evidence: MEASURED inputs the engine already holds that no row of the
        card carries — the instantiate meter's graph counts, the driver
        attribution's module baseline and stack geometry, the boot heap
        walk's named unpooled bytes, the observed peaks under
        :data:`OBSERVED_PEAKS_KEY` and the arena's boot-only share under
        :data:`ARENA_RESIDUE_KEY`. Level 2 is written from these and from
        nothing else, so a caller that passes ``None`` gets a view with the
        drivers absent rather than a view with the drivers invented. A term
        this engine does not publish renders as UNAVAILABLE and says so;
        nothing here substitutes a zero for a measurement nobody took.
    physical_total_bytes: the card's PHYSICAL capacity, as NVML reports it.
        Larger than ``card["total_bytes"]``, which is what the process can
        address; the difference is the driver's own carve-out. Passed in
        rather than read here because this module folds a card and measures
        nothing, and ``None`` — NVML unavailable — renders as unavailable
        rather than as a carve-out of zero.

Returns:
    ``{groups, row_drivers, tenures, views, statements, warnings, total_bytes,
    free_bytes, resident_bytes, in_use_bytes, idle_bytes, promised_bytes,
    unclaimed_bytes, ghost_bytes, ghost_rows, closes, summary,
    summary_detail}``. ``groups``
    covers every row of the card, in display order, with zero-byte groups
    dropped; ``sum(g.bytes) == total_bytes`` whenever the card closes, so
    the view inherits the closing property rather than restating it. Each
    group carries ``items`` — its own rows with bytes on them, named and
    formatted, largest first — so opening a bar names the claims inside it
    rather than handing a reader one number to take on faith.
    ``tenures`` is the same groups sorted by WHAT KIND of number they are,
    with the section totals computed here so a console never sums a bar it
    was shown; a tenure with no group is dropped, so an engine that holds
    no reservation shows no reservation section. ``summary`` is the
    one-line reassurance a reader gets instead of arithmetic to check.
    ``views`` are the cross-cuts (:data:`OPERATOR_VIEWS`), beside every
    total and stating their overlap. ``statements`` is the headline: four
    or fewer sentences that put the card's bytes into exists / in use /
    idle / promised before a single bar is read.

Taxonomy names that fall into :data:`OTHER_GROUP`.

Empty is the contract: a name the taxonomy declares and this module has no
operator group for renders as "Unrecognised", which is a defect in the
mapping rather than a fact about the server.

Whether this group's bytes are physical that EXISTS right now.

Derived from :attr:`tenure` rather than declared, so the flag a console
dims a bar with and the section it files the bar under cannot disagree.

What KV basis this engine is serving under, read from the runtime.

FLIPPING the basis is NOT here. It belongs to the attention-backend swap
(``POST /v1/admin/attention_backend`` with ``calibration_path`` +
``oscar_rotation_path``), because the codebooks and the rotation they were fit
in are ONE UNIT -- a bundle is only valid in its own basis, and the codec
refuses a mismatch in both directions. Carrying them together on the backend
variant is what makes re-selecting a prepared Hadamard or OSCAR variant
instant rather than a recalibration, and it removes the surface on which an
operator could pair a bundle with a basis it was never fit in. A separate
"turn OSCAR on" switch would offer exactly that surface.

What is left here is the READ the swap surface does not provide: whether a
rotation is installed, and -- separately -- whether a bundle fit in it exists.
Those are different questions and only the second one predicts whether this
engine can actually serve.

Where this deployment keeps calibration artifacts.

Derived in ONE place so the status read and the flip cannot look in
different directories -- a status that reported on a path the flip does
not use would say "no rotation" about an engine that has one.

Re-exported from the calibration module that OWNS this derivation.

Kept as a name here because the admin surface reads better for it, but the
definition lives in one place -- see
:func:`arbi_serve.calibration.oscar_fit.default_rotation_path`.

What basis this engine is serving under, right now.

READ-ONLY, and it reports what the RUNTIME resolves rather than what a
config says it should be: the rotation is installed by assigning an env
var (or, for an embedded basis, into tkv's process cache), so the only
honest answer comes from asking tkv what it currently has. A status
endpoint that echoed the request back would report "enabled" for a flip
that silently failed to take.

``bundle_paired`` is the load-bearing field. A rotation installed WITHOUT
a bundle fit in it does not serve -- the codec refuses a bundle in the
wrong basis in both directions -- so "the rotation is on" is not the same
question as "this engine can serve under it", and an operator flipping the
switch needs the second one answered.

Engine-side typed output bus: per-step batch staging + marshaling.

Design (``docs/engine_core_process.md`` §3)
-------------------------------------------
The engine emits **data** — ``TokenOut`` / ``AudioOut`` / ``FinishOut``
structs from :mod:`arbi_serve.engine.proc.messages` — staged into the
step's :class:`OutputBatchBuilder` and handed to the consumer as one
``OutputBatchMsg`` per step. The
:class:`~arbi_serve.engine.client_request.OutputApplier` is the consumer
half; only the *marshaler* between them varies:

  * **inline** (single-loop): :meth:`OutputBus.flush` calls
    ``applier.apply(batch)`` directly — the engine loop is the consumer
    loop.
  * **thread** (``ARBI_ENGINE_OWN_THREAD``, the default serving mode):
    :meth:`OutputBus.bind` holds the HTTP loop; flush hands the batch
    object across via one ``loop.call_soon_threadsafe(applier.apply,
    batch)`` — no serialization, one hop per step.
  * **process** (P4): the same finalized ``OutputBatchMsg`` is
    msgspec-encoded onto the ZMQ output socket; the API child's recv
    thread decodes and feeds the identical applier. The bus's flush
    seam is where that marshaler slots in.

Perf invariants owned here
--------------------------
  * **One hop per step** — N tokens accumulate in the builder; a step's
    single :meth:`OutputBus.flush` fires at most one
    ``call_soon_threadsafe``. Emission sites never hop per token.
  * **Idle cheapness** — an empty builder makes :meth:`OutputBus.flush`
    a two-comparison no-op: no hop, no allocation on the idle tick (the
    builder's lists are reused; fresh lists are allocated only when a
    non-empty batch is handed off, because the consumer thread takes
    ownership of the old ones).
  * **FIFO** — batch order is commit order; consecutive batches carry a
    monotonic ``seq`` the applier gap-checks.

Residual delivery: when the HTTP loop is gone (shutdown teardown) or was
never bound, flush applies the batch inline best-effort so a trailing
finish emitted by the last step still reaches an attached consumer —
the deactivation-residual semantics of the thread bridge.

Mutable engine-side staging for one step's ``OutputBatchMsg``.

Engine-thread-only (append in the commit loop, take at the flush
point) — no lock. ``take`` transfers the filled lists into the
immutable wire struct and installs fresh ones, so a later step's
appends can never race the consumer reading the handed-off batch.

The engine's output seam: builder + marshaler in one object.

Emission sites (``run_step`` / ``detokenizer`` / the SPMD driver)
append structs; the step's single flush point marshals the finalized
batch to the :class:`OutputApplier`. The marshaler is selected by
:meth:`bind` / :meth:`unbind`:

  * unbound (default): inline — ``applier.apply(batch)`` direct call
    (the legacy single-loop mode and every sync test harness).
  * bound to the HTTP loop: thread — one
    ``call_soon_threadsafe(applier.apply, batch)`` per step, the
    batch crossing by reference.

The ``call_soon_threadsafe`` here is one of the few permitted
boundary crossings (``tests/test_transport_boundary.py``).

Finalize the staged step into a batch, or ``None`` when empty.

The empty case is the common idle tick — no allocation happens
there (invariant: idle cheapness). ``seq`` increments only for
emitted batches so the applier's gap check is exact.

Switch to the process marshaler (design §3, P4b).

``emit`` receives each finalized ``OutputBatchMsg`` — exactly one
call per non-empty step, on the flush caller's thread (the engine
loop); serialization happens on the runtime's output IO thread.
Mutually exclusive with the thread marshaler: process mode has no
in-proc HTTP loop to hand batches to.

Hand the staged step to the consumer — at most one hop.

No-op when nothing is staged (the idle tick), so callers invoke
it unconditionally on the hot path. The thread marshaler's
``call_soon_threadsafe`` is fired right before the engine thread
blocks on the per-step GPU event-sync (the callers' placement),
so the HTTP loop applies the batch during the GPU window. When
the bound loop is closed (shutdown teardown), the batch applies
inline best-effort so trailing finishes are never stranded.

Live value domains for overridable params whose bound only the engine knows.

:data:`arbi_serve.config_overrides.PARAMS` carries a static ``choices``
vocabulary for the params whose valid set is the same on every boot.  A
numeric param can also have a real ceiling, but one that is a property of
THIS boot -- the drafter checkpoint that was loaded, the attention codec
that was selected -- so no static registry can state it.

``mtp_n_draft`` is the case that motivates the module: the served depth K
is bounded by the drafter's block width (``block_size - 1`` draft
positions per denoise pass), by the blocks a bundled MTP head ships, and
by the tkv verify kernel's ``block_m - 1``.  Every one of those is read
from the running engine here, so a client can offer the operator the
depths this boot can actually serve instead of a blank number box whose
out-of-range values only fail at apply time.

A bound this module cannot resolve is OMITTED, never guessed: a client
that gets no domain for a param falls back to free entry, which is the
behaviour it had before the domain existed.

Draft positions one DFlash block holds, and where that number came from.

A checkpoint that declares its trained ``block_size`` bounds K at
``block_size - 1`` (``build_dflash_drafter`` refuses anything above
it).  One that declares none leaves the block runtime-owned, and the
loader caps it at the width such heads are known to have been trained
for; that cap is the bound here.

Speculation depths this boot can serve, or ``None`` when unbounded here.

The floor is 1 whenever a drafter source is named -- ``MtpConfig``
refuses an external drafter at K=0 -- and 0 otherwise, where 0 is the
"propose nothing" setting rather than an invalid one.

Every live-derived value domain, keyed by param name.

Best-effort by construction: a param whose bound this engine cannot
resolve is absent from the mapping, and a client without an entry
keeps free entry for that param.

Param classification for the graceful residency swap scheduler.

Two buckets:

  - **Sampling** — per-request, cheap, no engine state change. Already
    surfaced by :class:`arbi_serve.engine.request.SamplingParams`. Two
    requests with different sampling params can sit in the same engine
    step; the sampler reads the per-request fields.

  - **Engine** — per-(model, params) configuration. Things that today
    are set at boot via env vars or CLI flags: ``k_bits``, ``v_bits``,
    ``recipe``, ``calibration``, ``attention_backend`` choice. Changing
    any of these requires either a backend rebuild or a full model
    swap; we group requests by the hash of the engine subset and the
    swap orchestrator decides how expensive the transition is.

The cost ladder used by :func:`cheap_param_diff` is the same as the
swap orchestrator's dispatch table — kept here so the scheduler can
score a candidate transition without instantiating the orchestrator.

Split a flat param dict into ``(sampling, engine)`` sub-dicts.

Unknown keys are dropped silently — clients can pass extra metadata
without it accidentally affecting either bucket. The split is
canonical: a key never appears in both sub-dicts.

Stable 16-hex-char SHA256 over the engine subset of ``spec``.

Sampling-only differences yield identical hashes (cheap per-request
overrides don't fragment the group queue). Key order is invariant.
Unknown keys are dropped before hashing so client-side metadata
can't poison the group key.

Classify the cost class of moving from ``old`` to ``new`` engine
params.

Tiers (highest cost wins on tie-breaks):

  - ``"model"``    — different weights / dtype / max_context → full
    model swap (~1–1.5 s).
  - ``"backend"``  — same model, different bit widths or backend
    selection → backend rebuild + cudagraph drop (100–300 ms).
  - ``"recipe"``   — same model + same backend bits, but
    calibration / recipe override changed → reload calibration,
    re-attach to existing per-layer ops (0–100 ms).
  - ``"sampling"`` — only sampling-bucket fields differ → no engine
    state change (per-request handled).
  - ``"none"``     — fields identical or the only differences are
    keys outside both buckets.

Reads ``old`` and ``new`` as flat dicts; sampling vs engine split
happens internally, so callers can pass the full request param
dict.

Stage-by-stage VRAM census across a park / wake cycle.

A park releases physical in four independent steps (per-tensor sleep state,
named pools, the growable KV slab, the captured-graph execs) and a wake takes
them back in the mirror order. Each step asks the DRIVER for physical
separately, so when a wake does not restore the pre-park free-VRAM state the
only way to say WHICH step lost the bytes is to read the driver between them.

Emits one line per stage: driver free, the cuMem allocator's mapped total, the
growable slab's mapped bytes, and the default caching allocator's
reserved/allocated — plus the delta in driver free since the previous stage, so
a stage that gives back less than it took names itself.

The per-stage line is always emitted — a park is a rare admin operation, and
this is what turns a future wake OOM into a named term instead of a bare driver
error. The wider reads (per-pool table, forward-arena tenant inventory) walk the
allocator's records and the collector, so they are opt-in via
``ARBI_PARKWAKE_CENSUS=1``.

Whether the WIDE reads run (``ARBI_PARKWAKE_CENSUS=1``).

The per-stage census line is unconditional; this gates the per-pool table
and the arena tenant inventory, which cost a records walk and a gc scan.

Bytes the pools' OWN cuMem regions currently have mapped.

The growable KV slab plus the state pools' sentinel-alias arenas: both
drive cuMem directly, so neither is inside the per-tag totals above, and a
park/wake stage that returns or retakes them has to see them here or the
census cannot close its cycle.

Log one census line for ``phase``. Never raises.

``dfree`` is the change in DRIVER free VRAM since the previous stage, which
is the number that closes the cycle: what a park stage returned, a wake
stage has to take back, and no more.

Log every named pool's mapped / reserved / live bytes.

Wider than :func:`census` (one line per pool) — run at each end of the park
and of the wake, where the question is which pool came back a different
size than it went away.

Name every live tenant of ``scratch.forward_arena`` and who owns it.

The arena is per-step scratch, so anything alive at a park boundary is a
tenant that outlives a step. That is what keeps the pool permanently
non-empty, and a non-empty pool is one the idle release
(:func:`~arbi_serve.engine.inprocess_capture.release_idle_forward_arena`)
must refuse — so free physical this process is short of stays stranded in
the pool. Reuses the refusal path's attribution
(:mod:`arbi_serve.engine.arena_tenants`), so a tenant reports the same
dotted owner path here as it would in a boot log.

Fold the small persistent scratch buffers into one slab-backed pool.

``scratch.rope`` (the static cos/sin tables) and ``scratch.penalty_accum`` (the
sampler's additive penalty accumulator and its occurrence counts) are each a few
MiB, allocated once at boot, live for the engine's lifetime, and live at the
same time as each other. A pool per buffer pays the caching allocator's segment
overhang per pool (see :mod:`arbi_serve.runtime.pool_slab`), and a private
``torch.cuda.MemPool`` never hands that overhang back.

This seam allocates ONE slab in ``scratch.persistent_fold``, moves the rope
tables into it, points the accumulator's allocation hook at it, and destroys the
now-empty ``scratch.rope`` pool so its segment returns to the driver. The bytes
land back in KV at the deferred resize, which reads the driver's free after this
runs.

Both consumers keep their own VRAM-ledger row and their own cap key: each carve
is charged to its own sub-tag, which partitions the host pool's mapped bytes
rather than absorbing them (see
:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.charge_subtag`).

``scratch.cublas_workspace`` is deliberately NOT folded. Its block is allocated
by cuBLAS itself, with no Python tensor to move; it is cleared at the pre-freeze
reclaim seam and re-pinned lazily into whatever pool is active at the next GEMM;
and the address captured graphs baked into their GEMM launches lives in that
pool, so handing its segment to another consumer would let a replay write into
storage something else now owns.

BOOT-ONLY, and strictly before the capture sweep: the fold rebinds the rope
tables' storage, which a captured graph that had already baked their addresses
would replay out of.

Bytes the penalty accumulator will take, from the plan's own formula.

Priced against the RESOLVED serving width, so the slab spends what the
accumulator actually allocates rather than the auto ceiling the KV budget
was held out against.

Drop the slab a PREVIOUS model's build left on ``eng``.

``eng._persistent_slab`` is PER-MODEL: its sub-tag budgets are closed forms
of that model's vocab, width and dtype. The engine holds it across a build,
so a second build into the same engine (hot swap, reload, a residency
member) inherits the previous model's slab unless it is retired — and a
slab whose penalty budget is already fully carved refuses the new model's
first carve.

``PoolSlab.release`` frees nothing a live tensor still references: a parked
member's rope tables alias their views and keep the storage. What it drops
is the slab's own reference and the two sub-tag ledger charges — which are
keyed on the unqualified tag, so the process holds ONE row per sub-tag and
the incoming build's charge would replace this one regardless.

Fold the persistent scratch buffers into one slab; return bytes freed.

Runs after the serving width is resolved (the accumulator's size depends on
it) and before the capture sweep. Stashes the slab on the engine as
``_persistent_slab`` for the accumulator's allocation hook to carve from.

Returns 0 and leaves every pool untouched on any guard miss: no fold pool,
no cuMem backing, no rope tables, a rope pool holding something other than
the tables this seam can account for, or a fold that would save nothing.
A miss leaves ``eng._persistent_slab`` at None, never at a slab a previous
model built — see :func:`retire_persistent_slab`. The accumulator then
takes its own pool, which is the documented no-fold path.

The ``(shape, dtype) -> Tensor`` hook that carves ``subtag`` from the slab.

``None`` when no slab was built, so the caller keeps its own pool.

Idempotent per ``(subtag, shape, dtype)``: a config-override variant build
shares the donor member's pools, so it re-installs against a slab the
baseline already carved to its budget. Handing back the same buffer is
what that re-install wants, and a second carve out of a bump allocator
that never frees would raise :class:`SlabBudgetExceeded` instead.

Option B (B3) — frozen, static Phase-2 serving allocator (D2).

After the B2 handoff (capture done, KV grown into the reclaimed transient at
the stable VA), serving is a static role: weights + KV + persistent graphs are
all resident and nothing dynamic remains to overflow. This module locks that
in so serving is provably static and cannot OOM by construction:

  1. **Cap every cuMem-backed pool to its current live mapped size.** The
     boot ``graph_pool`` map-time cap (M2, ``pool_caps.arm_boot_caps``) is
     disarmed by the time we run; we re-arm a cap on every cuMem-backed tag
     at exactly its currently-mapped bytes. A serving allocation that tries
     to grow any pool past its frozen size is denied at map time — a
     contained ``cudaErrorMemoryAllocation`` (a torch OutOfMemoryError the
     caller can catch / log), never a card OOM. (Re-arming the C ``g_any_cap``
     gate is paid only at cuMem map time — i.e. when the caching allocator
     creates a new segment in a cuMem pool. Steady-state decode transients
     reuse already-mapped segments inside their pools — every named pool is
     cuMem-backed since sleep-to-floor, see ``engine.CUMEM_BACKED_POOLS`` —
     so the hot path never touches the gate. ``model.lora`` is capped at its
     operator budget instead of its live size, and the per-step scratch pools
     at their live size plus the serving-floor row that holds their growth
     free — see :data:`FREEZE_CAP_BUDGETED_POOLS`; ``comms.nccl`` and
     ``scratch.forward_arena`` stay uncapped because no row of ours states
     their ceiling, see :data:`FREEZE_CAP_EXEMPT_POOLS`.)

     We deliberately do not also arm the torch-MemPool-level
     ``enforce_size_target`` here: that check fires on a pool's ``use()`` exit
     and would raise on legitimate intra-pool reuse — e.g. the first serving
     step lazily allocates a few-KB per-num_tokens residual_buf inside
     ``graph_buffers_pool``, reusing free space in the pool's already-mapped
     cuMem segments (not new physical) — breaking serving. The cuMem map cap
     (1) is the correct by-construction guarantee: it denies only attempts to
     map new physical past the frozen size (the card-OOM-class event) while
     allowing reuse within already-mapped segments.

  2. **Disable further cudagraph capture + JIT** for the serving phase via
     ``eng._serving_frozen``. Kernels are already compiled + captured at boot;
     the boot pre-capture functions early-return when frozen (defensive — the
     serving path has no on-demand capture trigger today, but a resume hook
     or a future shape-miss must not re-enter capture against a frozen pool).

  3. **Assert the can't-OOM-by-construction invariant** at handoff:
     ``Σ(pool live mapped sizes) + runtime floor ≤ physical``. This is the
     static-server guarantee — every byte is accounted, nothing competes.

cuMem + in-process-capture only. bf16 / non-cuMem boots, or boots where the
in-process capture path did not drive (so KV was not grown at a stable VA), do
not freeze (the ``should_freeze_for_serving`` gate excludes them). Not gated on
``vram_mode``: when the in-process capture handoff ran, the static-server
freeze is correct memory management and applies for bench and serving alike.
Boot-only, one-time; the serving hot path is untouched.

Log a ``[boot-profile]`` span for one freeze sub-step.

The freeze publishes several independent reports (the VRAM ledger, the cap
posture table, the post-freeze class gate, the residency attribution). They
close inside a single boot phase, so without a per-step span the phase line
cannot say which report a slow freeze spent its seconds in.

The freeze cap for a :data:`FREEZE_CAP_BUDGETED_POOLS` pool, or ``None``.

``live_bytes`` plus the bytes of the pool's floor row, read from the
itemisation :func:`~arbi_serve.engine.inprocess_capture.serving_floor_for_grow`
publishes on the boot state — the same rows the ledger and the arena watch
read, so the cap cannot disagree with them. A pool whose row name is empty
is capped at ``live_bytes`` (its allotment is zero by proof).

``None`` when the pool names a row but no itemisation was published: the
allotment is then unknown and the caller leaves the pool uncapped rather
than stating a number nothing measured.

NEVER BELOW ``admission_bound_bytes``. The floor's row is a SHORTFALL —
the part of one step's arena share that the pool's own residency does not
already cover — so it goes to 0 exactly when the boot decides the residency
covers the bound, and ``live + 0`` then states that the pool may never map
another byte. Those are two different claims about the same pool made by
the same boot: the floor sheds a step's arena share against the residency
*because it stakes the floor on the pool being able to serve what a step
asks*, and admission is credited those same bytes and admits slates on
them. A cap at ``live`` refuses the pool exactly the capacity the floor
already spent and admission already promised — and refuses it as a
``nullptr`` from the allocator shim, which torch prints as a card OOM
(:func:`~arbi_serve.engine.pool_caps.drain_cap_denials` is what tells the
two apart). The bound is not a margin and not a constant: it is
:func:`~arbi_serve.engine.inprocess_capture.forward_arena_step_bound_bytes`
at the structural corner of admission's own model, published by the floor
that used it. 0 means the boot published none, and then the cap is the
floor's arithmetic alone.

``ARBI_DEBUG_FREEZE_MEMDUMP``'s output at freeze; returns rows attributed.

Two artifacts, of unequal worth:

  * an allocator pickle, which carries per-allocation frames only on a
    RE-ENTERED build — :func:`arbi_serve.engine.build.build` arms
    ``_record_memory_history`` under ``jit_detector.is_serving()`` because
    the cold boot's compile sweep OOM-kills the recorder on host. On the
    boot that actually sizes KV it names nothing;
  * the default-pool LIVE-owners table, which reads the allocator snapshot
    and the Python heap and needs no recorder, so it populates on a cold
    boot.

Only the table fires the flag's counter. A counter fired by the pickle
passes a cold boot whose table is empty — the one state this diagnostic
exists to make impossible. Best-effort throughout: a diagnostic never wedges
a boot.

``{pool name: reserved-but-free bytes}`` inside each named pool.

overhang". It is NOT an additive term: every named pool is a cuMem-backed
torch ``MemPool``, so these bytes are already mapped physical inside the
which is a number an operator can act on, and it removes the only reason the
same physical ever needed a second name.

A pool that HOLDS CAPTURED CUDAGRAPHS is floored at its cuMem TAG-MAPPED
bytes instead. torch's block view is blind to what a captured graph holds:
``torch.cuda.graph(pool=...)`` drops the capture-time tensors' Python refs,
so their blocks read ``inactive`` even though the graph replays out of them
and the physical can never be reused. Taking torch's reading at face value
operator would correctly try to reclaim and correctly fail to.

ONLY graph-holding pools. Tag-mapped bytes include a pool's mapped-but-FREE
segments, so flooring every cuMem pool at them would report 0 held-free for
all of them — erasing exactly the signal this column exists for (the forward
motivated would never have been found).

NVML framebuffer total for ``device``, or 0 when NVML is unavailable.

``cudaMemGetInfo``'s total is already NET of the framebuffer the driver
keeps for itself, so an operator comparing the boot log against the number
on the box sees a gap the log cannot explain. Reading NVML lets the ledger
state the card, the driver's share and the CUDA-visible total as one line.

Read through the shared NVML session rather than a private
``nvmlInit``/``nvmlShutdown`` pair: the handles are refcounted per process,
and the admin endpoint's capacity block quotes THIS number, so a second
reader is a second place the card's size can be different.

Emit THE boot VRAM ledger — one table whose rows sum to the card.

Replaces the old "serving pool table", which printed a subset and a signed
reconcile gap the reader had to close by hand. Every row here is measured,
the rows are disjoint, and ``driver.residual`` is the remainder of the
whole-card identity rather than of a driver sub-identity. Runs once at
freeze; zero hot-path cost.

Logged at INFO. The previous table was DEBUG, which meant the one surface

The ``driver.residual`` warning, or ``None`` when inside the alarm.

``driver.residual`` is the BOOT-LEDGER row: the whole-card identity's
remainder. It legitimately absorbs driver-resident cubin growth, so a
non-zero value is not by itself a leak. The similarly-spelled
``unpooled.unregistered_pool`` is a DIFFERENT quantity on the caching-allocator
surface — a private ``MemPool`` nobody registered — and that one IS a leak
signal. The alarm text below tells the reader to check it, precisely
because the two answer different questions.

The SIGN is the diagnosis and the two directions need opposite hunts:
NEGATIVE means the rows over-explain the card (two of them count the same
pages), POSITIVE means driver-resident physical grew past what this
configuration is known to need. Reporting a magnitude alone points the
operator at the wrong one. Pure so both are testable.

``modules_baselined`` says whether ``driver.modules_loaded`` was held at a
measured baseline. It governs what a POSITIVE residual can mean: held, the
growth had nowhere else to go, so the unregistered-allocation causes are in
scope; seeding, they are not, because the seeding boot books its own growth.

The published itemisation of ``floor_bytes``, or ``()`` when it is not it.

``serving_floor_for_grow`` publishes its terms on the boot state, and the
ledger reports whichever floor its caller measured. Those are the same
number on the freeze path and a DIFFERENT number on the unfrozen-boot path,
which reports the bare allocator floor. Handing the itemisation of one to the
other would render a breakdown of bytes that reserve does not contain, so
the terms travel only when they add to exactly the floor being reported.

PREDICTED (budget plan) vs MEASURED (VRAM ledger), row by row.

The KV-capacity ledger prints an ``unmodelled residual`` and tells the reader
to find it. This is the surface that finds it: the VRAM ledger says where
every byte IS, ``KvBudget.per_pool_predicted_bytes`` says where the plan
thought they would be, and the row-wise difference is the residual with a
name on every part. Two failure shapes show up here and they need different
fixes — a pool with NO predictor at all (the plan books 0), and a pool whose
predictor is simply low. Reporting only the net hides both.

Pure arithmetic over two dicts: no CUDA reads, no sizing effect.

Emit the HOST PINNED ledger beside the VRAM one. Idempotent, best-effort.

Its own table rather than rows in the VRAM ledger, because that ledger's
contract is that its rows sum to the card and host bytes inside it would
break the identity it exists to hold. But it prints in the same place, at
the same moment, for the same reason: a boot log that states every byte the
engine reserved on the CARD and nothing about what it page-locked on the
HOST is how a reservation nobody priced got the process OOM-killed.

Emitted on EVERY boot, including one with no CUDA and no VRAM ledger — the
host reservation is real either way, and a boot that printed no host row is
indistinguishable from a boot that reserved nothing.

Emit the boot VRAM ledger on a boot that does NOT freeze. Best-effort.

The freeze path is the one that locks a static layout, but it is NOT the
only path that serves: ``--no-cuda-graphs``, ``--no-cumem-pools`` and any
boot where the in-process capture handoff did not drive all skip it. Those
boots used to print no VRAM accounting at all — which fails the bar this
on the configurations an operator reaches for when something is already
wrong. The rows are the same and they still sum to the card; only the caps
are absent.

True iff the B3 freeze should lock this engine's allocator post-handoff.

Gated to the same regime that drove the B2 in-process single-capture
handoff: cuMem-backed pools + the in-process capture path actually ran (so
KV was grown at the stable VA and the pools hold their final, static
footprint). Any other boot (bf16, non-cuMem, subprocess-measure fallback)
is left unfrozen. Not gated on ``vram_mode`` — the freeze is correct memory
management whenever the handoff ran.

Lock the allocator for static Phase-2 serving (B3 / D2). Idempotent.

Called right after the post-capture headroom gate passes and the boot
caps are disarmed. Returns a small dict of the frozen sizes for logging /
test assertions. No-op (returns ``{"frozen": False}``) when
:func:`should_freeze_for_serving` is False.

Raises :class:`~arbi_serve.engine.memory_budget.MemoryBudgetError` if the
can't-OOM invariant (``Σ live mapped + floor ≤ physical``) does not hold —
that would mean the handoff over-grew KV, a build-logic bug we must fail
loud on (it never happens by construction: KV grows only into free space
bounded by ``free − floor − margin``).

Lift the B3 freeze so a hot-swap reload (B4) can re-enter

A reload tears down the model + pools and re-runs the two-phase boot for
the new model/config; the frozen caps + enforce flags belong to the old
model's static footprint and must be cleared before the new boot maps its
weights / captures. Clears the cuMem caps, drops ``_serving_frozen``, and
relaxes ``enforce_size_target`` on the surviving pool wrappers. Idempotent.

By-construction boot-time pool caps (M2/M3 of GMU-0.99 no-OOM).

The cuMem allocator enforces a per-tag map-time hard cap (see
``runtime/cumem_allocator.py``). This module arms that cap on the cudagraph
capture pool right before the capture sweep and disarms it before serving:

  * **M2 — arm.** Cap the ``graph_pool`` tag at ``current_mapped +
    (free_now − floor − margin)``: the pool may map the VRAM that is free right
    now except the runtime floor. An allocation past the cap is denied at map
    time → a contained in-pool OOM → that capture bucket falls to eager (TP1) or
    raises loud + symmetric (TP>1), never a card OOM.

    Both the captured-graph working set and the un-captured warmup forward route
    through the cuMem ``graph_pool``: the capture via
    ``torch.cuda.graph(pool=graph_pool.id)`` (graph-private allocations go
    through the MemPool's registered pluggable allocator), the warmup via
    ``graph_pool.use()`` (``runtime/capture/decode.py`` wraps the warmup
    forwards in it). The warmup transient frees back to the pool's free-list
    between buckets, and the captured allocations reuse those pages, so peak pool
    usage is ``max(warmup, captured)`` — not their sum. Bounding the pool to
    ``free − floor`` therefore makes ``post_capture_free ≥ floor`` hold by
    construction at any GMU — including 0.99.

  * **M3 — graceful gate.** With caps armed, a capture "OOM" is a contained
    cap denial, so the post-capture gate must not refuse on it (see
    ``verify_post_capture_headroom``); it verifies the floor invariant
    (which holds by construction) and proceeds, logging eager-fallback count.

  * **disarm.** ``disarm_boot_caps`` lifts the cap after the gate so the
    serving hot path pays zero overhead (the C ``g_any_cap`` gate goes to a
    single dead branch).

Boot invariant: ``free ≥ floor`` is asserted at arm time (fail loud) — the
card must be able to fund the runtime floor. ``gpu_memory_utilization`` is not
part of that test: it caps the engine's TOTAL device usage, while the runtime
floor is one of the engine's own reservations inside that cap, so every value
up to and including 1.0 arms normally.

ENFORCED vs OBSERVED
--------------------
A cap is only a guarantee if it can actually deny, and an over-tight cap costs
KV exactly as surely as an absent cap costs safety. Every cap therefore carries
an explicit posture, and :func:`log_cap_report` states it LOUDLY at every
report point (OBSERVE is never a silent default):

  * **ENFORCED** — ``capture.cudagraphs`` (this module). The value is
    ``already_mapped + (free_now − floor − margin)``: derived from a MEASURED
    ``mem_get_info`` at arm time, so it denies only an allocation that would eat
    the runtime floor — precisely the allocation that card-OOMs today. It cannot
    refuse a capture that would otherwise have fit in free VRAM.
  * **ENFORCED** — the Phase-2 freeze caps (``phase2_freeze.freeze_for_serving``)
    with a NON-ZERO live size, and the operator-configured ``model.lora``
    budget. Each is a measured live-bytes value for a pool that is static for
    the serving phase; ``comms.nccl`` alone is left uncapped.
  * **ENFORCED** — the freeze's budgeted scratch caps
    (``phase2_freeze.FREEZE_CAP_BUDGETED_POOLS``): live mapped bytes plus the
    serving-floor row that holds the pool's post-freeze growth free. Never
    tighter than what the floor promised, so a map the floor paid for is not
    refused, and a pool growing past the floor is a contained in-pool denial
    rather than a card OOM.
  * **OBSERVED** — any freeze cap whose measured live size is ZERO. "This pool
    must never map another byte" is not a claim a zero measurement supports: an
    empty-at-handoff pool that fills lazily would be refused on its first real
    request. Those caps are armed in observe mode and every breach is reported
    at ERROR, naming the pool to promote.

WHO READS THE DENIALS AFTER BOOT
--------------------------------
An ENFORCED cap denies by returning ``nullptr`` from the allocator shim
(``runtime/_cumem_shim.py``) BEFORE ``cuMemCreate`` is reached, so the driver is
never asked. Torch formats a ``nullptr`` from its pluggable allocator the same
way whatever produced it, which means **a cap denial and a genuine card OOM
print the identical message** — including a "GPU 0 has ... free" figure that is
the CARD's and says nothing about the pool that refused. A reader of that
message alone cannot tell the two apart, and every reading of one as driver
pressure is a diagnosis the evidence did not support.

The fact that separates them is recorded at the point of denial
(:class:`~arbi_serve.runtime.cumem_allocator.TagCapBreach`) and was, until
:func:`drain_cap_denials`, only ever printed by the two BOOT calls to
:func:`log_cap_report`. Nothing read it at serving time, so the one number that
names the refusing pool sat in the allocator for the life of the process.
:func:`drain_cap_denials` is that reader: the serving OOM path
(:func:`~arbi_serve.engine.memory_pressure.note_oom`) drains it on every
refusal, and :func:`log_cap_report` publishes the BASELINE the drain measures
against, so a boot-time denial (a capture bucket contained into eager fallback
— the mechanism working) is never charged to a serving request.

Arm the cuMem ``graph_pool`` hard cap for the capture sweep.

Returns True iff a cap was armed (cumem on, driver available, cuda graphs
on). Sets ``eng._pool_caps_armed`` so the post-capture gate knows the
capture was contained (a card OOM is impossible by construction; any
over-budget bucket fell to eager). Never raises on a benign no-op; raises
``MemoryBudgetError`` only when MEASURED free VRAM is already below the
runtime floor, which no ``gpu_memory_utilization`` can fix.

Not gated on ``vram_mode``: the by-construction containment cap is correct
memory management and always runs when cuMem + cuda_graphs are on. The
bench-vs-serve distinction (fail-loud vs warn on an incomplete/under-
captured engine) lives solely in ``verify_post_capture_headroom`` — bench
still fails loud on ``capture_oom_failures > 0`` (eager fallbacks would
skew a benchmark) even though the card itself never OOMs.

Report, then lift, all boot caps. Idempotent.

The report runs BEFORE the lift so the boot log carries what each cap
actually did — how much of it was used, whether it ever denied anything,
and (the failure mode this machinery exists to make impossible) whether it
was INERT because its tag was never allocated.

Log the per-tag byte ledger + every armed cap's posture and usage.

Returns the :class:`~arbi_serve.runtime.cumem_allocator.CapRow` list (also
stashed on ``eng._pool_cap_report`` for triage/tests).

Loudness contract:
  * every ENFORCED cap logs at INFO with its live usage,
  * every OBSERVED cap logs at WARNING — an unenforced cap is never allowed
    to look like a guarantee,
  * an OBSERVED cap that was BREACHED logs at ERROR and names the byte
    overshoot (it would have denied; promote or raise the value),
  * an ENFORCED cap whose tag was NEVER ALLOCATED logs at ERROR as INERT and
    raises in bench mode — a cap that cannot bind is not protection,
  * counter drift (the running cap counter disagreeing with a live scan of
    the allocations) logs at ERROR: the gate would be budgeting a fiction.

What the armed caps say about one allocation refusal.

Three determinate states, and the point of the type is that they stay
distinguishable:

* ``denials`` non-empty — an armed cap refused a map in this window. The
  refusal is CONTAINED and the pool is named. Whatever the card's free
  VRAM says, the driver was not the one that refused.
* ``denials`` empty with ``caps_armed`` > 0 — the gate was armed, watched
  this window, and refused nothing. That is a real null control: the
  refusal came from the driver.
* ``denials`` empty with ``caps_armed`` == 0 (or no allocator at all) — the
  C ``g_any_cap`` gate is disarmed, so the shim's precheck never ran and a
  cap COULD NOT have refused. Also the driver, by elimination.

A VALUE snapshot of every reported breach counter.

``CapRow.breach`` hands out the allocator's LIVE record, so a stashed row
re-read later reports the counter's value then, not at report time. The
baseline has to be ints.

Which armed cap refused a map since the last drain, if any.

Pure bookkeeping: dict reads under the allocator's lock, no CUDA call, no
logging inside the allocator, nothing that can re-enter the allocator. It
is therefore safe on the one path that has to call it — an OOM handler,
where a second exception replaces the fact being diagnosed.

The window is closed by the read: each drain consumes the refusals it
reports, so a second OOM reports ITS OWN refusals rather than re-reporting
the first one's. The window OPENS at the last :func:`log_cap_report` (the
Phase-2 freeze report is the last one a boot writes), which is what keeps a
contained boot-time capture denial from being read as a serving fault.

Σ ``limit_bytes`` over the ENFORCED caps in a :meth:`cap_report`.

The physical the caps, all reached at once, can let their pools hold.
OBSERVED caps deny nothing, so they bound nothing and are left out.

State the bound the per-tag caps add up to against the whole card.

Each cap bounds one pool. Only their sum, beside the serving floor, says
whether the caps can all be reached at once and still leave the floor on
the card: ``Σ ENFORCED cap limits + floor ≤ physical``. One line at INFO
when it holds; at ERROR when it does not, and in bench mode that raises
(the posture an INERT cap takes) — a benchmark must not report from an
engine whose caps over-commit the card. Returns the signed slack.

One operator sentence: which of the three states this is, and why.

Written to be read beside torch's OOM text, which cannot carry any of
it — see the module docstring.

Post-freeze allocation accounting — the gate behind ``gpu_memory_utilization`` 1.0.

The Phase-2 freeze (:mod:`arbi_serve.engine.phase2_freeze`) locks the VRAM
layout: every cuMem pool is capped at its live mapped bytes and the KV pool is
sized so ``serving_floor_for_grow`` bytes stay DRIVER-free. From that moment the
engine has no headroom beyond that floor, so every allocation it can still make
must either come out of the floor or be impossible.

``gpu_memory_utilization`` does not change that. It caps the engine's TOTAL
device usage, while the serving floor is one of the engine's own reservations
INSIDE that cap, applied as ``max(serving floor, (1 − util) × total, member
floor)`` (:func:`arbi_serve.engine.inprocess_capture.serving_floor_for_grow`).
At 1.0 the utilization leg goes to zero and the measured serving floor is the
whole protection — which is why the floor has to be complete rather than merely
large. This module makes "complete" a checked property.

Two gates, one static and one measured.

:func:`assert_post_freeze_classes_reserved` walks :data:`POST_FREEZE_CLASSES`:
every allocation class REACHABLE for this config must name either the serving-
floor term that reserves it or the reason it cannot allocate. A class that is
reachable and does neither fails the boot, so a new serving allocation is caught
at the freeze rather than by the first request that reaches it.

That walk is over allocation SITES, and a site-shaped registry cannot see a
whole ALLOCATOR that is free to grow: the freeze caps every cuMem pool at its
live mapped bytes except the ones in
:data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_BUDGETED_POOLS` (capped at
live plus a serving-floor row) and
:data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_EXEMPT_POOLS` (uncapped),
and each of those may map new physical after the layout is locked.
:func:`_registry_self_check` therefore derives a second obligation from those
two lists — every such pool must be named by a covered entry's ``pool`` — so
the lists cannot drift apart and a pool that can still grow cannot exist
without a stated reserve or proof.

:func:`assert_post_freeze_floor_intact` closes the loop with a measurement. The
boot's own post-freeze work — the serve-kernel warmup and the readiness forward —
is the floor's first consumer, and it reaches the JIT, sampler, codec and
spec-decode classes a live request reaches. It also runs past the last
driver-residency bracket, so it is the only reading of what those classes cost.
Consuming more than the floor proves an unreserved class exists; the boot then
refuses instead of serving a layout that will OOM under load.

Fragmentation
-------------
A sufficient total does not by itself make a large late allocation serviceable.
The floor is held as DRIVER-free VRAM, not as caching-allocator reserve, and
that is what makes it serviceable: ``cudaMalloc`` draws from the driver's own
physical pool at allocation granularity, so a single large late request is
placed whole as long as the floor covers it. PyTorch's ``reserved − allocated``
overhang is deliberately NOT credited against the floor
(:func:`arbi_serve.engine.inprocess_capture.serving_floor_for_grow`) — those
bytes are split across size classes and private cuMem pools and cannot host a
large contiguous transient. ``expandable_segments`` is force-disabled
process-wide (:mod:`arbi_serve`) because cuMem pools are mutually exclusive with
it, so segments are size-classed and the overhang can be neither merged nor
relied on.

The consequence is a requirement on the floor, not a margin: it must cover the
LARGEST SINGLE post-freeze allocation, not only their sum. The serving-floor
terms are per-class peaks rather than aggregates, so this holds for every class
carrying a reserve; a class whose transient scales with a request rather than
with the config cannot satisfy it and is recorded here as unreserved.

One allocation class the engine can still reach after the memory freeze.

Args:
    name: taxonomy-style ``owner.thing`` identifier, stable across boots.
    site: ``module:line`` of the allocation itself — the evidence for the
        claim, and where a reviewer checks it is still true.
    reachable: predicate over the engine deciding whether THIS config can
        reach the class at all. Must not touch CUDA.
    reserved_by: the serving-floor term that holds the bytes free, named as
        :func:`~arbi_serve.engine.inprocess_capture.serving_floor_for_grow`
        names it. Empty when the class is covered by a proof instead.
    proof: why the class cannot allocate after the freeze — a pre-freeze
        allocation, a capped pool, or a structural impossibility. Empty when
        the class is covered by a reserve instead.
    pool: the registered named pool the allocation lands in, when it lands
        in one. Read by :func:`_registry_self_check` against
        :data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_BUDGETED_POOLS`
        and :data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_EXEMPT_POOLS`,
        so a pool the freeze lets grow past its live size cannot exist
        without an entry here stating what covers it. Empty for a class
        that allocates from the
        default caching allocator or outside every allocator.

At most one of ``reserved_by`` / ``proof`` carries the coverage; a class with
both would be an unreviewable double claim, and one with neither is the
defect :func:`assert_post_freeze_classes_reserved` refuses on.

Gemma-4-class assistant drafter: reads the verifier's KV, writes none.

The own-slab drafter sweep declines these heads, so their chain is captured
through the driver's own ``(B, K)`` ladder — inside the capture pools, whose
freeze caps bound it.

Reject a registry that cannot be read as an accounting.

Names must be unique, so a class is looked up by exactly one entry.

Every pool in :data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_BUDGETED_POOLS`
and :data:`~arbi_serve.engine.phase2_freeze.FREEZE_CAP_EXEMPT_POOLS` must be
named by a COVERED entry's ``pool``. Those pools are the complete set of
allocator surfaces the freeze leaves able to map new physical, so one
without an entry is a post-freeze consumer with nothing stated to hold its
bytes free — the exact defect this module refuses on, and the one a
hand-written registry of allocation SITES cannot see. Deriving the
obligation from the freeze's own lists makes the omission unrepresentable:
adding a pool there without stating what covers it fails at import.

Post-freeze classes this config reaches that carry no reserve or proof.

A class whose ``reachable`` predicate raises is treated as REACHABLE: an
unreadable engine is not evidence that an allocation cannot happen.

Refuse a freeze that would lock a layout an unreserved class can overrun.

Raises :class:`~arbi_serve.engine.memory_budget.MemoryBudgetError` naming
every reachable class with neither a covering serving-floor term nor a proof
it cannot allocate, and the site each one allocates at.

Measure the floor against the boot's own post-freeze allocations.

The serve-kernel warmup and the readiness forward run AFTER the freeze and
reach the JIT, sampler, codec and spec-decode classes a live request
reaches. They also run past the last driver-residency bracket, so the free
VRAM they consumed is the only reading of what those classes cost. It is
drawn from the floor the KV grow held out, so consuming more than the floor
proves a class outside it, and the boot refuses.

No-op when the freeze recorded no measurement (an unfrozen boot, or CUDA
unavailable) — the gate reports what was measured and never invents a
reading.

Engine <-> API-child process boundary: wire layer + process runtime.

The process half of the EngineCore split (``docs/engine_core_process.md``):
typed messages (§5), the msgpack codec, the ZMQ socket/framing helpers
(§2), the explicit admin RPC registry (§6, ``ADMIN_METHODS`` /
:func:`dispatch_utility`), the engine status snapshot builder (§5.2,
:func:`build_stats_msg`), and the process plumbing itself: the
engine-side IO runtime (:mod:`.runner`), the API-child entry
(:mod:`.api_child`), the child supervisor + §7 sequencing
(:mod:`.supervisor`), and the API-side
:class:`~arbi_serve.engine.proc.client.EngineClient` (§6). This package
is the only place allowed to import zmq (the
``test_transport_boundary.py`` AST lint gates it).

``api_child`` / ``supervisor`` are not re-exported here: importing the
package must stay light (the child imports it before deciding what it
is), and their entry points are reached by qualified module path from
``cli/bootstrap.py`` / mp-spawn.

Explicit admin/control RPC registry (``ADMIN_METHODS``).

Design (``docs/engine_core_process.md`` §6)
-------------------------------------------
Every admin reach-in crosses the process boundary as a
:class:`~arbi_serve.engine.proc.messages.UtilityCallMsg` whose ``method``
is looked up here — an explicit registry, never ``getattr`` — and whose
JSON args are validated against the entry's typed arg schema before the
engine-side callable runs. Unknown method / bad args raise the typed
errors below, which the dispatcher maps to
``UtilityResultMsg(ok=False, error=...)``: loud failure, never a silent
no-op.

Execution model: :func:`dispatch_utility` runs on the caller's loop —
the engine loop in production, so ``CriticalSection``
acquire/drain/mutate happens entirely in one process+loop. The registry
itself is loop-agnostic (thread/inline modes route the same call through
``run_on_engine``): nothing here touches sockets, threads, or a specific
event loop.

Marshaling rules (§6):

* :class:`~arbi_serve.models.layer_spec.StateKind` args cross as their
  ``.value`` string and are reconstructed engine-side
  (``StateKind(value)``); enum-keyed result dicts are normalized back to
  ``.value`` keys by :func:`json_safe`.
* Every result is JSON-safe by construction — :func:`json_safe` fails
  loudly on anything it cannot represent, so a live object leaking into
  a result is a boot-visible bug, not wire corruption downstream.

Methods whose engine-side logic lives inline in an HTTP route body
(debug-stats reads, timeline rows, the live memory snapshot, LoRA
listing, profiler toggles) are also provided as helper functions in this
module, mirroring the route logic exactly, so those routes stay thin
``engine_client.call(...)`` shims.

Normalize a method result to JSON-representable data, loudly.

Enums (notably :class:`StateKind`, both as values and as dict keys)
become their ``.value``; tuples/sets become lists. Anything outside
the JSON data model raises ``TypeError`` — a live engine object in an
RPC result is a registry bug that must surface at the call, not as
codec corruption on the wire (§5.3: live objects never cross).

Identity normalizer for a result this process built out of JSON data.

:func:`json_safe` exists to catch a LIVE ENGINE OBJECT leaking into an
RPC result, and it earns its cost on a method that returns whatever an
engine attribute happens to hold. It does not earn it on a result this
module constructed itself, key by key, out of ``int`` / ``float`` /
``str`` / ``None``: there the walk re-derives, on every poll, a fact
that was settled when the data was built — and it is a FULL walk of
every key and every value, which on the request timeline is thousands
of dicts.

So the guarantee moves from the call to a TEST. A method registered
with this normalizer must be covered by an assertion that
``json_safe(result) == result`` — the walk still runs, once, in CI,
where a regression is a red test rather than a per-poll tax. See
``tests/test_timeline_serialised_once.py``.

One typed argument in a method's schema.

``kind`` is the JSON-level type the wire value must already have —
coercion is validation plus the enum reconstructions, not stringly
casting (``"4"`` does not become 4):

* ``"str" | "int" | "float" | "bool" | "list" | "dict"`` — exact
  JSON types (``float`` also accepts int; ``int`` rejects bool).
* ``"any"`` — any JSON value (heterogeneous payloads such as
  ``set_dynamic_config.value``).
* ``"state_kind"`` — a ``StateKind.value`` string, reconstructed to
  the enum engine-side.

``default`` is :data:`_REQUIRED` for mandatory args. ``allow_none``
lets an explicit ``None`` cross for nullable args.

One explicit ADMIN_METHODS entry.

``resolver`` maps the engine to the bound callable/coroutine that
implements the method; ``args`` is the typed schema the JSON args
dict is validated against; ``normalizer`` post-processes the raw
return into JSON-safe data (default :func:`json_safe`).

Whether the engine's drain/swap lock is currently held.

The engine-side source for the admin routes' 409 pre-check
(``_require_no_swap_in_progress``): in process mode the API child
cannot read ``eng._swap_lock`` directly, so it asks. Advisory by
nature (the authoritative serialization is the lock itself); read
defensively so an engine without the lock reports "not locked",
matching the route helper's historical behavior.

Flag-truth snapshot: contract evaluation over the live counters.

The counters live in the engine process (module globals bumped at the
perf-path execution sites — see :mod:`arbi_serve.flag_truth`), so this
read is implemented as an RPC, exactly like the debug-stats surfaces.
Rank 0's counts are the serve counts under rank-symmetric SPMD.

``GET /v1/admin/request_timeline`` body
(server/routes/admin/timeline.py). The ring is a module-global deque
written by the engine finish path — engine-process state post-split.

Also merges in live per-request rows for anything still queued or
mid-flight (:func:`in_flight_rows`), read straight off
``eng.scheduler`` at call time — this is per-request "is it stuck
right now", not an aggregate gauge, and adds nothing to the request
hot path since it only runs on this poll.

The ``driver.*`` rows as of NOW, falling back to the boot-measured ones.

Held boot rows are correct for boot and stale under load: the local-memory
pool grows when the driver raises the kernel-stack limit for a kernel first
launched while serving, and cubins keep loading on first dispatch. Every
such byte lands on the caller's residual unless the terms are re-measured,
which is why the panel's whole-card residual grows away from the boot
ledger's zero. ``driver.residual`` itself is excluded: the consumer computes
the WHOLE-card residual and would otherwise add the two. (``driver.residual``
is the ledger-side remainder; the allocator's ``unpooled.unregistered_pool`` is a
different row on a different surface and is not touched here.)

``driver.foreign_process`` is re-read here too. It is the one row that is not
ours at all, and the card is device-wide: a co-tenant that arrives after boot
is resident memory no held row can name.

Skipped (held rows returned) while a pool ``use()`` window or an engine
build is open: an allocation inside one holds the allocator mutex with the
GIL released, so an allocator query from this thread deadlocks the engine.
Taken as a LEASE, not a predicate read — checking and then querying leaves
a window an engine thread can enter, and that window is a measured hang.
The probe itself runs unsettled (``probe_device(settled=False)``) for the
same class of reason.

cuMem bytes the allocator tracks but has UNMAPPED, or 0 when unreadable.

``tracked_bytes - mapped_bytes`` over the process-singleton pluggable
allocator: exactly the allocations whose physical went back to the driver
through ``cuMemUnmap`` + ``cuMemRelease`` while ``torch.cuda.memory_reserved``
kept counting their address space. It is the MEASUREMENT behind the
``address_space.released_pool_va`` row's claim to hold nothing.

``dev``'s PHYSICAL capacity, or ``None`` when NVML cannot say.

NVML is the only source that reports the whole card;
``torch.cuda.mem_get_info`` reports the addressable part and nothing else.
Returning ``None`` rather than the addressable total is the point: a caller
must be able to say the card's size is unknown, and cannot do that if this
function hands back a number that looks like an answer.

NVML's own ``reserved`` field is deliberately not returned beside it. It
does not reconcile the two totals — measured 483 MiB against a 454 MiB gap
on an Ada card — so the carve-out worth reporting is the difference between
the two totals actually in play, which the caller has both halves of.

MEASURED inputs the operator view's Level 2 is written from.

Every entry is a number some part of the engine already took: the
instantiate meter's own sample sizes and bracketed deltas, the driver
attribution's held module baseline and the stack geometry behind
``driver.local_memory``, the boot heap walk's named unpooled bytes. None of
it is derivable from the card, which is why the view cannot compute it and
why it travels as an argument rather than as a constant in
:mod:`arbi_serve.engine.operator_card`.

Best-effort by construction. A term this engine cannot produce is ABSENT,
and the view then renders that driver line missing rather than rendering a
guess — the failure mode a "0" would produce is a reader concluding the
capture ladder is free.

``GET /v1/admin/memory/live`` body (server/routes/admin/memory.py).

Reads the CUDA allocator + named-pool module globals, which live in
the engine process — the reason this becomes an RPC.

Self-sufficient: every row a consumer needs to close the card identity
comes from this one snapshot, including ``state.attn_kv.mapped`` (the
growable KV slab, invisible to the allocator snapshot). ``unresolved``
names every contribution this snapshot could NOT establish, so a consumer
can refuse to render a closed breakdown instead of folding the missing
bytes into its residual.

``unpooled_owners`` is READ here, not computed: the boot's serving freeze
walks the heap and publishes the rows, because naming those storages holds
the GIL for on the order of a second and this method executes on the engine
loop — computing it here would stall every sequence in flight. The resolver
falls back to walking only when no freeze published an answer for this build
generation, and says so at WARNING when it does.

``GET /v1/admin/distortion_analysis`` body
(server/routes/admin/distortion.py) with the eval-prompt resolution
engine-side: ``eng.eval_prompts`` is engine-side state, so the
route ships ``num_prompts`` and the fallback list lives here.

Runs a real forward pass (long-running) — mandatory engine-process
execution; the process client's DRAIN_GATED timeout class covers
it. ``eng.distortion_analysis`` itself is sync, so it is offloaded
via ``asyncio.to_thread`` so the long analysis never stalls the
loop it dispatches on (it has no CriticalSection; it coexists with
serving by design).

Resolver for a method the Engine class itself exposes.

Explicit per-entry binding — the registry decides which attribute,
so an unknown RPC name can never reach ``getattr`` (fail-loud
contract). Missing attribute = mis-built engine → AttributeError at
the call, which is correct.

Execute one admin RPC against ``eng`` on the caller's loop.

The single entry the process dispatcher (and the thread/inline
``run_on_engine`` path) routes every ``UtilityCallMsg`` through.
Unknown method / bad args raise :class:`AdminRpcError` subclasses —
the dispatcher maps them (and any method-raised exception) to
``UtilityResultMsg(ok=False, error=...)``. Returns JSON-safe data
(a dict for every method except ``aunload_lora``, which returns
``None`` per its route's 204 contract).

Exception class names along ``type(exc).__mro__``, concrete first.

Stops before ``Exception`` (and anything above it) — the names are
the HTTP-mapping keys the route layer matches against its historical
``except SomeError`` clauses, and every one of those is a proper
``Exception`` subclass. Shipping the MRO (not just the leaf name)
preserves subclass semantics across the wire: an engine-side
``PeftValidationError(ValueError)`` still matches a route mapping
keyed on ``"ValueError"``.

Execute one wire ``UtilityCallMsg`` → ``UtilityResultMsg``.

The canonical engine-loop executor behind
``EngineProcRuntime.utility_handler`` (§6): decode the JSON args,
:func:`dispatch_utility`, encode the JSON-safe result. Every failure
— registry :class:`AdminRpcError` or a method-raised exception —
becomes ``ok=False`` with ``error=str(exc)`` plus the exception's
:func:`exception_type_names`, so the API-side route layer can map
the failure to the same HTTP status its in-proc ``except`` clauses
produce. Never raises: an admin failure must reach the caller as a
result, not kill the engine loop.

API-child process entry: uvicorn + the API-side ZMQ IO threads.

Design (``docs/engine_core_process.md`` §1, §2, §7)
---------------------------------------------------
:func:`main` is the mp-``spawn`` target the engine's supervisor launches
(§1 step 2). Order of operations, each load-bearing:

  1. ``PR_SET_PDEATHSIG=SIGKILL`` — an engine hard-kill can never orphan
     a listening HTTP socket (§1; the PID watcher is the non-Linux and
     race-window fallback).
  2. ``configure_logging`` — the child owns its own loguru/uvicorn log
     bridge (log config does not survive spawn).
  3. **Bind ROUTER + PULL before uvicorn starts** (§2: the child is the
     stable side — the engine connects whenever its build finishes; ZMQ
     queues anything that arrives before our IO threads run).
  4. Build the child app (:func:`arbi_serve.server.app.build_api_app`) —
     tokenizer + routes, no engine import, no CUDA.
  5. Run uvicorn programmatically (``uvicorn.Server``) with its default
     signal handlers suppressed — SIGTERM policy is ours (§7: the child
     defers to the engine's drain sequencing).

The API-side IO (one thread per socket, §2):

  * **output drain** — PULL recv → sentinel match / type-byte decode →
    one ``loop.call_soon_threadsafe(client.handle_message, msg)`` per
    whole message (coalescing is already per-step by construction).
    ``client`` routes to caches / RPC futures / the injected
    :class:`~arbi_serve.engine.proc.client.MessageSink` applier seam.
  * **intake send** — ROUTER recv learns the engine DEALER identity from
    ``HelloMsg`` (§2), then drains the submit outbox
    (:meth:`ApiChildRuntime.submit` — HTTP handlers enqueue through
    ``EngineClient.submit``) with the identity prefix.
  * **engine-PID watcher** — polls the engine PID (+ heartbeat
    staleness as the trigger to check); a vanished engine marks the
    client dead and stops uvicorn with a nonzero exit (§7: orchestrator
    restarts the container).

The API child's transport half: sockets, IO threads, liveness.

Owns the two bound sockets (§2: the child binds, the engine
connects) and the :class:`EngineClient` handed to the HTTP layer.
Separated from :func:`main` so tests can drive the full socket
behaviour without uvicorn / FastAPI / a tokenizer on disk.

Child SIGTERM: mark draining, wait for the engine's clean
``EngineDeadMsg`` (or the grace window + slack), then stop uvicorn.

The engine owns the drain (§7) — a container-stop SIGTERM reaches
both processes, and the child must not race the engine's final
flush out from under in-flight SSE streams.

API-child process entry (mp ``spawn`` target; see module doc).

``cfg_bytes`` is the pickled :class:`~arbi_serve.config.ServerConfig`
(fully picklable — frozen dataclasses of primitives); ``endpoints``
the run-scoped IPC pair; ``engine_pid`` the spawning engine's PID for
the liveness watcher.

Queue one intake message for the engine (any thread; cheap).

Messages enqueued before the engine's ``HelloMsg`` arrives wait
in the outbox — the ROUTER cannot address a peer it has not
heard from (§2).

Engine-PID liveness watcher (§7 hard-kill path).

pdeathsig already covers parent death on Linux; this closes the
non-Linux gap and the case where the engine is not our direct
parent's last thread standing. Heartbeat staleness alone (PID
alive) is not death — /health degrades to 503 via
``client.heartbeat_stale`` instead.

Record engine death (any thread) + stop serving. First wins.

Ordering contract: ``engine_dead_event`` is published last, so
that observing it set means death has been fully handled, not
merely noticed. Concretely, a thread that sees the event set may
rely on all of:

  * ``exit_code()`` is final (``_engine_dead_clean`` assigned);
  * the child has already been asked to stop serving
    (``_request_exit`` has returned);
  * the death notice has already been posted to the HTTP loop,
    so ``client.mark_engine_dead`` → ``sink.on_dead`` will run
    (posting is ordered; the loop callback itself is async, so
    waiters that need the sink still wait on the sink).

Exactly-once is enforced by ``_dead_latched`` under
The
side-effects run outside the lock: the latch alone gives
exactly-once, and holding a mutex across
``loop.call_soon_threadsafe`` and an arbitrary caller-supplied
exit requester is how a shutdown path deadlocks.

Latch a fatal boot-contract violation and stop serving.

Wired to the ReadyMsg tokenizer-fingerprint validation: a child
whose tokenizer diverges from the engine's would silently encode
with one vocabulary and detok with another — every request
corrupt, no crash. Exit nonzero so the orchestrator restarts the
container instead. First latch wins; callable from any thread.

Same ordering contract as :meth:`_note_engine_dead`:
``boot_fatal_reason`` is the public signal (``exit_code()`` keys
off it), so it is assigned last — a thread that reads it set may
rely on the exit having already been requested. Exactly-once is
enforced by the separate ``_boot_fatal_latched`` under
``_terminal_lock``, since the public signal is not set until the
work is complete and cannot double as its own idempotency guard.

API-child engine facade: the ``app.state.engine`` surface in
process mode.

The HTTP routes were written against the live in-proc :class:`Engine`.
In process mode the engine lives in another process, so the child's
``app.state.engine`` becomes this facade — the routes' engine surface
re-expressed over the typed wire:

* **submit** — :meth:`asubmit` delegates to
  :func:`arbi_serve.engine.submission.asubmit_process` (SubmitMsg over
  the intake socket, SubmitAck future resolution). The facade owns the
  child-side halves: the request-id counter (§5.1 — ids are allocated
  API-side), the :class:`~arbi_serve.engine.client_request.OutputApplier`
  (fed by the output socket through
  :class:`~arbi_serve.engine.proc.client.ApplierMessageSink`), and the
  :class:`~arbi_serve.engine.loop_bridge.SubmitAckRegistry`.
* **cancel** — :meth:`cancel_by_id` sends a ``CancelMsg`` (the engine's
  tombstones handle a cancel racing its own submit).
* **admission** — :meth:`admission_check` recomputes the route-level
  gate from the cached ``StatsMsg`` / ``ReadyMsg`` (§6: no RPC on the
  admission path). Approximate by design (stats lag one heartbeat);
  the engine-side scheduler + the HTTP semaphore remain the strict
  bounds.
* **admin** — hot-path-free reach-ins (``aswap_attention_backend``)
  route through :meth:`EngineClient.call` utility RPC (§6).
* **read surface** — ``generation_defaults`` load from the child's own
  model dir (the same loader the engine build uses); ``backends`` /
  ``active_specs`` / the health fault latches mirror the cached
  ``ReadyMsg`` / ``StatsMsg`` so helpers like
  ``engine_unhealthy_reason(eng)`` read the same truth the probes do.

* **residency routing** — ``stable_va`` is a
  :class:`~arbi_serve.engine.residency_view.CachedResidencyView` fed by
  the residents snapshot (``ReadyMsg`` trailing fields at boot,
  ``residency_changed`` control events after — see
  :meth:`apply_residency_snapshot`). The routes' ``_maybe_switch_model``
  runs unchanged against it (the mode-agnostic
  :class:`~arbi_serve.engine.residency_view.ResidentRouting` seam);
  ``aswitch_model`` goes through utility RPC. The view is approximate by
  design — the engine-side RequestFactory re-validates every routed
  request (``SubmitMsg.target_resident``), so a stale cache can never
  mis-serve.

What is deliberately absent (routes degrade to their documented no-op /
error paths): ``model`` (route-side placeholder expansion and speech-out
gating instead read the weightless ``mm_bindings`` / ``omni_output_spec``
properties below),
a real ``lora_store`` (presence checks are optimistic; the engine-side
factory remains the authority and rejects an unknown LoRA with the same
400).

Route-level LoRA presence shim (process mode).

The route pre-check (``_resolve_lora`` → ``lora_store.has``) exists
to fail fast in-proc; in process mode the loaded-adapter set lives
engine-side. ``has`` answers True so the request proceeds to the
engine, whose RequestFactory validates against the real store and
nacks an unknown LoRA with the identical 400 — one authority, no
stale mirror. ``GET /v1/loras`` and load/unload go through utility
RPC, never through this shim.

Install a residents/groups snapshot (HTTP loop).

Called with the ``ReadyMsg``-derived dict at boot/respawn and
with every ``residency_changed`` control-event payload. The
routing attributes are swapped atomically enough for the
single-threaded HTTP loop (no awaits between writes).

Stable-VA resident switch via utility RPC (§6 DRAIN_GATED).

The registry runs ``Engine.aswitch_model`` under the engine
critical section (drain → park active → wake target). Engine-side
``KeyError`` (unknown/dropped record) is re-raised as ``KeyError``
so ``_maybe_switch_model``'s ``except KeyError → 404`` mapping
holds. The cached view updates via the trailing
``residency_changed`` event (emitted before the RPC result on the
FIFO output socket, so it is already applied when this returns).

Cancel by id → ``CancelMsg`` on the intake socket.

Engine-dead is swallowed here (nothing left to cancel; the
death fan-out already error-finished the client half) — callers
are cleanup paths that must never raise.

Route-level admission from the cached StatsMsg/ReadyMsg.

Approximate (stats lag ≤ one heartbeat) — the engine scheduler
and the HTTP semaphore stay the strict bounds. ``priority`` is
accepted for signature parity; the priority-aware queue shaping
runs engine-side where the real queue lives.

Hot-swap via utility RPC (§6 — engine-loop CriticalSection).

Serves the per-request ``?attention_backend=`` swap
(``_request_ctx._maybe_swap``); the admin route goes through the
``admin_call`` dispatch instead. An engine-side
``ValueError`` is re-raised as ``ValueError`` so the caller's
``except ValueError → 400`` mapping holds.

API-side engine client: utility RPC, cached status, liveness.

Design (``docs/engine_core_process.md`` §6, §7)
-----------------------------------------------
:class:`EngineClient` is the only engine-facing object HTTP-side code
sees in process mode (``app.state.engine_client``). It keeps every
transport detail out of the route surface:

  * **utility RPC** (§6): :meth:`call` builds a ``UtilityCallMsg`` with a
    fresh ``call_id`` and awaits the matching ``UtilityResultMsg`` via an
    ``asyncio.Future`` (vLLM pattern). Method registry + typed dispatch
    are engine-side; this client just carries JSON blobs.
  * **cached status**: the latest ``StatsMsg`` / ``ReadyMsg`` are cached
    on receipt so ``/health`` and metrics gauges read memory, never RPC
    (§6 "no RPC on the probe path").
  * **liveness** (§7): :attr:`engine_alive` flips on ``EngineDeadMsg`` /
    sentinel / PID-watcher notice; every pending RPC future fails loudly
    at that moment rather than hanging.

Threading: everything here runs on the HTTP event loop. The api_child
output-drain thread hands whole decoded messages across via one
``call_soon_threadsafe(client.handle_message, msg)`` each (§2);
:meth:`handle_message` then updates caches, resolves RPC futures, and
forwards the data-plane messages to the injected :class:`MessageSink`
(the applier seam wired to ``OutputApplier``).

The applier seam: consumer half of the output pipeline (§3).

:class:`ApplierMessageSink` binds the data plane to
:class:`~arbi_serve.engine.client_request.OutputApplier` (look up the
``ClientRequest``, append text, fire ``new_token_event``, record
finishes); its status-plane hooks handle tokenizer reload on
``model_switched`` and SSE error frames on death. Every method runs on
the HTTP event loop and must not block. :class:`NullMessageSink`
(fail-loud on the data plane) keeps the plumbing testable without an
applier.

Placeholder sink for callers that have not wired a real ``OutputApplier``.

Status-plane messages (ready/stats/control/dead) are no-ops here —
the :class:`EngineClient` already caches them before forwarding. A
data-plane message (``OutputBatchMsg``) with no applier is a wiring
bug: tokens would be silently dropped, so it raises.

The production data-plane sink: feeds ``OutputApplier``.

``on_batch`` is exactly ``applier.apply`` — the process marshaler's
consumer half is byte-for-byte the thread/inline modes' applier
(design §3: identical code in all three modes). The applier's
``seq``-gap check and unknown-request tolerance apply unchanged.

Status-plane events take optional callbacks (all run on the HTTP
loop): production wiring connects ``on_control`` to the
tokenizer-pool reload (``model_switched``) and ``on_dead`` to the
SSE error-frame fan-out; unwired they are logged no-ops — the
:class:`EngineClient` has already cached / latched the message
before the sink sees it, so nothing is lost.

Engine-side utility failure, re-raised at the API-side call site.

Carries the failure structurally so the admin route layer can map it
to the same HTTP status its in-proc ``except`` clauses produce:
``engine_error`` is the engine-side ``str(exc)`` verbatim
(HTTP ``detail`` byte-identity), ``error_types`` the exception class
names in MRO order (concrete first) from
:func:`~arbi_serve.engine.proc.admin_registry.exception_type_names`.

API-child handle to the engine process (transport-opaque).

``submit_fn`` is the intake send path (``api_child`` wires it to the
ROUTER outbox); it must be callable from the HTTP loop and cheap.
``sink`` is the applier seam (default :class:`NullMessageSink`).

One admin/control RPC round-trip: returns the decoded result.

Raises :class:`UtilityCallError` on an engine-side failure (the
registry's loud unknown-method / bad-args errors surface here),
:class:`EngineDeadError` if the engine dies mid-call, and
``asyncio.TimeoutError`` past ``timeout_s``.

Route one decoded output message (HTTP-loop side, §2).

Updates the caches, resolves RPC futures, forwards data-plane
messages to the sink. Unknown types raise — a message class that
reaches here without a route is a contract bug, not a skip.

Latch engine death (§7): fail pending RPCs, notify the sink.

Called with the decoded ``EngineDeadMsg``, with the synthesized
equivalent for the raw sentinel frame, or by the PID/heartbeat
watcher on a hard kill. Idempotent (first notice wins).

Msgpack codec for the engine <-> API-child wire messages.

Design (``docs/engine_core_process.md`` §5)
-------------------------------------------
One :class:`MsgpackEncoder` per sending IO thread, one typed
:class:`MsgpackDecoder` per top-level message class on the receiving
side (the type-byte frame selects the decoder; see
:mod:`arbi_serve.engine.proc.messages`). Neither is thread-safe — each
IO thread owns its own instance.

Zero-copy multi-frame layout: ``encode()`` returns a list of buffers.
``bufs[0]`` is the msgpack blob; each tensor/ndarray at or above
``size_threshold`` contributes its raw backing store as an extra buffer,
referenced from the blob by index — the payload bytes are never copied
into the msgpack stream, and ZMQ sends each buffer as its own frame
(``copy=False``). Small arrays inline via the RAW_VIEW ext type instead,
so a tiny payload doesn't pay per-frame overhead. Only
:class:`~arbi_serve.engine.proc.messages.MediaRef` carries tensors — the
per-step token path is text + ints and encodes to a single frame.

Security contract: pickle never touches the wire unless the dev-only
``ARBI_ALLOW_INSECURE_SERIALIZATION`` flag is set (pickle deserializes
attacker-controlled bytes into arbitrary code execution). Default
behaviour for a type outside the message contract is a loud
``TypeError`` at encode time and a loud error at decode time — a
foreign object on this boundary is a bug in the emitting code, not
something to smuggle through.

Typed msgpack decoder — one instance per top-level message class.

Not thread-safe: ``aux_buffers`` is instance state shared with the
dec/ext hooks for the duration of one :meth:`decode`.

``share_mem=True`` (default) decodes arrays as views into the
received frames — valid only while the caller keeps those frames
alive; pass ``share_mem=False`` for owned copies.

Engine-side bindings: the seams ``run_engine_process`` injects.

This module is the one place the process transport meets the engine
loop (``docs/engine_core_process.md`` §1/§3/§6/§7):

* :class:`EngineBindings` — the ``intake_handler`` / ``utility_handler``
  / ``stats_provider`` seams of
  :class:`~arbi_serve.engine.proc.runner.EngineProcRuntime`, plus the §7
  ``stop_intake`` / ``drain`` / ``final_flush`` hooks and the
  :meth:`EngineBindings.run` loop driver. Intake messages are decoded on
  the runtime's IO thread and marshaled onto the engine loop
  (``call_soon_threadsafe`` — this module is transport-boundary
  allowlisted); the RequestFactory, admin registry, cancel and wakeup
  all execute there, exactly like the thread mode's coalesced drain.
* :func:`run_single_rank_engine_process` /
  :func:`run_distributed_engine_process` — the bootstrap entries binding
  :func:`~arbi_serve.engine.proc.supervisor.run_engine_process` to a
  real engine (single rank) or the torchrun rank-0
  ``DistributedEngineDriver``.

Per-flow wiring (the request-factory by-reference contract):

=================  ==================================================
in-proc payload    process-mode replacement (here)
=================  ==================================================
``client``         ``None`` — the consumer half lives in the API child
``features``       decoded from ``SubmitMsg.media``
                   (:func:`~arbi_serve.engine.proc.media.media_refs_to_features`)
``tenant_ctx``     ``None`` — factory reconstructs from ``msg.tenant``
=================  ==================================================

The run loop drives ``run_forever`` on this process's main-thread
asyncio loop — no ``EngineThread`` in process mode: the process is the
engine (§1). The OutputBus is switched to its process marshaler
(``bus.bind_process``): each per-step flush hands one finalized
``OutputBatchMsg`` to ``runtime.emit`` with a piggybacked ``StatsMsg``
(§5.2), preserving the empty-batch no-op.

OTEL for the engine process (no FastAPI here).

The engine-process-local gauges (weight attribution, named pools,
allocator walks — the ``_bind.py`` absent-in-child set) exist iff
this process configures its own exporter; the child exports the
StatsMsg-sourced subset + HTTP spans separately.

``ARBI_ENGINE_PROC=1``, torchrun rank 0 (bootstrap entry).

Rank 0 stays the torchrun engine primary (NCCL head, §1) and spawns
the API child instead of hosting uvicorn. The API child is spawned
first (inside ``run_engine_process``, before any CUDA init); the
collective ``driver.build()`` then rendezvouses with the worker
ranks exactly as the in-proc path does. ``driver.run_forever``
drives the per-step plan broadcast as the run loop; the seams bind
to ``driver.engine``.

SubmitMsg → RequestFactory → SubmitAckMsg (engine loop).

Mirrors the thread mode's coalesced intake closure
(``submission._asubmit_marshaled``): the common factory has
zero awaits, so it is driven to completion synchronously inside
this callback (publish ordering == arrival ordering); a request
the factory will await for (grammar compile; a pending
stable-VA resident switch) wraps in an engine-loop task and its
ack resolves when the await finishes.

UtilityCallMsg → registry → UtilityResultMsg (engine loop, §6).

``run_utility_call`` is the canonical executor: it never raises —
every failure crosses as ``ok=False`` + ``error_types`` (the
API-side route layer maps those to the corresponding HTTP statuses).

Re-arm a respawned API child (supervisor's respawn hook).

The fresh child's ReadyMsg/StatsMsg caches are empty — without a
re-emitted ReadyMsg it would 503 ``build_in_progress`` forever
(ReadyMsg is otherwise boot-once). Requests that were in flight
through the dead child are orphaned client-side (their
ClientRequests died with it); the engine keeps decoding them to
completion and the new child's applier logs-and-drops their
batches — bounded by their max_tokens, never a leak (the engine
frees on finish as always).

§7 step 2: drain in-flight requests on the engine loop.

Reuses the in-proc drain body (``server.lifecycle.drain_and_exit``
with a non-exiting hook): sets ``_terminating`` + ``_draining``,
waits ≤ ``grace_s``, force-cancels the remainder. Returns True
when fully drained.

Drive the engine on this thread's asyncio loop (§1 step 3-4).

``obj`` is the built engine (single rank) or the
``DistributedEngineDriver`` (rank 0); ``runtime`` is already
started (Hello sent). Order: bind loop + process marshaler →
start ``run_forever`` → MTP boot sweep (needs the live loop) →
emit ReadyMsg (the child flips ready and admits traffic) → kick
the background cache warms → await the loop.

The OutputBus process marshaler: one emit per step (§3).

Piggybacks the status snapshot on every non-empty step flush
(§5.2) so a busy engine never needs a separate heartbeat send —
the runtime's idle timer covers the quiet stretches.

Emit the residents/groups snapshot control event.

Registered on ``Engine._residency_changed_hooks`` — runs on the
engine loop at every resident-set / active-key / group-registry
mutation so the API child's cached routing view
(:class:`~arbi_serve.engine.proc.api_engine.ProcEngine`) follows
engine truth. Same payload shape as the ``ReadyMsg`` trailing
fields (one builder — ``residency_snapshot``).

Tokenizer fingerprint for the engine <-> API-child boot handshake.

Design (``docs/engine_core_process.md`` §4/§5.2): the API child loads its
own tokenizer from the model dir while the engine keeps one for detok.
Both instances must be byte-equivalent — a divergent pair silently
encodes prompts with one vocabulary and detokenizes with another (the
cross-arch hot-swap footgun, now split across two processes). The
``ReadyMsg.tokenizer_fingerprint`` carries this hash; the child compares
it against its own model dir at ready time and treats a mismatch as a
fatal boot error (never a warning — mismatched vocabularies produce
garbage, not degraded output).

The fingerprint is a sha256 over the sorted ``(filename, file-sha256)``
pairs of every tokenizer-defining file present in the model dir
(:data:`TOKENIZER_FILES` — vocab, merges, config, special tokens, chat
template). Content-addressed and order-independent, so it is stable
across hosts / mtimes / path prefixes; two dirs agree iff the files that
define tokenization agree.

Stable content hash of ``model_dir``'s tokenizer-defining files.

Returns ``sha256:<hex>``. Raises ``FileNotFoundError`` when none of
the tokenizer files exist — fingerprinting a dir with no tokenizer
at all is a caller bug (wrong path), not an empty fingerprint.

Multimodal features-dict <-> :class:`MediaRef` converter seam.

The engine ``Request`` carries preprocessed media as a
``dict[modality_name, features]`` (:class:`MultiModalFeatures` /
:class:`AudioFeatures` — see :mod:`arbi_serve.multimodal.inputs`), while
the wire :class:`~arbi_serve.engine.proc.messages.SubmitMsg` carries a
``list[MediaRef]``. This module is the one conversion point:

* thread / inline modes never convert — the features dict crosses by
  reference alongside the ``SubmitMsg`` (like the pre-made
  ``ClientRequest``; see :mod:`arbi_serve.engine.request_factory`), so
  realtime multimodal keeps its zero-copy in-proc path.
* process mode encodes the ``MediaRef`` list through the codec's
  tensor ext type (zero-copy extra frames) and reconstructs the features
  dict engine-side via :func:`media_refs_to_features`.

Encoding scheme (round-trip exact):

Each tensor field of a features dataclass becomes one ``MediaRef`` whose
``kind`` is the modality name and whose ``meta["field"]`` names the
dataclass field. Small non-tensor payloads (the audio int lists) ride as
JSON strings in the same ref's ``meta``. Unknown modalities or fields
fail loudly in both directions — a silently dropped media payload would
surface as a model attending to garbage.

Typed wire messages for the engine <-> API-child process boundary.

Design (``docs/engine_core_process.md`` §5)
-------------------------------------------
Every crossing between the engine process and its API child is one of
the ``msgspec.Struct`` messages defined here — nothing else goes on the
wire. The engine emits **data** (``TokenOut`` / ``FinishOut`` / ...),
never bound consumer callbacks, so the same message crosses by reference
in thread mode (:mod:`arbi_serve.engine.output_bus`) and serialized in
process mode (§3, "one typed pipeline, three marshalers").

Struct config — ``array_like=True`` (positional msgpack arrays, no field
names on the wire), ``omit_defaults=True`` (applies to any object-like
rendering of these structs, e.g. debug JSON dumps; positional msgpack
always carries every field), ``gc=False`` (messages never form reference
cycles; skipping GC tracking keeps the per-step allocation cost flat).
Because ``array_like`` encodes fields positionally, field order in these
classes is wire format: append new fields with defaults at the end of a
struct, never reorder or insert — decoders fill missing trailing fields
from defaults, which is what makes appended fields compatible with
in-flight peers.

What deliberately does not cross (§5.3): weights, KV state, cudagraph
handles, pool objects, the engine ``Request``, asyncio primitives, live
config objects. Tensors appear only inside :class:`MediaRef` (multimodal
feature payloads) — the high-frequency token path is text + ints only.

Dispatch: the first ZMQ frame of every message is a one-byte type tag
(``MSG_TYPE_BYTES`` / :func:`msg_class_for`) so the receiving IO thread
picks its typed decoder without inspecting the payload (§2).

Base for every struct that crosses the engine<->API boundary.

Exists so the (array_like / omit_defaults / gc) wire config is
declared exactly once — a message type that forgot one of them would
silently produce an incompatible encoding.

Constrained-decoding spec, mirroring ``request.ResponseFormat``.

Compiled engine-side (§4: the xgrammar matcher gates sampling and
cannot cross), so the wire carries the *source* spec — type tag plus
exactly one populated payload field for the non-``text`` types.

One multimodal input payload on a :class:`SubmitMsg`.

The only message that may carry a tensor / ndarray: pixel values or
precomputed features ride the codec's ext-type path (zero-copy extra
frames; see :mod:`arbi_serve.engine.proc.codec`). ``uri`` covers the
fetch-engine-side alternative where only a reference crosses.

Full explicit wire mirror of ``request.SamplingParams``.

Every field of the engine dataclass is carried explicitly — the
engine-side ``RequestFactory`` must be able to reconstruct a
``SamplingParams`` without consulting API-side state, and a field
silently dropped here would surface as a request that ignores its
own knob. Defaults are byte-for-byte the engine defaults so
``omit_defaults`` keeps the common chat request tiny.

Narrowings vs the dataclass (documented, not silent):

* ``response_format`` crosses as :class:`GrammarSpec` (same fields).
* ``tool_choice`` is typed ``str | dict | None`` — the dataclass says
  ``Any`` but the only legal values are ``"auto"`` / ``"none"`` /
  ``"required"`` / an OpenAI function-selector dict.
* ``priority`` here is the string scheduling class of
  ``SamplingParams.priority`` ("interactive" / "batch");
  :class:`SubmitMsg.priority` is the separate integer knob from the
  design's intake contract.

One request submission. Tokenization already happened API-side
(§4) — the engine receives ids, never prompt text.

``request_id`` is allocated API-side: the client half
(``ClientRequest``) is created + registered with the OutputApplier
before this message is sent, so the engine-side ``RequestFactory``
uses the incoming id verbatim and never allocates one.

``kind`` is the single source of truth for the task mode
(``generate`` / ``embed`` / ``rerank``). ``SamplingParams.task`` from ``kind``.

Admin/control RPC (§6): ``call_id`` matches the eventual
:class:`UtilityResultMsg`; ``method`` is looked up in the explicit
``ADMIN_METHODS`` registry (never ``getattr``); args are a JSON blob
because admin payloads are heterogeneous and cold-path.

First message on the intake DEALER (§2): teaches the API child's
ROUTER the engine's identity. There is no reply — endpoints were
passed at spawn, so this is readiness, not discovery.

Sampled-token logprob + top-N alternatives for one step.

Parallel arrays (not a list of pairs) keep the msgpack encoding flat;
``top_texts`` is detokenized engine-side like every other text on
this boundary.

Engine status snapshot (§5.2) — piggybacked on every non-empty
step flush and sent standalone at >=1 Hz when idle, doubling as the
heartbeat. The API child's ``/health`` and metrics gauges read the
cached latest snapshot, never engine memory (§6).

Sent once after ``abuild()`` succeeds; the API child flips
``/health`` to ready and admits traffic on receipt. Carries every
model-derived fact the HTTP layer needs so it never reads engine
memory.

Engine → API acknowledgement of one :class:`SubmitMsg`.

The wire form of the SubmitAck: the engine-side RequestFactory
resolves every submission exactly once — ``ok=True`` when the
request published (or was dropped by a racing cancel tombstone, in
which case the ``cancelled`` FinishOut still follows), ``ok=False``
with the HTTP ``status_code`` + message for every
:class:`~arbi_serve.engine.client_request.SubmitRejected`
pre-publish validation failure. The API child resolves the pending
:class:`~arbi_serve.engine.loop_bridge.SubmitAckRegistry` future
from it — the same future thread/inline modes resolve in-proc.

Engine-side process runtime: the ZMQ IO threads + outbox + heartbeat.

Design (``docs/engine_core_process.md`` §2, §3, §7)
---------------------------------------------------
:class:`EngineProcRuntime` is the engine process's half of the boundary:

  * **input IO thread** — DEALER connect to the intake endpoint, speaks
    first (``HelloMsg``, §2), then loops ``recv_typed`` dispatching by
    type byte to the injected ``intake_handler`` (``SubmitMsg`` /
    ``CancelMsg`` / ``WakeupMsg``) or ``utility_handler``
    (``UtilityCallMsg`` — the registry consumer). The handlers are
    seams wired to the engine's intake-coalescing path; this
    module knows nothing about the engine loop.
  * **output IO thread** — PUSH connect to the output endpoint, drains a
    plain ``deque`` + ``threading.Event`` outbox (§2) fed by
    :meth:`emit` — the ZMQ marshaler seam the engine-side per-step flush
    hands ``OutputBatchMsg`` / ``StatsMsg`` / ``ReadyMsg`` /
    ``ControlEventMsg`` to. Serialization happens on this thread, so it
    overlaps GPU work.
  * **heartbeat** — the output thread's outbox wait doubles as the >=1 Hz
    idle timer (§5.2): when nothing carrying stats has been sent within
    ``heartbeat_interval_s``, it emits a ``StatsMsg`` built by the
    injected ``stats_provider()`` so the API child's staleness detector
    never trips on a merely-idle engine.
  * **death** — :meth:`send_engine_dead` is the last-gasp path (§7):
    typed ``EngineDeadMsg`` when the encoder still works, the raw
    ``ENGINE_DEAD_SENTINEL`` frame when it may not.

Thread ownership: each ZMQ socket is touched by exactly one thread. The
DEALER belongs to the input thread (its rare engine->API sends — Hello /
re-Hello after a child respawn — ride an intake outbox drained by that
same thread). The PUSH belongs to the output thread while it runs;
:meth:`stop` / :meth:`send_engine_dead` join it before touching the
socket, so ownership hands over, never overlaps.

The engine process's IO runtime for one API child connection.

Owns the two boundary sockets + their IO threads. The engine loop
interacts through exactly three seams:

  * ``intake_handler(msg)`` — called on the input IO thread with each
    decoded ``SubmitMsg`` / ``CancelMsg`` / ``WakeupMsg``; must be
    cheap + thread-safe (publishes onto the engine loop via the
    intake-coalescing path).
  * ``utility_handler(msg)`` — called on the input IO thread with each
    ``UtilityCallMsg``. Production wiring: marshal onto the
    engine loop and ``await``
    :func:`arbi_serve.engine.proc.admin_registry.dispatch_utility`
    (a coroutine — it must run on the engine loop so ``CriticalSection``
    semantics hold, §6), then :meth:`emit` the ``UtilityResultMsg``
    (ok → ``json_safe`` result bytes; ``AdminRpcError`` / method
    exception → ``ok=False, error=str(exc)``).
  * ``stats_provider()`` — called on the output IO thread to build the
    idle-heartbeat ``StatsMsg``. Production wiring: ``lambda:
    build_stats_msg(eng)``
    (:func:`arbi_serve.engine.proc.stats.build_stats_msg` is
    defensive-by-contract — tolerates a mid-build engine and never
    raises, exactly what an off-loop heartbeat needs).

plus :meth:`emit` (any thread; typically the engine loop's per-step
flush). :meth:`emit` is the process marshaler of
``engine/output_bus.py`` §3: it accepts exactly what
``OutputBatchBuilder.take()`` finalizes (an ``OutputBatchMsg`` with
builder-owned ``seq``), wired as ``bus.flush`` handing the
batch here instead of the thread-mode ``call_soon_threadsafe``.
Lifecycle: :meth:`start` / :meth:`stop` / :meth:`send_engine_dead`.

Stop both IO threads and close the sockets (linger-bounded).

Idempotent. The outbox is drained by the output thread before it
exits, so messages emitted before stop() still ship (LINGER
bounds how long a dead peer can hold the close).

Queue a fresh ``HelloMsg`` (API-child respawn re-identification).

A respawned child binds a brand-new ROUTER that has never seen
this DEALER's identity; the supervisor calls this after a respawn
so the child can route intake-direction traffic again.

Queue one output message (the ZMQ marshaler seam, §3).

Called from the engine loop's per-step flush with
``OutputBatchMsg`` / ``StatsMsg`` / ``ReadyMsg`` /
``ControlEventMsg`` / ``UtilityResultMsg``. Cheap: one deque
append + event set; serialization happens on the output thread.
Fail-loud on a non-output message type or a stopped runtime — an
emit after stop would silently drop output.

Last-gasp death notice (§7). Safe to call in any runtime state.

Queues ``EngineDeadMsg`` behind any already-emitted output (the
clean-SIGTERM path must flush final batches first), stops the IO
threads, and — if the typed send could not be confirmed (thread
already dead, encoder failure) — fires the raw
``ENGINE_DEAD_SENTINEL`` frame directly on the PUSH socket, whose
ownership has been handed over by the join in :meth:`stop`.
Idempotent: only the first call sends.

ZMQ socket + framing helpers for the engine <-> API-child boundary.

Design (``docs/engine_core_process.md`` §2)
-------------------------------------------
Two ``ipc://`` endpoints in a run-scoped 0700 tempdir, both bound by the
API child (the stable side — it exists before the engine finishes
building), engine connects:

  * **intake** — API child ROUTER, engine DEALER. Carries ``SubmitMsg``
    / ``CancelMsg`` / ``UtilityCallMsg`` / ``WakeupMsg`` toward the
    engine; the engine's DEALER speaks first (``HelloMsg``) so the
    ROUTER learns its identity.
  * **output** — API child PULL, engine PUSH. Carries
    ``OutputBatchMsg`` / ``StatsMsg`` / ``ReadyMsg`` /
    ``UtilityResultMsg`` / ``ControlEventMsg`` / ``EngineDeadMsg``.

Wire framing: every message is a multipart ``(type_byte, *codec_frames)``
— the one-byte type tag (``messages.MSG_TYPE_BYTES``) rides its own
frame so the receiving IO thread dispatches to a typed decoder without
touching the payload. ROUTER additionally sees the DEALER identity as
frame 0; PULL/DEALER receivers see the type byte first.

Zero-copy send safety: ``send_multipart(copy=False)`` queues pointers
into the caller's buffers — mutating or freeing a buffer before ZMQ's
IO thread ships it corrupts the message. :func:`send_typed` therefore
follows the vLLM pattern: large sends go out ``copy=False, track=True``
and the ``(tracker, buffers)`` pair is parked on the caller's pending
deque until the tracker reports delivery; small sends just copy (the
copy is cheaper than the tracking).

TCP is deliberately unsupported: the API child is always same-host
(spawned by the engine), so a ``tcp://`` endpoint here is a config bug.

Create, configure, and bind/connect one boundary socket.

``hwm=0`` (unlimited) on both directions is the correct default for
this boundary: the peers are same-host and the producers are
naturally paced (one output batch per engine step; intake paced by
the API-side admission semaphore), so an HWM drop/block would only
convert a transient stall into silent message loss or an engine-loop
stall. Only ``ipc://`` / ``inproc://`` endpoints are accepted (§2:
no multi-node API tier).

The run-scoped IPC rendezvous: one private tempdir, two endpoints.

Allocated by the engine (CLI) before spawning the API child; both
endpoint strings are passed to the child at spawn, so there is no
discovery step. ``cleanup()`` is idempotent and safe while sockets
are open (unlink semantics) — the engine calls it on exit.

Allocate the run-scoped IPC directory + the two endpoint paths.

``mkdtemp`` creates the directory 0700 (owner-only) — the ipc socket
files inherit that confinement, so another local user cannot connect
to (or squat on) the engine's intake. The mode is re-asserted rather
than trusted, since a permissive umask override here would silently
open the engine's admin surface to every local user.

Arrange for ``sig`` to be delivered when this process's parent dies.

The API child calls this first thing after spawn so an engine
hard-kill can never orphan a listening HTTP socket — the kernel
SIGKILLs the child the instant its parent exits, no cooperation
required.

Linux-only (``prctl(PR_SET_PDEATHSIG)``); returns ``True`` when armed,
``False`` on non-Linux platforms where the engine-PID liveness watcher
(``api_child``) is the only orphan guard. A failed prctl on Linux
raises — a child that silently believes it is death-guarded when it
is not would be exactly the kind of silent fallback this repo bans.

Caveat (inherent to pdeathsig): the setting fires on the death of the
spawning thread's task, and is cleared across ``execve`` of a setuid
binary — neither applies to our mp-spawn child, which calls this from
its main thread and never execs.

Send one ``(type_byte, *buffers)`` multipart message.

``buffers`` is the codec's ``encode()`` output (blob + aux frames).
Small payloads are sent ``copy=True``. Large payloads require a
``pending`` deque: they go out ``copy=False, track=True`` and the
buffers are parked on the deque until delivered — dropping the last
reference to a zero-copy buffer before delivery is use-after-free,
so refusing to send large without a guard is the safe default.
``identity`` prefixes the DEALER address frame when sending from a
ROUTER.

Receive one message from a PULL / DEALER socket.

Returns ``(type_byte, payload_frames)`` with frames received
``copy=False`` — feed ``payload_frames`` straight to
``MsgpackDecoder.decode`` (zmq.Frame satisfies the buffer protocol).
The one raw, non-typed frame on the wire —
``messages.ENGINE_DEAD_SENTINEL`` — is returned verbatim as the
first element; match it before type-byte dispatch.

Engine status snapshot builder — :func:`build_stats_msg`.

Design (``docs/engine_core_process.md`` §5.2 / §6)
--------------------------------------------------
The engine piggybacks a :class:`~arbi_serve.engine.proc.messages.StatsMsg`
on every non-empty step flush and sends one standalone at >=1 Hz when
idle (the heartbeat). The API child's ``/health`` + ``/health/ready`` +
OTEL gauges read the cached latest snapshot — never engine memory — so
this builder is the single place the health/readiness reach-ins
(``server/health.py``, ``engine/infra_health.py``) are flattened to data.

Fault-latch semantics: the three reasons ship raw (sticky fatal, sticky
swap, self-healing infra + its ``ts``/``count``); the API child applies
the ``INFRA_UNHEALTHY_WINDOW_S`` decay itself from ``infra_failure_ts``,
mirroring ``engine_unhealthy_reason`` precedence (fatal → swap → infra).

Defensive by contract: the builder runs from the heartbeat while the
engine may be mid-build (pool/scheduler/model_runner still ``None``,
attributes not yet assigned). Every read tolerates a missing/None
attribute and yields a valid booting-state message — the builder never
raises into the engine loop.

Engine-side API-child management: spawn, monitor, respawn, SIGTERM.

Design (``docs/engine_core_process.md`` §1, §7)
-----------------------------------------------
The engine is the primary process; the API/HTTP server is its child
(§1 — arbi's TP is torchrun-SPMD, so the engine owns the process).
This module owns that child:

  * :func:`spawn_api_child` — ``multiprocessing.get_context("spawn")``
    (the child must never inherit a forked CUDA context), passing the
    pickled ``ServerConfig`` + IPC endpoints + the engine's PID.
  * :class:`ApiChildSupervisor` — a monitor thread on the child's
    ``Process.sentinel``; on unexpected exit, respawn with the same
    endpoints, bounded by :data:`MAX_RESPAWNS` in
    :data:`RESPAWN_WINDOW_S` (§7: "3 respawns / 10 min, then fatal").
    ``atexit`` + ``weakref.finalize`` kill-child guards on both normal
    and GC'd teardown so a dead engine can never leave a zombie child
    squatting on the HTTP port.
  * :func:`run_sigterm_sequence` — the §7 shutdown order as one
    auditable function: stop intake → drain ≤ grace_s → final flush →
    ``EngineDeadMsg(clean=True)`` → SIGTERM child → join → return.
  * :func:`run_engine_process` — the process-mode serve tail skeleton
    (bootstrap calls it when ``ARBI_ENGINE_PROC`` is on). Everything
    transport-side is implemented; the engine-facing seams
    (``intake_handler`` / ``utility_handler`` / ``stats_provider`` /
    ``build_engine`` / ``run_loop``) are injected by the caller.

Default fatal action: log critical and exit nonzero.

``os._exit`` (not ``sys.exit``): this runs on the monitor thread,
where an exception would die silently instead of stopping the
engine; the orchestrator's restart policy is the recovery path (§7).

Owns the API child process: spawn, sentinel monitor, respawn.

``spawn_target`` is ``api_child.main`` in production (injectable for
tests, which substitute a torch-free stub). ``on_respawn`` is called
after every successful respawn (the engine runner re-sends
``HelloMsg`` there — the fresh child's ROUTER has never seen the
DEALER identity). ``on_fatal(reason)`` runs when the respawn budget
is exhausted (default: exit the engine process nonzero).

Execute the §7 SIGTERM order, engine-side. Blocking; call once.

1. ``stop_intake()`` — reject new ``SubmitMsg`` with a shutting-down
   error (engine-side seam bound to the intake path).
2. ``drain(grace_s)`` — wait for in-flight requests, ≤ ``grace_s``;
   returns True when fully drained (False = grace expired and the
   remainder was force-cancelled — logged, not hidden).
3. ``final_flush()`` — emit the last ``OutputBatchMsg``es through
   ``runtime.emit`` so no committed token is lost.
4. ``runtime.send_engine_dead(clean=True)`` — the child flips 503
   and stops accepting; this also flushes + stops the runtime's IO.
5. SIGTERM the child, join it (supervisor stops respawning first).

The caller exits the process after this returns.

Process-mode serve tail (§1 sequence): the ``ARBI_ENGINE_PROC`` path.

Order (design §1): allocate the run-scoped IPC endpoints → spawn the
API child first (before any CUDA init — it serves ``/health`` =
booting while the engine builds) → build the engine in this process's
main thread → start the runtime (Hello) → the engine's ReadyMsg goes
out through the run-loop wiring → run the loop → §7 sequencing on
exit.

The seams (this function is the only place they meet):

  * ``build_engine(cfg) -> engine`` — construct + ``build()`` the
    engine (single-rank) or driver (multi-rank rank 0).
  * ``make_runtime_seams(engine) -> dict`` — must return
    ``intake_handler`` / ``utility_handler`` / ``stats_provider``
    (the :class:`EngineProcRuntime` kwargs) bound to the engine's
    intake-coalescing path, the admin registry
    (``utility_handler`` marshals onto the engine loop, awaits
    :func:`~arbi_serve.engine.proc.admin_registry.dispatch_utility`,
    emits the ``UtilityResultMsg``), and
    ``lambda: build_stats_msg(eng)``
    (:func:`~arbi_serve.engine.proc.stats.build_stats_msg`); it may
    also return ``on_stop_intake`` / ``on_drain`` /
    ``on_final_flush`` overrides for the §7 sequence and
    ``on_child_respawn`` (re-emit the cached ReadyMsg so a
    respawned child's empty caches flip ready again).
  * ``run_loop(engine, runtime) -> None`` — drive ``run_forever``,
    emitting ``ReadyMsg`` first and per-step batches via
    ``runtime.emit``; returns when the engine stops.

All three are required: refusing to run half-wired is the fail-loud
contract — a missing seam raises before anything is spawned.
Startup-failure path (§7): if ``build_engine`` raises, the already-
serving child gets ``EngineDeadMsg(clean=False, reason=<the error>)``
so it reports the boot error verbatim, and both processes exit
nonzero within the boot timeout.

SIGTERM/SIGINT → the §7 sequence, then exit 0.

Installed with ``signal.signal`` (the engine loop may or may not be
asyncio at this level); the sequence itself runs on a dedicated
thread so a long drain never blocks inside a signal frame.

Profile the activation peak. The peak is reserved as MEASURED — no pad.

Runs the model_runner's owned probe and reconciles the peak across TP ranks
(MAX so the smallest budget wins). Returns ``(peak_activation_bytes,
scratch_bytes, raw_peak_bytes)``; all three are the same measured value, the
triple kept so callers can name the role they mean. Sets
``eng.boot_state.profiled_activation_peak_bytes`` and
``eng.boot_state.single_shape_peak_activation_bytes``.

Engine-side inputs of the ``capture.io_buffers`` prediction.

Everything :func:`~arbi_serve.engine.memory_budget.capture_io_buffer_terms`
needs that is not already a scalar in the caller's frame: the captured
``(B, S)`` shape set, the verify slate width and block-table page count,
the sampler-scratch gates, the pre-warmed ``residual_buf`` rung set, the
DFlash tap width and the EXL3 reconstruct geometry.

Resolved the same way the boot itself resolves them, from the same
functions, so the prediction and the allocation cannot pick different
answers. Each lookup that reads a checkpoint or an attached-later engine
field is best-effort: a term that cannot be resolved is reported as ZERO
and shows up in the freeze reconcile as an under-prediction on this pool,
which is strictly better than a boot that dies inside the budget.

Closed-form per-pool VRAM predictions for the KV budget.

Every entry is the analytic upper bound for its pool. ``spec.drafter_slots``
is seeded at 0 — the DFlash two-pass refines it against the realized
``num_pages`` (see :func:`_bisect_dflash_drafter_slots`).

``graph_pool_is_measured_total`` says which KIND of number
``graph_pool_predicted`` is, because the two do not mean the same thing:

  * ``False`` — the COLD estimate
    (:func:`~arbi_serve.engine.build_graph_pool.cold_boot_graph_pool_upper_bound_bytes`),
    which models the ``capture.cudagraphs`` pool alone. It goes on that row
    and ``capture.io_buffers`` keeps its own analytic prediction.
  * ``True`` — a MEASURED value from the budget cache, which is the SUM of
    both capture pools (``total_cudagraph_bytes`` reads both tags). Handing
    that sum to the ``capture.cudagraphs`` row while the io row also carried
    its analytic prediction booked the io pool TWICE. It is split into its
    two rows instead (:func:`~arbi_serve.engine.build_graph_pool.
    resolve_capture_pool_rows`), from the per-tag reading the cache carries
    when one was available and by subtraction when it was not.

Bisect the DFlash ``spec.drafter_slots`` two-pass fixed point.

Breaks the chicken-and-egg between the drafter's per-ceiling KV cost and
``num_pages`` (what the budget derives). Bisects for the unique fixed point
of ``S = g(f(S))`` where ``S`` is the KV serving ceiling,
``f(S)`` is the drafter's KV footprint at effective context ``S`` — both the
persistent ``spec.drafter_slots`` pool and the num_pages-bounded per-step
context-assemble transient reserved in the serving floor — and
``g(f(S)) = serving_kv_ceiling_pages``. ``g∘f`` is monotone non-increasing so
bisection converges. Folding the transient into ``f`` (not just the slots)
matters: a fixed raw-ceiling transient reserve would disagree with the
converged (lower) ceiling, so the grow couldn't leave the realized
transient headroom and KV would starve to a handful of pages.

``seed_budget`` is the no-drafter run (``spec.drafter_slots = 0``) used for
the ``hi`` bound. ``slots_for_num_pages`` is the slots half of ``f`` (``None``
→ raw-ceiling upper bound), used only to report the final reserved bytes.
``budget_for_ceiling(S)`` re-runs the budget with both the slots pool and the
serving-floor transient sized for ``S``. Returns ``(final_slots, cap_pages,
budget)`` where ``cap_pages`` is the proven-feasible serving ceiling both KV
and the real ``DraftKVSlots`` are clamped to.

Resolve the MTP draft depth with the capture ladder priced first.

Builds the ``non_kv_bytes_at(k)`` cost model the depth resolver needs:
for a candidate depth ``k``, everything that is not the paged-KV pool
and that scales with ``k`` — the cold-boot cudagraph private-pool
reserve (the verify ``(B, k+1)`` + drafter ``(B, k)`` rungs, whose
working sets a captured graph retains for the whole boot), the
depth-scaled pools (``spec.gdn_rollback`` on a GDN-hybrid is the big
one), and the vocab-scale stochastic-verify serving tail.

``cfg.mtp.n_draft`` is pinned around each probe because the shape
derivations (``_drafter_chain_buckets``,
``config_verify_capture_shapes``) read the depth off config — there is
no attached driver to ask this early. The pin is restored before
returning; the resolver makes the real one.

Returns the resolved depth (``0`` when MTP is off).

Profile the activation peak and derive ``num_pages``.

Boot sequence:

  1. Snapshot ``weights_bytes = torch.cuda.memory_allocated`` —
     every weight tensor is on-device by this point.
  2. Probe ``torch.cuda.mem_get_info`` for total VRAM. The usable
     budget is ``total × gpu_memory_utilization``.
  3. Build a transient scratch :class:`MultiStatePool` + per-layer
     attn ops + per-kind metadata builders sized just large enough
     to absorb one ``max_batched_tokens`` prefill batch.
  4. Reset CUDA peak, run :meth:`ModelRunner.forward` on a
     synthetic prefill batch of random ids, synchronize, and
     capture ``max_memory_allocated``.
  5. Subtract ``weights_bytes`` and the scratch pool's slab bytes
     from the peak — the residue is the activation peak.
  6. Drop the scratch state, ``empty_cache``, and verify the
     allocator returns to the pre-profile baseline.
  7. Hand the four numbers to :func:`compute_kv_budget` to derive
     ``num_pages``. Per-page bytes come from the active backend's
     own ``bytes_per_token`` (TKV / bf16 / etc all expose it) —
     we never recompute bit widths here.
  8. Under TP > 1, ``all_reduce(MAX)`` the activation peak so the
     most pessimistic rank wins, then ``all_reduce(MIN)`` the
     resulting ``num_pages`` so every rank's pool size is
     identical (the page-allocator requires it).
  9. Validate any user-supplied ``--num-pages`` override against
     the budget. Override > budget → :class:`MemoryBudgetError`.

Sets ``eng.boot_state.profiled_num_pages`` and ``eng.boot_state.budget`` (the full
:class:`KvBudget` for logging / admin endpoints). The call is a
no-op for non-paged-KV models: MLA / Mamba / GDN sizing is
handled inside their respective pool ctors and is not driven by
``num_pages`` in the same way.

``spec.drafter_slots`` bytes for a given realized ``num_pages``.

The DFlash drafter mirrors a request's committed context in its own
persistent draft-KV slabs, so they must size to the effective
per-request context the realized paged-KV pool serves at the
resolved concurrency — not the raw ``cfg.cache.max_context`` RoPE
ceiling (262k at "auto"). ``num_pages=None`` returns the raw-ceiling
upper bound (the pre-budget seed); a concrete ``num_pages`` narrows
it to ``derive_effective_max_context`` — exactly what
``build_dflash_drafter`` passes the real ``DraftKVSlots``.

``capture.io_buffers`` bytes at a given realized ``num_pages``.

The pool's one num_pages-dependent tenant is the DFlash draft graphs'
shared context buffers: one ``(sweep width, n_kv, layer_len_i,
head_dim)`` K and V pair plus the int64 key positions per draft layer,
i.e. the ``DraftKVSlots`` row geometry re-materialised at the capture
sweep's batch width. Its layer widths come from the effective context
the REALIZED pool serves, so it narrows with the candidate ceiling
exactly like ``spec.drafter_slots`` and must ride the same two-pass —
priced at the raw ``max_context`` bound it would oversubscribe the
budget before the refine could run. Zero off the DFlash capture path.

DFlash draft-step transient bytes bounded by a realized ``num_pages``.

The dominant term (context-assemble copy) mirrors the persistent
``spec.drafter_slots`` geometry, so it narrows with the servable page
count exactly like :func:`_dflash_slots_for_num_pages`. ``None`` returns
the raw max_context upper bound (the pre-budget seed). Zero off the
stochastic / DFlash path.

Plan-time serving activation floor at a given draft-transient size.

The verify-tail / bf16-scratch / tkv-staging / prefill terms are
num_pages-independent; only the DFlash draft transient tracks the
candidate ceiling, so it is the one input the caller varies.

Plan-time serving activation floor with the draft transient bounded to
``num_pages`` — the num_pages-dependent input the drafter two-pass
bisection sweeps so the KV ceiling stays self-consistent with the reserve
left free at grow.

Re-run the budget with both the drafter slots pool and the draft-step
transient reserve sized for the candidate serving ceiling ``num_pages``.

The two-pass fixed point (:func:`_bisect_dflash_drafter_slots`) sweeps
this: at each candidate ceiling the persistent ``spec.drafter_slots``
(:func:`_dflash_slots_for_num_pages`) and the serving-floor draft
transient (:func:`_serving_floor_for_num_pages`) both narrow together, so
the converged ceiling leaves room for the drafter's slabs and its per-step
context-assemble — no post-capture grow KV-starve. ``None`` seeds the
raw-ceiling upper bound.

Member-aware split of profile-time VRAM into (weights, external, foreign).

``baseline`` is the ``(memory_allocated, memory_reserved)`` snapshot taken
right before this build started (``build_member_into_engine`` sets it for
residency-pool member builds; ``None``/zeros for a solo boot). The split:

  * ``weights``  = this build's own torch allocations
    (``alloc_now − base_alloc``). A parked member's allocations stay
    counted by ``memory_allocated`` (live tensors behind stable VAs even
    when cuMem-unmapped) — they must not be booked as this member's
    weights.
  * ``external`` = physical THIS PROCESS holds that its own caching-allocator
    reservation does not account for: the CUDA primary context, the driver's
    cubin/kernel-stack residency, the regions we cuMem-map directly, and (pool
    mode) the other members' still-mapped physical — exactly what stays
    resident while this member serves.
  * ``foreign``  = the rest of the card: VRAM OTHER PROCESSES hold. Real,
    resident, and not ours to reclaim, so the budget still books it — but as
    its own ``driver.foreign_process`` row, never inside our context term.

``own_used_bytes`` is this process's own residency per NVML's per-process
walk. ``None`` means the walk could not identify us, and then the card
cannot be split at all: ``external`` falls back to the device-wide
remainder (the conservative reading — it holds out every byte on the card,
so the budget stays safe) and ``foreign`` is 0 because we cannot name a
single byte as another process's. The caller says so; silence there is the
failure this returns a distinguishable shape to prevent.

With ``baseline=None`` and an unavailable per-process reading this reduces
to the single-model arithmetic (``weights = memory_allocated``, ``external
= (total − free) − memory_reserved``).

Measure profile-time VRAM and split it into weights / external / foreign.

Syncs the device, reads ``mem_get_info`` / ``memory_allocated`` /
``memory_reserved`` / NVML's per-process walk, and applies the member-aware
split via :func:`attribute_profile_memory`. Returns ``(total_bytes,
raw_alloc_bytes, weights_bytes, external_reserved_bytes,
foreign_process_bytes)`` where ``raw_alloc_bytes`` is the unattributed
allocator snapshot the activation probe subtracts from its peak.

GPU-time profiling for the decode loop — ``ARBI_PROFILE`` mode selector.

cProfile is the wrong lens for arbi-serve's problem: decode is
host/launch-bound (Self CPU >> Self CUDA) and cProfile is blind to GPU
time. This module is the GPU-aware replacement, gated behind the
``ARBI_PROFILE`` flag (``off`` | ``cprofile`` | ``torch`` | ``nvtx``):

  * ``off`` (default) — :func:`make_profiler` returns ``None`` and the
    NVTX range helpers are no-ops. Zero runtime overhead: no profiler
    object is built, no NVTX push/pop reaches the driver.

  * ``torch`` — :class:`TorchProfilerController` wraps the decode loop in
    ``torch.profiler.profile(activities=[CPU, CUDA], schedule=...,
    on_trace_ready=tensorboard_trace_handler(dir))``. ``prof.step()`` is
    called once per decode step. One Kineto trace is written per rank
    (the handler stamps the worker name) — HTA-loadable.

  * ``nvtx`` — :class:`NvtxProfilerController` emits
    ``torch.cuda.nvtx.range_push/pop`` ranges around the decode step and
    the MTP sub-phases, plus ``cudaProfilerStart()/Stop()`` around the
    ``[ARBI_PROFILE_START, ARBI_PROFILE_STOP)`` iteration window. ~zero
    GPU overhead; wrap externally with ``nsys ... -c cudaProfilerApi``.

  * ``cprofile`` — :class:`CProfileController`. Host/Python-only; kept
    for parity with the prior cProfile-only flag, but it cannot see GPU
    time. Use ``torch`` / ``nvtx`` for the real (launch-bound) picture.

The controller object is created once at engine boot (:func:`make_profiler`,
called from the run loop) and threaded through the per-step path. The
``nvtx_range`` context manager is the hot-path seam the run loop and
``spec_decode/mtp.py`` use to bracket sub-phases — it is a near-free
no-op when the active mode is not ``nvtx``.

Live control: :func:`start_torch_profile` / :func:`stop_torch_profile`
let the ``/start_profile`` + ``/stop_profile`` HTTP endpoints flip a
torch profiler on the running server without a restart. Under TP>1 the
driver broadcasts the start/stop so each rank dumps its own trace.

Create + verify ``profile_dir`` is writable; fail loud otherwise.

Per AGENTS.md no-silent-fallbacks: an enable-time misconfiguration
(read-only path, file-where-a-dir-should-be) must raise here, not
degrade to dropping the trace silently at ``on_trace_ready`` time.
Returns the absolute path on success.

The per-step profiling seam the run loop drives.

Implementations are mode-specific. The run loop calls
:meth:`enter` once before the loop, :meth:`step` once per decode
step, and :meth:`exit` on teardown. :meth:`nvtx_range` brackets a
named sub-phase (only the nvtx controller emits ranges; the others
yield a no-op context).

``torch.profiler`` over the decode loop — one Kineto trace per rank.

The schedule lands the ``active`` window past cudagraph capture in
warm steady-state (``skip_first`` ~ 50). ``repeat=1`` so the loop is
profiled exactly once. ``on_trace_ready`` writes a
``tensorboard_trace_handler`` dump tagged with ``worker_name`` =
``rank{N}`` so each rank's trace is a distinct, HTA-loadable file.

NVTX ranges + ``cudaProfilerStart/Stop`` over an iteration window.

Meant to run under ``nsys profile ... -c cudaProfilerApi
--cuda-graph-trace=node --trace=cuda,nvtx,nccl``. ``cudaProfilerStart``
fires at the ``start``-th decode step and ``cudaProfilerStop`` at the
``stop``-th, so nsys captures only the warm steady-state window. The
NVTX ranges (decode_step / drafter_step / verify / all_reduce) are
always emitted while active so they show up in the nsys timeline.

Build the per-engine :class:`ProfilerController` for ``ARBI_PROFILE``.

Returns ``None`` when the mode is ``off`` (the production path) — the
run loop then skips every profiler call and there is zero overhead.
Fails loud (via :func:`ensure_profile_dir_writable`) when a trace-
writing mode is selected but the dir is not writable.

Start a live torch profiler on ``eng`` (idempotent-ish; raises if running).

Reads ``ARBI_PROFILE_DIR`` / ``ARBI_PROFILE_WITH_STACK`` /
``ARBI_PROFILE_RECORD_SHAPES`` for the trace config. Fails loud if the
dir isn't writable. Stores the live ``torch.profiler.profile`` on
``eng._live_torch_profiler``.

Reasoning-marker parser registry + model auto-detection.

The active parser is **model-tied**: resolved at engine build from the
``--reasoning-parser`` config (default ``auto`` → sniff the chat template)
and stored on the engine, so a hotswap to a different model re-resolves it
through ``eng.build()``. Mirrors :mod:`arbi_serve.engine.tool_parsers`.

A model whose chat template carries no reasoning markers resolves to
``None`` — the split is skipped entirely and ``content`` behaves exactly
as before.

Return the :class:`ReasoningParser` for this model, or ``None``.

``override`` (the ``--reasoning-parser`` value) wins when it names a
registered parser; ``"none"`` disables the split outright. ``auto`` /
``None`` sniff the chat template for each registered family's marker
pair — the template is the model's own declaration of the framing it
was trained to emit. No marker pair in the template → ``None`` (the
model has no reasoning surface).

Assert both markers survive a ``skip_special_tokens=True`` detokenize.

The split consumes decoded text, and the incremental detokenizer
decodes with ``skip_special_tokens=True``
(:mod:`arbi_serve.tokenizer`). A checkpoint that registers ``</think>``
as a *special* token would therefore drop it before the split ever ran
— every answer would silently land in ``reasoning_content`` with an
empty ``content``. That is exactly the silent-corruption mode this
project refuses: probe the real round-trip through the real tokenizer
and fail the build loudly instead.

Pluggable reasoning-trace (``<think>``) markers, per model family.

A thinking model emits its chain-of-thought inside a delimited block and
its user-facing answer outside it. The OpenAI-compatible surface carries
the two apart: the trace goes in ``reasoning_content``, the answer in
``content``. What differs between model families is only the delimiter
pair:

  - Qwen3.x / DeepSeek-R1 / MiMo:  ``<think>`` … ``</think>``
  - Seed-OSS:                      ``<seed:think>`` … ``</seed:think>``
  - Kimi:                          ``◁think▷`` … ``◁/think▷``

A :class:`ReasoningParser` declares one such pair. The split machinery —
non-streaming and incremental — lives in
:mod:`arbi_serve.engine.reasoning_parsing` and is marker-agnostic: it
reads the strings off the parser instance and never hardcodes them.

True when ``rendered_prompt`` leaves the model inside a think block.

This is the model-truth signal for whether generation starts as
reasoning, and it is read off the prompt the model's own chat
template just produced — not guessed from a request flag or a
per-model default.

A trailing unclosed start marker means "already reasoning". No
marker, or a closed pair, means generation starts as content.

Qwen3.x reasoning markers (``<think>`` / ``</think>``).

Shared verbatim by DeepSeek-R1, MiMo and the other ``<think>``-family
checkpoints, so this parser is the registry's default for any chat
template that carries the pair.

``<think>`` … ``</think>`` — Qwen3.5 / Qwen3.6 / DeepSeek-R1 framing.

Both markers are ordinary (non-special) tokens in the Qwen3.x
vocabulary, so they survive the ``skip_special_tokens=True``
detokenize (:mod:`arbi_serve.tokenizer`) and reach the split as plain
text. :func:`arbi_serve.engine.reasoning_parsers.validate_markers`
asserts that at build rather than trusting it.

Reasoning-block framing: split a thinking model's trace from its answer.

Marker-agnostic. The delimiter pair is read off the model-tied
:class:`~arbi_serve.engine.reasoning_parsers.base.ReasoningParser`
(resolved from the chat template at build); nothing here hardcodes
``<think>``. Mirrors the tool-call split in
:mod:`arbi_serve.engine.tool_parsing`, which this module runs upstream of:
reasoning is stripped first, and only the answer region is handed to the
tool extractor.

Whether the output already opens a reasoning block is not read from a
request flag: some chat templates pre-fill the opening marker into the
generation prompt, so the trace can arrive with no opening marker in the
output. :meth:`ReasoningParser.prompt_opens_reasoning` reads the rendered
prompt — the actual product of the model's own template and the request's
kwargs — and asks whether it ends inside an unclosed block. The caller
passes that answer in as ``opens_in_reasoning``.

Length of the longest suffix of ``text`` that is a proper prefix of ``marker``.

Lets the streaming split hold back a fragment that could still grow
into the closing marker, so a ``</think>`` straddling a chunk boundary
is never streamed out as reasoning text. Cheap: only the last
``len(marker)-1`` chars are inspected, and the common no-``<`` token
exits on one C-level ``find``.

Split ``text`` into ``(reasoning, content)``. Non-streaming path.

``reasoning`` is ``None`` when the response carries no trace at all —
no parser for this model, or the prompt did not open a block (thinking
off), in which case the text passes through as content untouched.

An unterminated block — the model hit ``max_tokens`` mid-trace — yields
all of ``text`` as reasoning and an empty content, losing nothing.

Incremental reasoning split for SSE streaming.

One instance per request (lives on the request's
:class:`~arbi_serve.engine.tool_stream.ChatToolStream`, never shared).

Contract mirrors
:class:`~arbi_serve.engine.tool_parsing.StreamingToolCallExtractor`:
:meth:`feed` is given the FULL accumulated output text every call — so
a stop-string trim that retroactively SHORTENS it is handled — and
returns only the forward progress:

  ``(reasoning_delta: str, content_so_far: str)``

``content_so_far`` is the accumulated post-marker answer region (empty
until the block closes), suitable to hand straight to the tool
extractor, which does its own delta tracking against it.

The closing marker is held back while it could still be completing, so
a ``</think>`` split across chunk boundaries never leaks half a marker
into ``reasoning_content``.

Flush any tail held back for a partial marker. End of stream.

An unterminated block (truncated at ``max_tokens``) can leave a
few chars withheld as a possible ``</think>`` prefix. They are real
trace text, so release them rather than dropping them on the floor.

Per-request state, sampling params, and lifecycle.

Each request's per-layer state lives in ``Request.state_handles``,
keyed by :class:`StateKind`: paged-KV layers store a page-id list,
Mamba / GDN layers store opaque state handles.

This class is heavily mutated by the scheduler / engine / detokenizer
(``state``, ``prompt_consumed``, ``output_token_ids``, ``output_text``,
``finish_reason``, …) on every step. The ``tenant`` and ``adapter_state``
fields are optional with non-breaking defaults for tenant / adapter wiring.

Lifecycle states for a request.

Allowed transitions:
  WAITING → PREFILLING → DECODING → FINISHED
  PREFILLING → WAITING (preempted; pages freed, prompt re-tokenizes)
  DECODING → WAITING   (preempted)
  DECODING → SUSPENDED (batch job parked; KV offloaded GPU→host→disk)
  SUSPENDED → DECODING (resumed; KV faulted back into fresh pages)

Constrained-decoding spec on a request.

``type`` is one of:
  - ``"text"``     — unconstrained (default; OpenAI legacy).
  - ``"json_object"`` — free-form JSON (compiled as JSON grammar).
  - ``"json_schema"`` — JSON validating against ``schema``.
  - ``"regex"``    — output must match ``regex``.
  - ``"grammar"``  — raw EBNF in ``grammar``.
  - ``"structural_tag"`` — xgrammar structural-tag JSON in
    ``structural_tag`` (used internally for tool-call enforcement;
    see :mod:`arbi_serve.sampler.tool_grammar`).

Mirrors OpenAI's ``response_format`` body field, extended with
``regex`` and ``grammar`` types for grammar-constrained decoding.

One in-flight request.

Lives in the engine's request map until ``state == FINISHED``,
after which the API layer streams the final SSE chunk and drops it.

``eq=False`` — a request is its identity, never its field values. Two
requests holding the same prompt are two requests. The generated
``__eq__`` would build a tuple of EVERY field on both operands for a
single comparison, which is what ``x in container`` /
``container.remove(x)`` calls per element scanned: the scheduler's
waiting/running deques and the graceful queue are all scanned that
way, so the field-tuple build dominated the scan. Identity ``__eq__``
makes those scans a pointer compare, and it restores ``__hash__``
(the generated ``__eq__`` suppressed it), so a request can live in a
set or key a dict.

``state_handles[StateKind.PAGED_KV]`` is the request's list of page
IDs (managed by :class:`FlatPageTable`); future state kinds add
their own handle types alongside.

True iff this request is a persistent duplex-lane session.

Mirrors :attr:`is_batch_priority`'s exact shape (a computed
property over :attr:`priority`, not a separate stored field) —
this is ``docs/nemotron-voicechat-duplex-design.md`` §7 step 3's
own resolution of the §6 decision-point-3 field-shape choice:
"reusing the ``is_batch_priority``-style pattern ... for
``is_duplex_frame``".

True when ``tools`` + ``tool_choice`` should constrain decoding.

OpenAI's default when ``tools`` are present but ``tool_choice`` is
omitted is ``"auto"``, which is treated as enforcing. Only an
explicit ``"none"`` opts out.

Build + pair a consumer half for a directly-constructed Request.

The production intake creates the :class:`ClientRequest` API-side
(``submission.new_client_request``) and the RequestFactory pairs
it; this helper serves test fixtures (and in-proc harnesses) that
construct a ``Request`` without going through submission.
``prompt_token_ids``, ``sampling`` and ``timing`` are shared
references (the prompt buffer is never mutated in place — see
``submission._encode_prompt``; the others are engine-write-once,
consumer-read). Idempotent: a second call returns the existing
pairing.

The consumer wakeup event — lives on :attr:`client`.

The engine never sets it; the OutputApplier fires it after
applying a batch. Exposed here so in-process await-the-result
consumers (bench harnesses, calibration drivers, scripts) that
hold the engine handle wait on the client event directly.

Fails loudly when no consumer half is paired: in process mode
``client`` is ``None`` engine-side and no engine-side code may
wait on consumer state — a caller landing here is reaching
across the process boundary.

Tokens occupying KV slots: prompt (consumed portion) + output.

Includes any in-flight optimistic MTP-verify advance
(:attr:`_mtp_optimistic_len`) so the scheduler + the next step's
positions see the assume-all-accepted length while a deferred
verify step is parked. The drain
(:func:`arbi_serve.engine.run_step.drain_pending_verify`) zeroes
it and appends the real accepted tokens before the next plan
build, so this advance is transient and never double-counts.

Also includes the async-output advance (:attr:`_async_optimistic_len`):
the K=1 host commit is deferred one tick, but ``total_length`` must
stay one tick behind the fed ``last_sampled_gpu`` so the next step's
decode position doesn't duplicate a KV slot. Zero at rest.

Includes :attr:`context_consumed` — tokens injected mid-decode by
:attr:`pending_context_token_ids`, which occupy KV slots and
advance RoPE positions exactly like prompt tokens but are never
appended to ``output_token_ids`` (they are context, not output).
Zero for every request that never injects, so ``is_prefill`` and
this property keep their meaning for all other callers.

Invalidate BOTH parked optimistic advances (MTP verify + async).

Called by every lifecycle event that moves a request off the slate
its parked step was built against — preemption, offload-suspend and
finish. Each of those rewrites the row's page-table state, so the
parked step's KV slots stop existing and its assumed length stops
being true; a parked advance that survives the move inflates
:attr:`total_length` for a row that no longer holds those slots.

One seam rather than a field-clear at each site: the two advances
are documented as symmetric and die for the identical reason, and
clearing only one is exactly the bug this replaced — every
lifecycle site zeroed :attr:`_async_optimistic_len` and left
:attr:`_mtp_optimistic_len` armed, so a preempted row's parked
verify went on to retire draft slots the preempt had already freed
(``radix_pagetable_spec`` raised "the engine over-trimmed").

Zeroing :attr:`_mtp_optimistic_len` is also the signal the verify
drain reads: a row parked with a non-zero advance that arrives at
the drain with zero was invalidated in between, and the drain must
skip both its slot retire and its token commit
(:func:`arbi_serve.engine.run_step.drain_pending_verify`).

Query-token count this request's next row occupies.

``1`` for an ordinary decode row (the pending sample alone) and
``1 + len(pending_context_token_ids)`` while injected context is
queued: the injection row re-feeds the pending sample at its own
position and then the caller's context tokens, so the sampled
token is never dropped. The slate builder emits this as the row's
``n`` and :func:`_gather_slate` re-derives it to validate.

True iff the request's sampling priority is the low ``"batch"`` class.

Mirrors :attr:`SamplingParams.is_batch_priority`; lets the
scheduler read the priority off the request without reaching
into ``req.sampling`` at every call site. Defaults to ``False``
(interactive) for any request whose sampling lacks the field
(older fixtures), so the priority-aware gates are zero-effect
on the all-interactive default path.

This session's frame cadence, in seconds.

``engine_default`` is ``batch_cfg.duplex_frame_interval_s``, used
whenever this request's own
:attr:`SamplingParams.duplex_frame_interval_s` is unset — so every
session that does not ask for a different cadence still paces on
the engine-wide 80 ms grid exactly as before. ``0.0`` is a
meaningful value (UNPACED), which is why the check is ``is None``
rather than a falsy test. Negative values are clamped to ``0.0``.

``getattr``, not a bare read: minimal ``sampling`` stand-ins in
CPU tests predate the field and must keep the engine default.

Whether THIS session's own frame deadline has come due.

A disabled persistent session is parked between bounded turns and
is never due, regardless of its retained deadline.

``next_duplex_deadline is None`` — never ticked — is due, so a
freshly admitted session is not stuck behind a deadline nobody
armed. An unpaced session (interval ``0.0``) re-arms to ``now``
every tick and is therefore always due again immediately.

Engine-side RequestFactory: :class:`SubmitMsg` -> :class:`Request` (P3).

Design (``docs/engine_core_process.md`` §3/§4/§5.1)
---------------------------------------------------
The API side always builds a :class:`~arbi_serve.engine.proc.messages.SubmitMsg`
(:mod:`arbi_serve.engine.submission`); this module is the engine half that
reconstructs the engine :class:`Request` from it and publishes it to the
scheduler. :func:`build_and_publish` runs on the engine loop (the API side
marshals through ``LoopBridge.run_on_engine_coalesced``) and absorbs, in
order: prompt-length + vocab-range guards, stop-id merge, MTP-K resolution
+ admission gates, LoRA validate + ref-acquire, the timing object,
``Request`` construction (+ ``ClientRequest`` pairing), OmniStreamState,
media-salted cache keys, the CPU-prep kick, timeout / backend / metrics
bookkeeping, grammar compile + attach (the one await), tenant binding, and
finally the publish body (detok worker, graceful residency queue or
scheduler add, wakeup, ``on_admitted``).

Failure contract (pick-one, per failure class):

* **Pre-publish validation** raises :class:`SubmitRejected` — resolved as
  the negative SubmitAck; the request is never published and never emits
  a ``FinishOut``. Grammar-compile failures are normalized here to a
  400-class :class:`SubmitRejected`.
* **Post-admission drops** (scheduler refusal, timeout, cancel — anything
  after the publish) keep emitting ``FinishOut`` through the output bus,
  exactly as P2's applier expects.

By-reference in-proc payloads
-----------------------------
``msgspec.Struct`` messages cannot carry live Python objects, so the
runtime-only companions of a ``SubmitMsg`` cross alongside the message in
thread/inline mode (the bridge intake closure captures them):

* ``client`` — the API-side pre-made
  :class:`~arbi_serve.engine.client_request.ClientRequest` (already
  registered with the OutputApplier). The factory pairs it onto
  ``Request.client`` so in-proc consumers (bench / calibration / live-GPU
  tests) keep using ``handle.client`` + ``new_token_event``. In process
  mode (P4b) nothing crosses and ``Request.client`` stays ``None``
  engine-side.
* ``features`` — the preprocessed multimodal features dict
  (``{"image": MultiModalFeatures, ...}``). Process mode converts through
  :mod:`arbi_serve.engine.proc.media` (``SubmitMsg.media`` tensor-ext).
* ``tenant_ctx`` — the rich :class:`~arbi_serve.tenant.context.TenantContext`
  from the auth middleware. Process mode reconstructs a minimal context
  from ``SubmitMsg.tenant`` (the id string).

Cancel-before-publish (tombstones)
----------------------------------
Request ids are allocated API-side, so a ``CancelMsg`` can legally arrive
before its ``SubmitMsg`` is processed. :meth:`Engine.cancel_by_id` records
unknown ids in the engine's :class:`CancelTombstones`; the publish path
checks it and drops a tombstoned submission with a ``cancelled``
``FinishOut`` (never scheduled, LoRA ref released). The set is TTL- and
size-bounded and fails loudly on overflow.

Bounded record of cancels that arrived before their SubmitMsg.

``Engine.cancel_by_id`` adds an entry when the id is not (yet) in
``eng.requests``; :func:`publish_and_admit` consumes it and drops the
racing submission with a ``cancelled`` ``FinishOut``. Entries expire
after ``ttl_s`` (a tombstone whose submission never arrives — e.g. a
cancel for a genuinely bogus id — must not accumulate); if the set
still exceeds ``max_entries`` after pruning, ``add`` raises — a flood
of cancels-for-unknown-ids means a protocol bug upstream and must
surface loudly, never silently drop tombstones (which would turn
into un-cancellable racing requests).

Engine-loop only (written by the marshaled cancel, read by the
marshaled publish) — no locking needed.

Engine-side: reconstruct a fresh :class:`SamplingParams`.

``kind`` (from ``SubmitMsg.kind``) is the single source for the task
mode — ``SamplingMsg.task`` is deliberately NOT read (deprecated;
kept on the wire for compatibility only). The returned object is
engine-owned: the factory mutates it (stop-id merge, mtp_k resolve)
without touching the API-side original.

Cached frozenset of the engine tokenizer's stop-token ids.

The engine-side stop ids are static for a given tokenizer, so a
per-request ``set(eng.tokenizer.stop_token_ids)`` rebuild is pure
GIL-bound waste. Cache the frozenset on the engine, keyed by the
``id()`` of the tokenizer's stop-id list so a model hot-swap (which
rebinds ``eng.tokenizer`` to a new list) transparently re-derives
the set instead of serving a stale one.

Returns an empty frozenset when the tokenizer exposes no stop ids.

The model's end-of-speech-segment control id (``<tts_end>``), or None.

Read from the served model's ``output_modalities()`` spec so the id is
never hard-coded here (text-only models / non-omni models return None).
Cached on the engine keyed by the model identity so the per-request lookup
is a dict read after the first call.

Always

A request for a non-active resident is PARKED synchronously in the
graceful queue (:func:`publish_and_admit` → ``_multi_group_queue``)
and the run loop performs the fast park/wake swap after the active
scheduler drains. So :func:`build_and_publish` never awaits a resident
switch and the awaitless intake drive is preserved for every
resident-routed request. Kept (inert) for API stability.

True when driving :func:`build_and_publish` will suspend.

The one gate the intake marshalers (thread-mode
``submission._asubmit_marshaled`` and the process-mode
``EngineBindings._handle_submit``) consult to pick task-wrap vs the
synchronous coalesced drive. Mirrors the factory's await points
exactly: the grammar attach and the P6 resident switch.

True when :func:`build_and_publish` will await a grammar compile.

The API-side intake closure uses this to decide between driving the
factory coroutine to completion synchronously inside the coalesced
drain (the common, awaitless path) and wrapping it in a task (the
grammar path — compile is the factory's one await). Mirrors the
attach gate in :func:`build_and_publish` exactly.

Every grammar request attaches, on every routing. ``Request.grammar_state``
has exactly one builder (:meth:`XGrammarLogitsProcessor.attach`); the SPMD
per-rank matcher is built on a different object
(``_MirroredRequest.grammar_state``) and is additive, never an alternative.
No routing flag selects between builders — so there is no routing in which
nobody builds one.

The wall-clock budget one arriving request is held to.

The request's own ``timeout_s`` when it set one, else the server
default. ``None`` disables the sweep for that request.

Resolved from ``eng.cfg.default_timeout_s`` at ARRIVAL — never
snapshotted at boot — so a live ``POST /v1/admin/config``
``default_timeout_s`` governs every request submitted after it, and
in-flight requests keep the budget they were admitted under.

Per-model state used to validate a request for ``target_resident``.

A request for a parked resident is built before the scheduler performs the
park/wake transition. Its model defaults and admission ceilings therefore
come from the target's captured state, while an unpinned request continues
to use the active engine attributes.

Validate ``submit`` and build the engine :class:`Request` (sync).

Everything :func:`build_and_publish` does except the grammar attach
(async) and the publish. Raises :class:`SubmitRejected` on any
client-input validation failure. Exposed separately so pure
validation tests can exercise the gates without an event loop and
so the sync :func:`~arbi_serve.engine.submission.submit` compat
entry can reuse it.

Undo build-time acquisitions when a post-build validation step
(grammar compile) rejects: LoRA ref + the in-flight CPU prep. The
request was never published, never entered ``eng.requests``, and
emits no ``FinishOut`` — the rejection travels via the SubmitAck.

Drop a submission whose cancel arrived first (never scheduled).

Emits the ``cancelled`` ``FinishOut`` through the output bus so the
registered ``ClientRequest`` resolves exactly like any other
terminal state, and flushes immediately (no step flush will carry
it — the request never enters the scheduler).

Publish ``req`` to the scheduler + wake the run loop (engine loop).

This is the intake boundary body. ``scheduler.add`` (or the
graceful residency queue add) mutates engine-owned admission state and
``_wakeup.set()`` wakes the run-loop coroutine, so this must run on
the engine loop; :func:`build_and_publish` already does (the API
side marshals), and the sync ``submit`` compat entry marshals this
function through the bridge.

Checks the cancel tombstones first: a submission whose cancel-by-id
arrived before this publish is dropped with a ``cancelled``
``FinishOut`` instead of being scheduled.

True when a request should be PARKED in the graceful queue for a
stable-VA resident swap (rather than admitted to the scheduler).

Parks only when residency is engaged, the graceful queue exists, the
target is a known resident, and it is NOT already active. A request
for the ACTIVE resident (or when residency/queue is absent) returns
False → the scheduler fast path (no swap).

Validate the request's stable-VA resident target (P6) — no switch.

Runs against the LIVE controller on the engine loop (the authority),
so a request naming a resident that vanished from the API child's
cached snapshot is rejected here rather than mis-served:

* unknown resident (record dropped since the child's snapshot) →
  :class:`SubmitRejected` 404, byte-matching the route's 404 path;
* residency not engaged engine-side at all → a "shouldn't be
  possible" topology mismatch; fail loudly (500-class), never
  serve the wrong model silently.

The Pure validation — never suspends — so the awaitless
intake drive is preserved.

Reconstruct + validate + publish one submission (on the engine loop).

The single engine-side intake entry (design §3): resident
enforcement (P6; awaits only on a pending switch) -> validation ->
:class:`Request` -> grammar attach (absent for the common request;
``submit_will_await`` mirrors both await points) -> publish.
Raises :class:`SubmitRejected` for every
pre-publish validation failure; the API side resolves it as the
negative SubmitAck. See the module docstring for the by-reference
in-proc payloads (``client`` / ``features`` / ``tenant_ctx``).

Mode-agnostic residency routing surface + the residents/groups snapshot.

The routing layer (``server/routes/_request_ctx.py``) resolves a
request's ``model`` field against "the resident set" without caring
which execution mode it runs in:

* **in-proc** (thread / inline): ``eng.stable_va`` is the live
  :class:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController`
  — authoritative engine state.
* **process mode**: ``eng`` is the API-child
  :class:`~arbi_serve.engine.proc.api_engine.ProcEngine` facade and
  ``eng.stable_va`` is a :class:`CachedResidencyView` fed by the
  engine's residents/groups snapshot (boot ``ReadyMsg`` fields +
  ``ControlEventMsg(kind="residency_changed")`` updates).

Both satisfy the tiny :class:`ResidentRouting` contract, so the route
has one code path (design rule: no per-mode branches in routes). The
marker-base-class form (not a ``runtime_checkable`` Protocol) is
deliberate: an ``isinstance`` check against a real class keeps a
``MagicMock`` engine (whose auto-vivified ``stable_va`` attribute would
satisfy any structural check) on the routing-disabled default path.

This module must stay import-light (no torch / engine chain): it is
imported by the torch-free route module and by the GPU-free API child.

:func:`residency_snapshot` is the one builder of the snapshot payload —
``build_ready_msg`` maps it onto the ``ReadyMsg`` trailing fields and
the engine bindings JSON-encode it into the ``residency_changed``
control event, so the two carriers can never drift.

Marker base for objects the model-routing layer may route against.

Contract (both implementations honor it):

* ``keys`` — ``list[str]`` of registered resident routing keys.
* ``active`` — the active resident's key, or ``None``.
* ``has(key)`` — exact-key membership.

API-child cached mirror of the engine's resident set (process mode).

Approximate by design (it lags the engine by one control event);
the engine-side RequestFactory re-validates every routed request
against the actual active resident (``SubmitMsg.target_resident``),
so a stale view can delay a switch decision but never mis-serve.

The underlying ``model.path`` a resident record was built from, or None.

Config-override variants (``cfg-<sig>``) and the boot member of the same
model all carry the SAME ``model.path`` in their captured engine-state
snapshot; only their runtime/capture config differs. Reading it lets a
caller tell "a config variant of the model I named" apart from "a genuinely
different model", and lets a request name a resident by the checkpoint path
it was loaded from.

Best-effort — returns None on any missing snapshot, so a record mid-teardown,
a CPU bookkeeping record, or the process-mode
:class:`CachedResidencyView` (which carries keys, not snapshots) simply does
not answer rather than answering wrongly.

Resolve one concrete model name to a registered resident key, or None.

The ONE identity map from a caller-supplied model name to a resident. Both
users of it go through here — the request router
(``server/routes/_request_ctx._resolve_resident_key``, which layers the
``"default"`` alias on top) and the runtime load/switch admin path
(``engine/runtime_resident.resolve_present_key``) — so a name that ROUTES to
a resident also LOADS to it, and neither can gain an alias the other lacks.

Resolution order — each step is unambiguous or it falls through:

  1. exact resident key (id) → that resident. Covers a configured
     ``served_name``.
  2. a resident key's path BASENAME (``"Qwen3.5-9B"`` for a resident keyed
     ``/models/Qwen3.5-9B``) → that resident, but ONLY when the basename is
     UNAMBIGUOUS across the resident set. An ambiguous basename (two
     residents share it) returns None so the caller refuses rather than
     guesses.
  3. the CHECKPOINT PATH a resident was loaded from → that resident. The
     inverse of (2), and the reason it exists is that ``--model`` takes a
     path while the resident is keyed by its served name: without this, the
     identifier the server was STARTED with is not an identifier it answers
     to, and ``--model /models/Qwen3.5-9B`` followed by a request naming
     that same path 404s with the model itself in the "Available" list.
     Matched against each record's captured ``model.path``, so this is an
     identity, not a guess: it resolves the path that record was actually
     built from and nothing else. A path whose BASENAME merely equals a
     resident key is deliberately NOT accepted — two checkpoints in
     different directories share a basename, and honouring that would answer
     a fully-qualified request with a different model.

We deliberately do NOT alias an arbitrary unknown name to the active
resident — that would silently mis-route a genuine multi-resident request.

Snapshot the engine's residents/groups routing state (JSON-safe).

The single source for both carriers of the snapshot (`ReadyMsg`
trailing fields at boot, ``residency_changed`` control-event payload
on every mutation). Strings/JSON only — no live objects (§5.3).

Defensive by contract (heartbeat-adjacent): tolerates a mid-build
engine with attributes still unset.

The one place that decides which of the three row shapes a slate row is.

``_gather_slate`` (:mod:`arbi_serve.runtime._batch_build_materialize`,
rank 0 / single-process) and ``derive_step_tensors``
(:mod:`arbi_serve.distributed.spmd_derive`, its SPMD rank-local twin)
both walk a ``slate: list[tuple[Request, int]]`` and, for every row,
answer the same question: given this request's committed state and
the caller-supplied row width ``n``, is this a PREFILL row, an DECODE row — and where does it start in the
sequence?

That decision is pure scalar arithmetic. It doesn't touch a page
table, a device, or a GPU tensor — so there was never a real reason
for the two builders to each hand-write it. INJECT
branch in ``_gather_slate`` but the SPMD twin's decode branch stayed
pinned to width 1 for one release, which would have silently
corrupted the flat batch (``len(flat_ids) != len(flat_positions)``)
the first time a duplex tool-response injection landed under TP>1.

:func:`resolve_row_shape` is now the ONLY place that arithmetic lives.
Both builders call it instead of re-deriving ``is_inject`` /
``start_pos`` / the width assertion themselves. Everything that
legitimately differs between the two builders — numpy vs. Python-list
token gather, GPU D2D sample overwrites, pinned-buffer staging,
CPU-prep fast paths, the host-vs-device execution boundary itself —
stays local to each; none of that is row-SHAPE logic.

Adding a fourth row shape means adding a :class:`RowKind` member and a
branch here. Both callers match on ``RowKind`` exhaustively and raise
on the fallback (see the ``else`` arm in each), so forgetting to wire
the new kind into either builder fails loudly the first time a slate
actually produces it — not silently, under whichever topology nobody
happened to test first.

The resolved shape of one slate row.

``start_pos`` is the position of the row's FIRST query token; the
row occupies flat positions ``[start_pos, start_pos + n)`` and, on
a PAGED_KV model, the ``n`` KV slots just allocated for it.

Resolve a slate row's shape from request state + the caller's ``n``.

Raises ``ValueError`` when ``n`` disagrees with the width the
resolved shape requires:

  - INJECT: ``n`` must equal ``1 + pending_context_len`` (the
    not-yet-forwarded sample plus the queued context). A mismatch
    means positions and the just-allocated KV slots would disagree
    by the gap — silent corruption, not a shape a caller should
    ever intentionally build.
  - DECODE: ``n`` must be ``1`` — with no pending prompt tokens and
    no pending context, a row is always exactly one token wide.
  - PREFILL: any ``n`` is caller-determined (a chunk size); not
    validated here.

A caller with a stale row-width computation gets a loud
``ValueError`` here instead of a flat-vector length mismatch that
silently desyncs positions from slots several lines downstream.

What happens to the PROCESS when the engine run loop leaves.

The run loop is the only thing that turns an admitted request into
tokens. In the in-process topologies it does not own the process: it
runs on a dedicated OS thread (:class:`~arbi_serve.engine.engine_thread.EngineThread`,
default) or as an asyncio task on uvicorn's loop
(``ARBI_ENGINE_OWN_THREAD=0``). An exception that escapes it therefore
kills a THREAD, not the server — uvicorn keeps answering, the socket
stays bound, and the container keeps its ``restart:`` policy unused
because nothing exited. That is the dead-but-listening state: every
admitted request fails in milliseconds and the deployment's restart
policy never fires.

This module is the seam that closes it. A run-loop exit is classified
once, in one place (:func:`classify_run_loop_exit`), and a TERMINAL one
arms the process exit that the code raising already promises — e.g. the
drafter circuit breaker's "Raising out of the engine loop so
orchestration restarts the server instead of serving degraded forever".

**Terminality is a property of the LOOP, not of the exception.** Once
``run_forever`` has unwound there is no supervisor inside the process
that rebuilds it: no later state change can make the engine step again.
So every exception that escapes is terminal, whatever raised it — and
correspondingly, none of the recoverable states are exceptions that
escape:

  * a graceful drain / admin shutdown / SIGTERM — the loop RETURNS
    (``eng._shutdown``), it does not raise;
  * an admin model reload — runs inside the loop's utility dispatch and
    either rolls back or arms its own exit
    (:mod:`arbi_serve.engine.swap_admin`); the loop keeps ticking;
  * sleep / release-memory, a park or an evict — the loop's quiesce
    branch keeps ticking. Sleep quiesces on ``eng._memory_released``, NOT
    on ``scheduler is None``: a phase-2 release frees the pools' physical
    and leaves the scheduler object in place, which is exactly why
    ``/health/ready`` had to grow an ``engine_asleep`` reason of its own
    (readiness saw a built engine and answered 200). ``scheduler is None``
    is the OTHER disjunct — a live rebuild — and it is likewise not an
    exit. Either way this module is never reached, because the loop has
    not unwound;
  * a transient 503, a cancelled request, a load shed, a per-request
    step failure, a single-rank infra OOM under the self-healing latch —
    all handled INSIDE the per-step ``try``: the slate is finished with
    ``reason="error"`` and the loop continues.

The one OOM that is NOT in that last group is the end of the
memory-pressure ladder (:mod:`arbi_serve.engine.memory_pressure`): an OOM
that lands while admission is already closed for memory pressure and no
step has succeeded since it closed means narrowing has run out, so it
raises ``StepMemoryExhaustedError`` out of the loop DELIBERATELY. It is
terminal for the same reason everything else here is — the loop has
unwound — and this module is what makes that verdict reach the process
rather than only the engine thread. A card genuinely too small restarts
into whatever smaller ``--max-batch`` / ``--max-context`` /
``--gpu-memory-utilization`` an operator sets, instead of serving
degraded forever.

Two more exits are deliberately non-terminal even though they arrive as
exceptions: :class:`asyncio.CancelledError` (the lifespan cancelling the
run task during teardown) and ``SystemExit`` / ``KeyboardInterrupt``
(something else already chose how this process ends). Anything raised
after shutdown was requested is teardown noise, not a fault — including
the ``RuntimeError`` that ``run_until_complete`` raises when
:meth:`EngineThread.stop` force-stops a loop that overran its join
timeout.

**Arming is explicit and positive.** Nothing exits until the serving
entrypoint calls :func:`arm_terminal_exit`, so an import, a unit test, a
bench harness, or an in-process calibration run that drives a run loop
can never take the process down: they get the CRITICAL log and nothing
else. The serve CLI is the only caller (see
:mod:`arbi_serve.cli.bootstrap`).

Topologies:

  * **own-thread, in-proc (default)** — the incident path. The thread
    dies, this module exits the process, and ``restart: unless-stopped``
    recycles the container within seconds.
  * **in-loop, in-proc** (``ARBI_ENGINE_OWN_THREAD=0``) — the run task's
    done-callback lands here with the same classification.
  * **process mode** (``ARBI_ENGINE_PROC=1``) — already correct without
    this module: the loop IS the engine process's main coroutine, so the
    exception unwinds out of ``run_engine_process`` and the process exits
    non-zero on its own. Arming is harmless there (nothing calls in).
  * **multi-rank TP** — rank 0 in-proc is the case this fixes: a dead
    rank-0 loop used to leave every worker rank blocked in the collective
    with no peer, invisible to torchrun. Exiting rank 0's process makes
    torchrun tear the workers down. Worker ranks already exit non-zero by
    unwinding out of ``asyncio.run(driver.run_forever())``.

Whether the engine has been told to stop.

``eng`` is the Engine in every topology — the distributed driver's
:meth:`shutdown` delegates to the engine's, so this one read covers
the driver path too. ``getattr`` because the run loop can also be
driven against a scheduler/engine stand-in that predates the flag.

Report a run-loop exit and, when TERMINAL, arm the process exit.

Returns the classification so a caller (and a test) can assert on it
without inspecting the process. The CRITICAL line is emitted BEFORE
the exit is armed, and the exit flushes every log sink on its way out
(:func:`~arbi_serve.calibration.child_exit.flush_log_sinks`), so the
reason reaches the log sink ahead of the death it explains.

Engine hot path: ``run_step`` entry point, run loop, terminal-state.

The single engine step entry point. The :func:`run_forever`
coroutine drives the engine step by step:

  1. Sweep request timeouts.
  2. Ask the scheduler for a slate ``[(req, n_tokens), ...]``.
  3. Build a :class:`StepPlan` (K-aware, frozen) from the slate.
  4. Dispatch through ``eng.spec_decode.run_step`` — the strategy
     decides whether to run the K=1 path (:class:`NoSpecDecode`)
     or split into MTP verify-pass + sub-slate
     (:class:`MtpStrategy`).
  5. Per-step: the strategy / runner materializes the batch, runs
     metadata builders, runs forward (captured-graph replay or
     live), and samples. Terminal-state callbacks
     (:func:`on_finished`, :func:`cancel`, :func:`shutdown`) and the
     per-token finish-condition checker (:func:`post_token`) live
     here.

The K=1 :func:`step` body lives here as the ``ModelRunner.execute``
consumer; the MTP verify body lives in
:mod:`arbi_serve.spec_decode.mtp` (the unified strategy module).

The :class:`StepPlan` is the shape on the public boundary. The
internal ``ScheduledBatch`` shape stays inside
:class:`ModelRunner`; the engine API is
:func:`run_step` taking a :class:`StepPlan` and returning a
:class:`StepResult`.

This package splits ``run_step`` into concern submodules —
:mod:`.output` (emission), :mod:`.finish` (finish/metrics lifecycle),
:mod:`.terminal` (``post_token`` + streamed finish), :mod:`.loop` (run
loop + timeout sweep), :mod:`.embed` (embedding/rerank step),
:mod:`.step_variants` (async step paths + spec-mode routing) — with the
K=1 step, MTP seed-forward, and async-output commit/drain kept here.
Every symbol the codebase imported from ``run_step`` remains importable
from this package.

Public async step entry point — builds a :class:`StepPlan` and
routes through the engine's :class:`SpecDecodeStrategy`.

One entry point for K=1 and K>=2. The plan carries the K dimension
explicitly; the strategy
(``NoSpecDecode`` / ``MtpStrategy`` / ``ExternalDrafterStrategy``)
decides per-step body via ``plan.spec_mode``.

The slate-shaped ``[(req, n_tokens)]`` is the scheduler emit
shape; the plan is built around it via
:meth:`StepPlan.from_slate`. The strategy consumes the slate
internally.

Run one slate through the model runner, then commit + post-process.

The runner (``eng.model_runner``) materializes the batch, runs the
captured-graph replay or live forward, and samples. The run loop
handles the post-step concerns: scheduler commit, per-token finish
checks, and the commit-done timing stamp.

MTP seed-draft: after a prefill or K=1 decode step finishes for an
MTP-opted request, run :meth:`MtpDriver.draft` over the just-
sampled token's hidden so the request's NEXT step (which will
route to :func:`mtp_step`) has a head-driven draft ready in
``req.mtp_next_drafts``. Without this seed the first MTP
verify step falls back to the cold-path "repeat last token"
drafter and accept rate drops to near zero on that step.

TP>1: the engine's :attr:`worker_bridge` (set by
:class:`DistributedEngineDriver` per-rank) broadcasts each forward
/ drafter / seed to workers from inside ``model_runner.execute``,
``MtpDriver.draft``, and ``advance_mtp_seed``. No per-step batch
hand-off is required here.

Materialize the batch + metadata for an MTP-opted seed step.

Stays OUTSIDE ``inference_mode`` — some metadata builders allocate
state tensors that persist in the pool across steps, and making those
inference tensors corrupts state for later steps. The async path
(:func:`step_mtp_seed_async`) builds the batch on the loop thread and
then offloads only the GPU-blocking forward + drafter section.

The bucket's just-sampled tokens as a device tensor, without a round trip.

``forward_exec.sample`` already stamps every ready row's
``last_sampled_gpu`` — a 0-d view of the token that never left the device —
so building the drafter chain's input from the HOST ints sends back values
the device is already holding. The views are keyed on the request object, so
stacking them cannot mis-order the bucket the way an index into ``sampled``
could.

Falls back to the pinned host staging when any row has no device view (a
sampler path that does not stamp one, or a CPU engine). Value-identical
either way: both spell the same tokens in the same order.

Hand each row of a seed bucket its drafts, on the device where possible.

The chain writes its ``(>=depth, B)`` block on the device and the next
verify plan needs those tokens ON the device — as verify input ids and as
the accept test's draft tensor. Pulling them here is a stream drain the
step does not need: the host waits for the chain to finish before it can
build the plan or launch anything else.

So the seed side stashes the block exactly as the verify side already does
(:func:`~arbi_serve.spec_decode.device_drafts.stash_device_drafts`), and a
consumer that genuinely needs the values pulls its own row. It is the same
mechanism, on the path that never got it.

DEPTH IS A SHAPE, NOT A VALUE. The chain may return fewer than ``req_k``
slots when a per-slot depth gate stopped it, and the verify plan takes the
slate's K from the published length — but that length is
``drafts[:req_k].shape[0]``, which the host knows without reading the
device. The pull only ever supplied the token ids.

Falls back to the host pull where the device route cannot run (TP>1, whose
worker mirror derives its plan from the host lists; a non-CUDA engine; the
flag off) — the same conditions the verify side checks.

Forward + sample + batched seed-drafter chain for an MTP-opted step.

The GPU-blocking body of the :func:`step` MTP-seed branch. Both the
sync path and :func:`step_mtp_seed_async` (which runs this in the
forward executor, OFF the event loop) share this one implementation —
byte-identical kernels. Returns the ``result`` shim consumed by
:func:`_commit_step_result`. The per-request ``mtp_next_drafts`` /
``mtp_next_draft_probs`` assignments here are plain attribute writes on
the slate's requests (no ``asyncio`` / queue side effects), so they are
safe to run in the worker thread; the streaming ``post_token`` tail
stays on the loop in :func:`_commit_step_result`.

Sample + observe + batched seed-drafter chain, from a forward's outputs.

The post-forward half of :func:`_run_mtp_seed_forward`, split from the
forward so the SAME tail serves a prefill chunk whose forward ran fused
with the step's verify rows (:mod:`arbi_serve.runtime.fused_mixed_step`).
``logits`` is one row per slate row (the last token's), ``hidden_last``
the matching ``(B, H)`` hidden. ``observe_batch`` is what a
tap-conditioned drafter's ``observe_seed_forward`` reads its row
boundaries from when they are not ``batch``'s own — the fused layout's
prefill rows start where its verify rows end; ``None`` is ``batch``.

Commit a finished forward's sampled tokens + run per-token post.

The async event-loop-friendly path (:func:`step_async`) runs the GPU
forward off-thread (in a thread executor — the loop stays free to
admit / stream) and then runs this commit + ``post_token`` tail back
ON the loop thread. ``post_token`` runs the inline detok and hands
the worker-mode ``detok_queue`` its ``put_nowait`` — both
engine-loop-only — so it MUST stay here, not in the worker thread.

Async output. When ``result.pending`` is set (the runner sampled on
GPU without the host sync), this step's commit is DEFERRED: the
PREVIOUS in-flight step (if any) is drained + committed first (it
overlapped THIS step's forward), then this step's :class:`PendingStep`
is parked on ``eng._pending_output``. The host materialize + commit
of step N thus runs at the top of step N+1, after N+1's forward was
launched.

``already_drained`` (host-ahead pipeline, ARBI_ASYNC_SCHEDULE): the
caller (:func:`pipelined_step_async`) already ran the previous step's
drain OVERLAPPED with this step's forward dispatch, so the drain here
is skipped — and so is the first-token inline materialize below (a
full sync on THIS step's still-running forward, which would stall the
pipeline for a whole forward on every prompt-completing step; the
first token rides the normal one-tick deferral instead — a bounded,
one-engine-tick TTFT tradeoff for the host-overlap win).

True iff any row in ``slate`` completes its prompt THIS step.

Such a row produces the request's first output token (the
prefill→decode transition) — its TTFT. Read off the request state
BEFORE :meth:`Scheduler.commit_state` mutates ``prompt_consumed``: a
row still ``is_prefill`` whose ``prompt_consumed + n_tok`` reaches the
full prompt length is finishing prefill this step. Cheap O(slate);
the common steady-state decode slate has no prefill row and returns
on the first iteration's ``is_prefill`` miss.

Commit materialized sampled tokens + run per-token post.

The synchronous commit body — shared by the eager path and the
async-output deferred drain. ``sampled`` is the materialized
per-row int list (``None`` for non-ready rows).

Materialize + commit the one in-flight deferred step (async output).

Called at the top of the next step (after its forward is launched)
and on shutdown / before any synchronous step (MTP verify) that must
see committed state. No-op when nothing is pending.

Abort/cancel safety: a request finished between sample and drain
(disconnect handler set ``state == FINISHED``, or a stop on an
earlier row of this same drain) must NOT have its deferred token
committed or streamed. ``scheduler.commit`` only appends to
``output_token_ids`` and is harmless on a finished request, but
``post_token`` would re-stream + re-finish it — so we drop the token
for any row already FINISHED at drain time.

Flush the ENTIRE in-flight async-output pipeline.

Only one step is ever parked (``_pending_output``), so this is a thin
wrapper over :func:`drain_pending_output` — kept as its own name for
callers at a terminal boundary (empty slate / shutdown / a synchronous
MTP step that must see committed state), independent of whether the
pipeline is ever deepened again in the future.

Host-ahead K=1 step (ARBI_ASYNC_SCHEDULE) — a non-blocking step
whose forward dispatch overlaps the previous step's drain.

Submits the WHOLE ``execute`` (build + metadata + forward + sample
launches) to the single-worker forward executor, waits only for the
runner's ``on_inputs_built`` signal (every request-state read and
page-table mutation of step N done), then — while the runner thread
dispatches the forward — runs the PREVIOUS step's deferred drain
(materialize + ``commit_output`` + ``post_token`` / detok / SSE) on
the engine thread. Request finishes discovered during the overlap are
deferred (``eng._deferred_finish_queue``) and flushed once the runner
returns; the step then parks (:func:`_commit_step_result` with
``already_drained=True`` — ``commit_state`` + the optimistic
placeholder advance).

The single-threaded-CUDA invariant holds: the ONE executor worker runs
every pipelined launch, the engine loop awaits the future before the
next step, and the engine thread's only CUDA touch in the window is
the drain's (thread-safe) event synchronize. ``_forward_in_flight``
brackets the whole window so the admission-time offload self-heal
defers, same as :func:`step_async`.

Engine hot path: pure-prefill embedding / rerank step.

The dedicated one-forward step for embedding / reranking instances
(:func:`embed_step`): run the decoder, last-token-pool (embed) or read
the yes/no logits (rerank), stash the result, and finish the request
immediately — no sampler, no decode loop, no token stream. Plus the
torch.compile de-specialization helpers the live compiled embed forward
needs (:func:`_despecialize_embed_batch`, :func:`_bucket_up`).

Round ``n`` UP to the next power of two (min 1).

Used to collapse the unbounded ``max_query_len`` / ``max_seq_len``
int values — which Dynamo would otherwise bake as per-value guards
(``batch_meta.max_query_len == N``) — into ~log2(max_context)
buckets. Safe because these feed flash-attn's ``max_seqlen_q`` /
``max_seqlen_k`` sizing hints, which only require an upper bound on
the true max length (a larger value sizes a slightly bigger kernel
grid; the varlen kernel still attends exactly ``seq_lens`` tokens
per row).

Mark per-step varying tensor dims dynamic + bucket int sizing
hints so the live compiled embed forward stops re-specializing.

Two specialization sources are removed:

1. SHAPE guards — ``input_ids`` / ``positions`` / ``slot_mapping`` /
   ``cu_seqlens_q`` / ``cu_seqlens_k`` / ``seq_lens`` vary in dim 0
   (token count / num_seqs), and ``block_table`` varies in BOTH dims
   (num_seqs × page-table width). We call
   :func:`torch._dynamo.maybe_mark_dynamic` on both the batch and
   its ``attn_meta`` (the block forward reads ``batch_meta.*``
   mid-trace, so the meta's tensors must carry the dynamic mark too;
   they alias the batch tensors but the mark is per-tensor-object).

2. INT guards — ``int(batch_meta.max_query_len)`` /
   ``int(batch_meta.max_seq_len)`` read inside the traced custom-op
   call bake a constant guard per distinct value. We bucket both to
   the next power of two, collapsing thousands of values to a
   handful.

Pure-prefill step for embedding / reranking instances.

Every request is a single forward: run the Qwen3 decoder, then
either last-token-pool the post-final-norm hidden state (embed) or
read the ``yes`` / ``no`` last-position logits (rerank). The result
is stashed on the request and the request is finished immediately —
no sampler, no decode, no token stream.

No persistent KV is needed: with single-chunk prefill every
sequence is fully presented in this one batch, so attention rides
the tkv / bf16 first-chunk bypass (nothing is read back across
steps). The pages the scheduler allocated for the prompt are freed
by :meth:`Scheduler.finished` below — they live for exactly this
step.

Defensive on chunking: a row whose prompt is NOT fully consumed
this step (only possible if an operator shrinks ``chunk_prefill``
below ``max_context``) is left PREFILLING and picked up again on a
later step; its hidden state is only read once the final token is
in the batch.

Mark an embed/rerank request FINISHED and free its resources.

Mirrors the run-loop's terminal-state handling: commit + free KV
pages via :meth:`Scheduler.finished`, drop the LoRA ref, tear down
any detok task, signal the waiter, and emit finish metrics.

Engine hot path: terminal-state finish + metrics lifecycle.

The in-flight gauge admission bump (:func:`on_admitted`), the
per-request terminal metrics + trace burst (:func:`on_finished`), and
the shared "drop a request out of the engine" free sequence
(:func:`_finish_request`) used by the drained-drop, run-loop error,
group-swap-failure, timeout, and embed/rerank finalize paths.

Bump the in-flight gauge and log the call's arrival.

Idempotent: a request is counted into the gauge at most once even if
``on_admitted`` fires again (e.g. an offloaded job resuming). The
matching ``on_finished`` only decrements requests that were counted,
so a request that aborts/errors

One access-log line stating the call and its initial disposition.

The pair a reader follows is ``→ call#N`` here and ``✓/⊘/✗ done#N``
from :func:`on_finished`; nothing else about the call reaches INFO.
It carries what is decided at admission: the kind of work, how much the
call contains, and whether it is interactive or batch traffic.

``cached`` renders only on a HIT. The prefix resolve runs inside the
scheduler's admission, which the parked and deferred admission paths
reach later or not at all, so a zero here is "not resolved yet" as often
as it is "no hit" — and a number that cannot tell those apart is not one
to print. The finish line reports the resolved figure.

Two kinds of request announce their arrival somewhere else and so log
it at DEBUG here. Embed / rerank fan one API call out to one engine
request per input, and batch-priority work arrives inside a container
(a Batch API job, or a client that asked to be treated as bulk) that
states its own size once. Both still emit their per-request finish line
at INFO — a call's outcome is never summarised away.

Best-effort served-model label for tracing spans.

Mirrors the resource-attribute derivation in
:mod:`arbi_serve.server.app` (``served_name`` or the path
basename). Defensive: any missing attribute degrades to
``"unknown"`` rather than raising into the finish path.

The request's exact KV-cache footprint in bytes.

``pages × per_page_bytes`` — the request's allocated paged-KV pages
times the pool's per-page byte cost. The page count is the request's
own page-id handle (``state_handles[PAGED_KV]``), which is exact and
per-request (no cross-request conflation); it falls back to the
page-count implied by the request's token span
(``ceil((prompt + output) / block_size)``) when the handle is absent.
Returns 0 when the per-page cost is unknown. A plain-int read on the
engine task — no CUDA work.

Compact batch-config fingerprint for the request-timeline ring.

Groups the timeline's per-(model, backend) stats by the settings that
move throughput, so an A/B over the SAME model+backend but different
batch settings reports as separate rows. Reads the CURRENT effective
config, so a runtime override re-fingerprints from the next request on.

Record terminal-state metrics + drop the in-flight gauge.

Idempotent — guarded by the ``_metrics_finished`` sentinel on the
request, so every finish path (stop / length / context / timeout
/ cancelled / error) can call this freely.

Drop ``req`` from the stable-VA graceful queue. Idempotent.

A request published for a NON-active resident is PARKED in
``eng._multi_group_queue`` under its ``(resident, "")`` group key and
is never handed to ``eng.scheduler`` — so ``Scheduler.remove`` cannot
see it. Every terminal path runs this: a stale parked entry both pins
the request and its prompt until the group becomes resident, and is
then re-admitted by :func:`_advance_resident`, where
:meth:`Scheduler.add` rejects it for not being WAITING. A no-op for a
request that was never parked.

Run the terminal finish/free sequence for one request.

The shared body of every "drop a request out of the engine" site that
routes its terminal signal through ``cancel_detok_task`` +
:func:`_emit_finish`: the drained-suspended drop, the run-loop error
finish, the group-swap-failure drop, the timeout sweep, the
embed/rerank finalize, and the client-disconnect cancel
(:func:`arbi_serve.engine.lifecycle._cancel_inner`). It drops the
parked graceful-queue entry, frees the request's KV (via the
scheduler), marks it FINISHED, releases its LoRA ref, tears down its
detok task, emits the terminal ``FinishOut`` (the consumer wakeup
rides the applier), drops it from the live request map, and emits
finish metrics. The emission stages onto the step's batch; the caller's
context flushes (the run loop's per-iteration flush, a drain site, or
an explicit off-step flush like the cancel body).

Per-site variations:

  * ``reason`` — the ``finish_reason`` to stamp, or ``None`` to leave
    the request's existing reason untouched (the drained-suspended path
    already stamps ``offload_lost`` upstream, so it passes ``None``).
  * ``scheduler`` — ``"finished"`` runs :meth:`Scheduler.finished`
    (commits full pages into the radix cache + releases quota: the
    completion path); ``"remove"`` runs :meth:`Scheduler.remove`
    (drop from whatever queue WITHOUT a radix commit: the timeout /
    cancel path), suppressing a missing-key error; ``"none"`` skips the
    scheduler entirely (the group-swap-failure rows were drained from
    the cross-group queue and never admitted to ``eng.scheduler``, so
    calling ``finished`` on them would commit/quota a request the
    scheduler never tracked).

Resource release is the invariant, not a per-caller choice: the LoRA
ref, the parked graceful-queue entry, the detok task and the CPU-prep
arrays are released on EVERY terminal path. Each release is
idempotent, so a site that has already released one pays a predicate
and nothing else.

Ordering: parked-queue drop, then the scheduler action, then the
state flip, then LoRA release, then detok teardown, then the consumer
signal, then the map drop, then ``on_finished``.

Engine hot path: the per-step run loop + terminal sweep.

The :func:`run_forever` coroutine drives the engine step by step
(:func:`_run_forever_inner` is the loop body, split out so the
heartbeat + profiler lifecycle owns one try/finally), the stable-VA
residency advance (:func:`maybe_advance_active_group`), the once-per-iteration
timeout sweep (:func:`sweep_timeouts`), and the multi-rank predicate
(:func:`engine_is_multi_rank`) the step-outcome seam keys off.

THIS IS ONE OF THREE RUN LOOPS. A single-rank boot runs this one; a TP>1 boot
runs :meth:`~arbi_serve.engine.distributed_driver_spmd._SpmdLoopMixin.
_run_loop_spmd` or the legacy rank-0 driver loop instead. What a step OUTCOME
means is therefore not this module's to decide — the classification, the health latches and the memory-pressure ladder
live in :mod:`arbi_serve.engine.step_outcome` and every loop calls the same
two functions. This loop keeps only what is genuinely its own: finishing the
slate, the duplex tick rollback, and its own teardown.

True when the engine spans more than one rank (TP or SPMD).

A CUDA OOM inside a forward that carries cross-rank collectives desyncs the
ranks — the peers already completed the collective or are blocked in it /
waiting on the SPMD shm control bridge — so the engine loop must FAIL LOUD
(poison + tear down) rather than run the self-healing infra latch, which
would hang the shm writer or trip the NCCL watchdog (600s → SIGABRT). Single
rank (world==1, no worker bridge) has no peer to desync and keeps the
self-healing behaviour. Pure attribute reads — no CUDA. See the ``run_forever``
infra-error handler and ``DrafterCollectiveFault``.

Hand back the output-copy ring slots a failed step left orphaned.

The step's parked objects (``_pending_output`` / ``_pending_verify``)
are the only holders the next drain will consume; a slot held by any
other object, or by none, belongs to the step that just raised and
would otherwise refuse the next acquire.

Drive the engine's per-step loop until shutdown is requested.

Rank 0 emits the engine heartbeat; ``ARBI_PROFILE`` (torch / nvtx /
cprofile) builds a :class:`ProfilerController` driven per decode step.
Both wrap the loop body; ``off`` + worker ranks build nothing, so the
production hot path is untouched.

Pick + swap to the next parked resident, then admit its requests.

Called from the run loop when the active scheduler has fully drained.
The selector picks the cheapest resident; the FAST recapture-free
park/wake swap (``eng.aswitch_model``) runs the transition; finally
the resident's parked requests are admitted to the scheduler. Under
TP>1 the ``swap`` callback lets the distributed driver broadcast the
switch to every rank BEFORE the local park/wake (rank-lockstep).

Fast stable-VA residency advance: swap via park/wake, admit parked.

The resident is the analog of a group; its queue key is
``(resident, "")``. The swap is ``eng.aswitch_model`` (recapture-free
park/wake) — or the caller-supplied ``swap`` callback, which the TP
driver uses to broadcast the switch to every rank in lockstep.

Engine hot path: per-step output emission seam.

The single text/audio/finish emission points plus the per-step batch
flush. Every terminal + streaming path routes its consumer-facing
structs through here (``_emit_token`` / ``_emit_audio`` / ``_emit_finish``)
and hands the staged batch over exactly once per step (``_flush_output``).

One request class is deliberately NOT on this boundary: the duplex lane's
persistent per-connection request — see :func:`_has_bus_consumer`.

The engine's :class:`~arbi_serve.engine.output_bus.OutputBus`.

Lazily constructs a bus + applier pair on a partial / stub engine
(sync test harnesses + the teardown drain paths call the emission
helpers with engines that skip ``__init__``); a real serving Engine
builds both in ``_setup_request_state``. The lazy pair uses the
inline marshaler, matching the single-loop behaviour those harnesses
exercise.

True iff ``req``'s output belongs on the typed output bus.

Every ordinary request reaches the engine through the submission
intake (:mod:`arbi_serve.engine.submission`), which builds its
:class:`~arbi_serve.engine.client_request.ClientRequest` and
``register``s it with the engine's ``OutputApplier`` before the
request is published — so the applier has a registry entry for the id
every struct staged here carries.

A DUPLEX-LANE request does not. It is built by
``new_duplex_request`` and admitted by ``admit_duplex_request``
(:mod:`arbi_serve.realtime.nemotron_voicechat_duplex_admission`),
which bypass submission entirely: no ``ClientRequest`` exists for it
and nothing ever registers one. That is by design, not an oversight —
a duplex connection's consumer is its own
``DuplexConnection.outbound`` queue, woken once per tick through
:meth:`~arbi_serve.engine.loop_bridge.LoopBridge.notify_http`,
carrying raw waveform / ``TurnEvent`` payloads the ``TokenOut`` /
``AudioOut`` structs deliberately do not model (that method's own
docstring records the decision).

So a struct staged for a duplex request has no consumer by
construction: the applier can only drop it, and it logs a WARNING per
dropped item. Emitting anyway cost one wasted struct + batch slot per
committed token on the duplex tick's own latency-critical path, and
flooded the log with ``"TokenOut for unknown request N dropped (late
batch after cancel/finish is legal)"`` for the entire life of every
duplex session — a message whose parenthetical is false here (nothing
was cancelled or finished) and whose volume destroys its value as the
genuine late-batch-after-cancel diagnostic it was written to be.

Gating on ``is_duplex_frame`` rather than on ``req.client is None``
is deliberate: in process mode (P4) the engine side holds no
``ClientRequest`` for ANY request — the applier lives API-side — so a
``client``-based test would silence every request's output there.

``getattr``, not a bare read: the emission seam is reached with
minimal CPU-test ``Request`` stand-ins (``test_distributed_driver``'s
SPMD commit fakes, the teardown drain harnesses) that model only the
fields they exercise. A stand-in that does not model duplex-ness is
not duplex — the same convention, for the same reason, as
``DuplexLane._prepare_tick_inner``'s ``duplex_tick_due`` read. A real
``Request`` always defines the property, so the default never applies
to one.

Emit one :class:`TokenOut` carrying ``output_text`` progress.

The single text-lane emission point: computes the delta between the
engine's ``output_text`` (already decoded + stop-scanned by the
caller) and what previous ``TokenOut``s carried
(``req._client_text_len``), then stages the struct on the step's
batch. A stop-string retro-trim that cut BELOW the emitted cursor is
encoded as ``trim_to`` (the applier truncates the client copy before
appending). Every committed text-lane token emits exactly one
``TokenOut`` — even an empty-delta one (unstable UTF-8 boundary,
special token) — so the applier's client-side token count advances
with the stream.

No-op for a duplex-lane request (see :func:`_has_bus_consumer`).

Emit an :class:`AudioOut` with the request's FRESH speech codes.

Reads the omni demux buffer past the ``_audio_codes_emitted`` cursor;
a control token that produced no new code emits nothing (the
consumer has no state change to observe). Engine thread only.

No-op for a duplex-lane request (see :func:`_has_bus_consumer`) — its
speech leaves through the lane's own per-tick TTS path, not here.

Emit the request's terminal :class:`FinishOut` (exactly once).

THE finish representation on this boundary: every terminal
transition — token-carried (stop id / max_tokens / context / stop
string) and non-token (error, cancel, timeout, embed/rerank,
drained-drop) — emits one ``FinishOut``; ``TokenOut.finish_reason``
is never populated. ``output_token_count`` is authoritative here
(it includes audio/control tokens that emit no ``TokenOut``).
Idempotent via ``req._finish_emitted`` so overlapping terminal paths
(a cancel racing a finish) emit a single record.

``reasoning_token_count`` is the boundary index accumulated per token by
:func:`arbi_serve.sampler.thinking_budget.count_reasoning_token` — an int
for a request that entered a ``<think>…</think>`` block (equal to
``output_token_count`` when it was truncated still inside one), ``None``
for every request that never did, so the API omits
``completion_tokens_details`` rather than reporting a misleading zero.

No-op for a duplex-lane request (see :func:`_has_bus_consumer`) — its
teardown is observed through ``DuplexConnection.closed`` /
``finished_reason``, which ``DuplexLane``'s own close paths set. The
``_finish_emitted`` latch is still stamped so the "exactly once"
contract reads the same on every request class.

Hand this step's staged output batch to the consumer — ONE hop.

Called at the single well-timed point per step — right before the
engine thread blocks on the per-step GPU event-sync — so under
``ARBI_ENGINE_OWN_THREAD`` the HTTP loop applies the batch during
the GPU window (GIL released), overlapped and off the per-token
critical path (invariant: one ``call_soon_threadsafe`` per step,
never per token). No-op when nothing is staged (the common idle
tick), so callers invoke it unconditionally on the hot path.
Single-loop / stub engines apply inline through the same seam.

Engine hot path: async step variants + spec-mode routing.

The event-loop-friendly K=1 step (:func:`step_async`, INLINE in steady
state / OFFLOAD under admission pressure), the MTP seed-draft step
(:func:`step_mtp_seed_async`), the deferred MTP verify drain
(:func:`drain_pending_verify`), the host-ahead pipeline boot + per-step
gates (:func:`_async_schedule_enabled`, :func:`_pipeline_step_eligible`,
:func:`_flush_deferred_finishes`), and the slate spec-mode / K pickers
(:func:`_spec_mode_for_slate`, :func:`_step_K_for_slate`,
:func:`_mtp_seed_needed`).

Pick the :class:`SpecMode` for a slate based on engine config + rows.

When the engine has an MTP driver attached AND any row in the
slate opts in to MTP, the slate runs through :class:`MtpStrategy`
(mode ``MTP``). Otherwise the no-spec-decode path (mode ``NONE``).

Returns the resolved :class:`SpecMode` enum value.

Pick the per-step ``K`` for the plan.

For ``SpecMode.NONE``, K=1 always. For MTP, the scheduler bucketed
the slate to a uniform ``mtp_k`` at admission (see
:meth:`Scheduler._bucket_mtp_k_uniform`), so any opted-in row's
``mtp_k`` is the step K. ``K`` here is the verify-pass spec K
(the number of speculative tokens proposed); the verify forward's
actual T-per-row is ``K + 1``.

True iff this step must run the MTP seed-draft forward
(:func:`_run_mtp_seed_forward` — hidden-state path,
``return_hidden_state=True``; its whole-forward portion replays a
hidden-retaining captured prefill rung when one matches).

Required exactly when some MTP-opted row SAMPLES a token this step: a
non-prefill row, or a prefill row on its FINAL chunk
(``prompt_consumed + n >= len(prompt_token_ids)`` — the ``_sample``
row-readiness predicate). A mid-prefill chunk samples ``None`` and is
excluded from the seed-draft ``ready_rows``, so its hidden state is
never consumed; routing it through the seed forward bypasses every
captured-graph rung (the whole-forward prefill replay in
``_forward_or_replay``) for no benefit. Slates where no opted row
samples take the plain ``execute()`` path — the same forward the
MTP-off config runs.

Tap-conditioned drafters (``observe_seed_forward``, e.g. DFlash)
consume EVERY prefill chunk's forward through the seed path — any
opted row keeps the slate on it.

Wait, in the worker thread, for this seed step's launches to complete.

The wait is the whole point of the offload: ``run_in_executor`` returns
the loop its control only when the worker callable does, so waiting HERE
is what keeps the loop free across the compute. The bare launches return
in microseconds and would hand the loop back before the GPU had started.

An event on the CURRENT stream, not ``torch.cuda.synchronize``. Both
release the GIL while they wait, so the loop is served either way; the
device barrier additionally waits on every other stream on the card — the
async-output copy stream, the recurrent-savepoint spill, NCCL side
streams — none of which this step is waiting for, and each of which is a
stall a faster forward would expose. Nothing downstream depends on the
difference: the commit that follows rests on stream ordering and on the
copies it drains itself, never on this wait, which is a scheduling device
and not a barrier.

The targeted wait is what the verify side of the same offload already
does (:mod:`arbi_serve.spec_decode.mtp_verify_offload`, and
``verify._batched_tolist``); the seed side is the copy that never got it.

Event-loop-friendly MTP seed-draft step (mirror of :func:`step_async`).

The MTP-opted seed step (prefill / K=1 decode that feeds the NEXT
verify step's drafts) runs a main-model forward + sample PLUS the
bundled seed-drafter chain — the SAME ~tens-of-kernel-launches GPU
section that, run inline on the asyncio loop, starves the FastAPI
intake / detok / SSE coroutines. ``step`` runs it inline; this routes
the GPU section through the forward executor so the loop stays free.

Mirrors :func:`step_async`: the batch + metadata build stays on the
loop thread (metadata builders allocate persistent pool tensors that
must not be inference tensors), then the forward + sample + drafter
chain (:func:`_run_mtp_seed_forward`, followed by
:func:`_drain_seed_launches` so the GPU completes in-thread) runs in
``eng._forward_executor`` — the single-worker thread the verify forward
+ K=1 path share, so the single-threaded-CUDA invariant holds. Falls
back to inline :func:`step` when offload is disabled or no executor is
wired. The commit + ``post_token`` tail runs back on the loop.

Routes a slate with NO MTP-opted row — or one where no opted row
samples this step (mid-prefill chunks; :func:`_mtp_seed_needed`) —
through :func:`step_async` (the plain K=1 offload), so a mixed legacy
sub-slate still benefits and mid-prefill chunks keep the captured
whole-forward prefill replay rung.

Materialize + commit the one in-flight deferred MTP verify step.

The spec-decode analogue of :func:`drain_pending_output`. Called at
the head of :meth:`MtpStrategy.run_step` — BEFORE that step's plan
build and forward dispatch, because the plan reads each row's
corrected ``output_token_ids`` / page-table length, which exist only
once the real accept count has landed. So unlike the K=1 drain (which
runs after step N+1's forward is already in flight —
``_commit_step_result`` / ``pipelined_step_async``), this one has
nothing enqueued behind it: every microsecond spent here is host time
the next forward dispatch is waiting on.

What that costs depends entirely on WHEN the acceptance copy was
issued, which is why :func:`start_verify_copy` issues it at the point
the acceptance becomes final rather than at park. Issued at park, the
copy stream's ``wait_stream(compute)`` chains it behind the rest of
the step — the drafter chain and, on a recurrent model, the
partial-accept rollback — and the event wait below becomes a wait
for that whole GPU tail. Issued early, it is a wait for a
four-integer copy that has long since landed.

Also called on the empty slate and on shutdown. No-op when nothing is
pending.

The parked step OPTIMISTICALLY appended ``K + 1`` placeholder tokens
per row (assume-all-accepted) and left every draft slot allocated.
Here we:

  1. block on the copy event + materialize the device ``n_accepted``
     ``(B,)`` and ``main_argmax`` ``(K+1, B)``;
  2. per row, TRUNCATE the ``K + 1`` placeholders, then re-run the
     exact synchronous commit: free the ``K - n`` rejected draft
     slots, append the real ``main_argmax[0..n]`` committed tokens,
     and ``post_token`` each (breaking on a mid-run finish — the
     identical stop semantics the sync path has);
  3. update the per-request + driver accept counters;
  4. drop the deferred output of any row already FINISHED at drain
     time (mirrors the finished-row drop in
     :func:`drain_pending_output` — no max_tokens+1 / stop+1 leak).

Invariant: ``post_token`` (which streams + runs finish checks) runs
ONLY here, after the optimistic placeholders are corrected — the
speculative tokens never reach the consumer.

Record the step's GPU tail, split at the drafter/reconcile boundary.

Reports the PREVIOUS step's events, not this one's, and that is the
whole point. Timing the current step here can only be done two ways,
both wrong: read the events unconditionally and they have not retired,
or read them only when they HAVE retired and the sample is selected by
the very property being measured — the short tails, or the steps whose
host side happened to run long. Either way the mean is not the mean.
Waiting on them instead would re-create the stall this drain exists to
avoid, and would change the timings of every pair that follows.

One step later both events have certainly retired, so every step is
sampled and none is waited on. ``tail_unretired_one_step_later`` fires
if that assumption is ever false, so a lost sample announces itself
rather than silently thinning the distribution.

Emits ``gpu_ms_accept_to_drafts`` (the drafter chain) and
``gpu_ms_drafts_to_park`` (the recurrent reconcile) — the split that
says which of the two is worth attacking — plus their sum.

True iff the scheduler has requests queued for admission this step.

Drives the adaptive forward-offload gate (:func:`step_async`). "Pressure"
= a request the event loop should be free to tokenize / admit while the
GPU forward runs: anything in ``scheduler.waiting`` (newly submitted,
not yet admitted) or ``_deferred_admits`` (admission rolled back last
step, retried this step). Both are cheap ``deque`` truthiness checks on
the hot path. ``suspended`` (offload-parked batch jobs) is deliberately
excluded — resuming those is the scheduler's job, not an event-loop
admission the thread-offload would help.

Event-loop-friendly K=1 step — INLINE in steady state, OFFLOAD under
admission pressure.

With ``ARBI_ENGINE_OWN_THREAD`` (default ON) the engine run loop runs on
its OWN OS thread, so the HTTP/SSE loop is a structurally-free second
thread for the connection accept + tokenize. But the engine loop's OWN
work — admitting a queued request (``asubmit`` → ``scheduler.add`` +
``_wakeup.set``, marshaled onto the engine loop via
``call_soon_threadsafe``) and flushing the per-token detok / SSE wakeups
of the CURRENTLY-DECODING requests — only runs when the run-loop
coroutine YIELDS. A synchronous ``step()`` does not yield until it
returns: while it runs, those marshaled callbacks are blocked.

For a steady-state decode step that is fine — the forward only enqueues
kernels (async output keeps sampling on GPU), so it returns in
microseconds and the per-step ``sleep(0)`` services the callbacks
promptly. But a step that carries a FRESH PREFILL runs a fat GDN/Mamba
recurrent forward (up to ``max_batched_tokens`` tokens, EAGER above the
top captured cudagraph rung) whose Python-side kernel-launch sequence
holds the engine thread (and the GIL) for tens of ms. Run inline, that
one step STALLS the admission of every queued request AND the streaming
of every concurrent decode: a prefill wave monopolizes the engine
thread while decodes and new arrivals pile up.

Adaptive policy: run INLINE (back-to-back kernel enqueues keep the GPU
stream saturated, no ``run_in_executor`` reschedule + launch bubble — the
steady-state-throughput-optimal path) UNLESS the scheduler has requests
queued for admission (:func:`_has_admission_pressure`). Only THEN offload
the forward onto the single-worker ``_forward_executor`` so the run-loop
coroutine suspends on the executor future — yielding the engine loop to
drain the queued admits and flush the in-flight decodes' tokens DURING
the fat forward instead of after it. A request arriving mid-step is
caught by the NEXT step's check (it is in ``waiting`` by then), so at
most one inline step precedes the offload — negligible vs the stall it
removes. The selection is purely WHICH THREAD enqueues the (identical)
GPU work: token-for-token bit-identical either way.

Concurrency invariant: the run loop ``await``s this coroutine before the
next step, so AT MOST ONE forward thread is in flight — the
single-threaded-CUDA invariant holds. ``_forward_in_flight`` is set
around the executor call so the scheduler's admission-time offload
self-heal defers (does not park a resident job mid-forward). The commit +
``post_token`` tail runs back on the loop thread (queue puts / event sets
must not run off-loop).

The MTP-opted path stays inline here — it threads per-request GPU tensors
(``mtp_next_draft_probs``) + the seed-drafter chain across the forward,
which has its own offload (:func:`step_mtp_seed_async`). The MTP
seed/verify slate is routed through that path, not this one.

Boot-level gate for the ARBI_ASYNC_SCHEDULE host-ahead pipeline.

Resolved once per engine (cached on ``eng._async_schedule_gate``) and
logged once. Declines — falling back to the serial loop, never
crashing — on: flag OFF (default), non-``generate`` task, TP>1 (the worker
bridge's per-step mirror ordering under a two-thread overlap is
unaudited), ``async_output`` OFF (the whole premise is the un-synced
in-flight forward), or a missing forward executor.

Per-step gate for the host-ahead pipeline (cheap, hot-path).

The caller (:func:`run_step_async`) already established
``SpecMode.NONE`` (no MTP verify routing). Beyond the boot gate,
decline when a deferred MTP verify step is parked (its drain is
entangled with live GPU state and must not run overlapped), when an
EXTERNAL drafter is attached (DFlash routes every slate — including
NONE-mode ones — through its own strategy for tap/slot bookkeeping;
bypassing it would silently drop drafting), or when the runner's
async output is off (a synchronous-sample execute would return only
after draining the GPU — nothing to overlap). MTP-opted rows on
their seed step keep the dedicated :func:`step_mtp_seed_async` path.

Execute the finishes queued during the pipeline's overlap window.

Runs on the engine thread AFTER the in-flight Entries are executed in
arrival order (per-request token order was already preserved by the
drain that queued them). The queue is consumed destructively so a
re-entrant flush (error path finally + normal path) is a no-op.

Engine hot path: per-token post-processing + streamed-finish terminal.

:func:`post_token` is the per-token finish-condition checker (stop id /
max_tokens / context / stop string), the omni demux route, and the
detok / streaming hand-off. The streamed-finish terminal sequence
(:func:`_finish_streamed_request` → :func:`_terminal_detok_handoff`)
renders the final token and emits the ``TokenOut`` + ``FinishOut`` pair.

Detokenizer hand-off for a token that FINISHED its request.

Worker path (``detok_queue`` present): enqueue with the FINISH
sentinel — the worker (running on the engine loop) decodes the final
token, emits its ``TokenOut`` and the terminal ``FinishOut``, and
exits. The applier applies text before finish, so readers observing
``finish_reason`` always see the complete text — no drain needed.

Inline path (``detok_queue is None``): when ``post_token`` already
rendered the final token's text (``pre_decoded`` — single decode per
token; it also keeps the ENGINE-side text complete before
``derive_finish_reason`` / the scheduler finish observation point),
re-entering ``_decode_inline`` here would decode nothing — but a
NON-idempotent test stub tokenizer would emit the last char twice,
so skip it. Then emit the final ``TokenOut`` (carrying the
pre-decoded delta) followed by the ``FinishOut``; both ride this
step's batch, and the applier makes the finish visible only after
the final text (the publish-order invariant, now structural).

Diagnostic: one-line summary of a finished speech-output request.

Gated by ``ARBI_DEBUG_OMNI_STREAM``. Surfaces the interleaved-stream shape
(text vs code counts, tts_start/tts_end offsets, EOS presence, tail ids) so
the audio/text desync + runaway failure modes can be told apart.

Terminal sequence for a request finished by :func:`post_token`.

The streamed-finish variant (stop-token / max-tokens / context). The
caller has already stamped ``finish_reason``. Frees KV via
:meth:`Scheduler.finished`, marks FINISHED, releases the LoRA ref,
drops the request from the live map, runs the terminal detok hand-off
(:func:`_terminal_detok_handoff` renders the final token and emits the
``TokenOut`` + ``FinishOut`` pair), then emits finish metrics.

Differs from :func:`_finish_request`: the final token rides the detok
hand-off rather than a plain ``cancel_detok_task`` +
:func:`_emit_finish`, and the request is dropped from the map BEFORE
the hand-off — the exact order the three ``post_token`` finish
branches shared.

Host-ahead pipeline (ARBI_ASYNC_SCHEDULE): while the overlapped drain
runs (``eng._deferred_finish_queue`` is set), the WHOLE finish is
QUEUED and executed after the in-flight ``execute`` returns — freeing
the request's pages / recurrent slots here would race the runner
thread's forward dispatch, which still reads the slate's
``state_handles`` and stamps ``last_sampled_gpu``. The terminal
``TokenOut``/``FinishOut`` emission rides the deferral too (it happens
inside the hand-off), so the consumer observes the finish when the
pipeline flushes — at most one forward dispatch later.

Process a freshly-sampled token.

The token-id-only finish checks (``stop_token_ids``, ``max_tokens``,
``context``) are decided synchronously here — they don't need
decoded text. The actual

Output crossing: this function emits typed structs (``TokenOut`` /
``AudioOut`` / ``FinishOut``) onto the step's output batch; the
OutputApplier feeds the consumer-side ``ClientRequest`` (text copy,
token count, audio lane, finish) and fires ``new_token_event``. The
engine-side ``req.output_text`` remains the detok / stop-scan
working buffer only.

When no detok worker has been started for the request (the inline
default and the sync test path), :func:`enqueue_token` runs the
decode inline on the engine thread and the ``TokenOut`` is emitted
here after the stop-scan.

Runtime add-resident + one-surface switch-or-load for stable-VA residency.

Gives a model loaded AT SERVE TIME (never named at launch) the same fast
park/wake swap-back a boot ``--pool-member`` resident gets. Two capabilities
reconciled behind one call:

  * :func:`aswitch_or_load_model` — "switch to model X, loading it if absent".
    Present resident ⇒ the existing fast park/wake (:meth:`Engine.aswitch_model`);
    absent ⇒ :func:`aadd_resident_member` builds it into the residency, then it
    is the active resident. A subsequent A↔B swap is the fast park/wake, NOT a
    full :func:`arbi_serve.engine.swap_admin.areload_model` reload.
  * :func:`aadd_resident_member` — the runtime analog of ONE
    :func:`arbi_serve.engine.stable_va_pool_builder.prepare_model_pool` member
    step: park+evict+drop-weights the active member (cross-model park), build the
    new DISTINCT model as the sole VRAM resident under its tag namespace,
    register it, and make it the active resident (its weights stay mapped; the
    outgoing member's are on its flat dump). Runs under
    ``eng.critical_section(ALL_KINDS)`` with the TP>1 cross-rank commit/abort
    barrier so an asymmetric build outcome rolls back symmetrically.

Residency is engaged LAZILY (:func:`_ensure_residency_engaged`): a single-model
boot never constructed the controller, so the first add wraps the running model
as the boot record and starts the graceful queue that drives request-triggered
resident swaps.

Engage stable-VA residency for the running model if it is not already.

A single-model boot (no ``--pool-member``) never constructed the controller.
Wrap the currently-serving model as the boot resident record and start the
graceful queue that request-triggered resident swaps route through — so the
first runtime add turns the single-model engine into a residency with two
members that A↔B swap via fast park/wake. Idempotent: returns the existing
controller untouched (queue already running) when residency is engaged.

Fails loud under ``--no-cumem-pools`` (``enable_stable_va_residency`` raises):
park/wake unmaps physical behind stable VAs through the cuMem allocator, so
there is no residency without it — never a silent degrade to a full reload.

Resolve the request to an ALREADY-registered resident key, or None.

Tries, in order, the explicit ``key`` / ``served_name`` / the ``path``
against the resident set. None ⇒ absent, so the caller builds it.

Each candidate goes through :func:`arbi_serve.engine.residency_view.
resolve_resident_key` — the SAME identity map the request router resolves a
``model`` field with (minus the ``"default"`` alias, which that layer adds:
a load names a concrete model, never "whatever is resident"). Sharing it is
load-bearing in both directions: a ``path`` that names an already-registered
resident must resolve to a SWITCH, because falling through would register a
second resident for a checkpoint already resident — the boot member is keyed
by its served NAME, so a load naming the boot ``--model`` path is exactly
that case.

Derive the new member's ``(cfg, flag_overlay)`` from the boot config.

Collapses the admin aliases (``path`` / ``calibration_path`` /
``attention_backends`` / ``default_backends``) and the per-member drafter
source (``enable_mtp`` / ``mtp_assistant_path`` / ``mtp_draft_model_path`` /
``mtp_draft_model_dtype`` / ``mtp_draft_model_max_seqs``) into ONE sparse
delta and applies it through the SAME overlay machinery a boot
``--pool-member`` and a live ``config_override`` use
(:func:`arbi_serve.config_overrides.apply_overrides`) — cross-arch semantics
(an explicit ``calibration=""`` ⇒ uncalibrated member) resolve identically.
Stamps the member's ``served_name`` (its routing key).

DISTINCT-model member: the boot model's per-model drafter SOURCE is cleared
from the base so the runtime-loaded model names its OWN assistant / draft
model / DFlash (or none); its ``mtp_*`` fields layer on top. Fail loud on a
bad drafter combo via :meth:`MtpConfig.validate_source`.

Build a DISTINCT model into the residency at runtime + make it active.

The runtime analog of one :func:`prepare_model_pool` member step. Under the
engine critical section (scheduler drained + paused) so no forward runs
against half-mapped physical during the park→build window:

  1. Park the active member CROSS-MODEL — discard its growable KV (drained;
     re-prefills on a later wake) and drop its weights to its flat dump (a
     distinct model's weights cannot co-reside), then evict its residual pool
     physical to host, so the new model builds as the true sole VRAM resident.
  2. Build the new model into the primary engine under its tag namespace
     (:func:`build_member_into_engine`), register it, and — on a positive
     cross-rank verdict — make it the active resident. Its weights stay
     mapped (it serves now); the outgoing member is parked with weights on
     flat, so a later switch back is the fast park/wake.

TP>1: every rank reaches this deterministically from the same broadcast op
and runs the identical sequence; the build's NCCL collectives rendezvous by
construction and a cross-rank commit/abort barrier
(:func:`arbi_serve.engine.member_build_barrier.confirm_build_across_ranks`)
makes an asymmetric build outcome roll back on every rank. TP=1 degenerates
to "local outcome is the verdict".

Drop a failed/aborted runtime add + re-wake ``prev_active``.

Transactional: leave the engine serving ``prev_active`` exactly as before the
add. Discards whatever the build mapped (orphan cudagraphs, released engine
state, the namespace-tagged physical, the registry record) before re-waking
``prev_active`` — otherwise its KV remap has no room. ``wake`` restores
prev's snapshot (incl. its cfg + weights-from-flat) so the engine is coherent.
Mirrors :func:`arbi_serve.engine.config_variant._rollback_member_build` minus
the config-variant base-cfg restore (the wake restores prev's cfg).

Switch the resident model to the named model, loading it if absent.

The one coherent runtime-model surface:

  * PRESENT (already a resident) ⇒ the existing fast park/wake
    (:meth:`Engine.aswitch_model`); a no-op when it is already active.
  * ABSENT ⇒ :func:`aadd_resident_member` builds it into the residency and
    makes it the active resident (add-then-wake).

Engages residency lazily on the first call (:func:`_ensure_residency_engaged`)
so a single-model boot gains the capability without a restart. This is the
fast path; the destructive full reload
(:func:`arbi_serve.engine.swap_admin.areload_model`) stays the explicit
fallback for a topology change a park/wake cannot express.

Boot realizes the widest serving step, so the KV sizing reads what it leaves.

The serving grow floor holds free VRAM back from KV. Its terms price what a
step allocates through a pool we own. The CUDA caching allocator also takes
fresh segments OUTSIDE every named pool the first time a step presents a block
size no boot phase realized, and those segments come out of the same free VRAM.
Nothing that hand-builds a batch can name them: which block sizes a step
presents is decided by :func:`~arbi_serve.runtime.batch_build.build_batch` and
by the kernels the real step path dispatches, and no boot phase went through
either.

So boot goes through both. This module drives ``max_batch`` synthetic requests
through the SAME entry the run loop drives -- :meth:`Scheduler.add`,
:meth:`Scheduler.schedule`, :func:`~arbi_serve.engine.run_step.run_step_async`
-- which routes through the engine's spec-decode strategy and materialises
every batch through the shipped builder: ``build_batch`` for a slate,
``reconstruct_verify_batch`` for an MTP verify plan. The batch a served step
presents is therefore the batch the realization presents, and the allocator
takes its segments HERE, before the one ``mem_get_info`` the KV grow sizes
from. :func:`_batch_materialize_counter` is the receipt.

MEASURED UNCENSORED, WHICH IS THE PART THAT COULD NOT BE DONE WHILE SERVING. A
reading taken from inside the floor cannot exceed the floor it is measured
against: it saturates, and a saturated reading reads exactly like a fitting one
while the card runs at zero free. The drive here has the whole post-capture card
free, so what it reports is what a step takes.

The drive then hands back what the caching allocator holds in freed blocks. That
overhang is not a requirement -- it is what an allocator with a whole card free
never had to re-use -- and leaving it resident would spend KV pages on segments
no step asks for. What cannot come back (a cubin the driver loaded, a pool that
mapped and holds, state the engine keeps) stays, and is outside the free reading
by construction.

DRIVEN AGAINST THE CARD THE GROW WILL LEAVE, which is the second half of the
same argument. A drive with the whole card free is a drive the caching
allocator never has to re-use inside: it answers every new block size with a
fresh segment, and what it takes is what an empty card permits rather than what
a step needs. Serving runs against a card whose free VRAM is the grow floor and
nothing more. So the sizing is a CONVERGENCE, not a single reading
(:func:`realize_under_serving_pressure`): measure, work out the KV size that
reading buys, put the card into exactly that state, and drive again. The state
is reproduced with the SAME physical the grow maps -- a cuMem
:class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion` holding the bytes
the extra KV pages would hold (:func:`_pressure_ballast`) -- so the drive is not
approximating the post-grow card, it is running on it.

NO SIZE IS KEPT THAT WAS NOT DRIVEN AT ITS OWN PRESSURE. Each pass proposes a
page count and the NEXT pass is the proof: it drives with that many pages'
worth of physical held. A pass that cannot complete leaves the last page count
a drive did stand at, so the loop can only ever hand back a size some drive in
this process survived. That is the never-OOM side of the ledger; the pages the
convergence recovers are the never-overbudget side, and neither is a constant.

PARAMETRIZED ON THE STEP SHAPE, because that is what the quantity is a function
of. The rows, the token budget, the prefill chunk, the speculation depth and
the sampling routes are read from the live config
(:func:`realization_shape`), so a boot at a different shape realizes a
different quantity by construction rather than by a key that has to remember to
carry the shape.

WHAT IT DOES NOT PRESENT is as load-bearing as what it does. The realization
issues plain text generation: no logprobs, no media, no speech. Those consumers
keep their own additive serving-floor rows; this measurement is MAX-ed against
them, never substituted for them.

The step shape the realization presents, and its every input.

Every field is read from the live config. Two boots that agree on all of
them realize the same step; a boot that changes any of them realizes a
different one, which is what makes the measured quantity a function of the
configuration rather than of which process happened to run first.

What ``scratch.forward_arena`` costs SERVING, measured at the KV seam.

The pool's reserved reading at this seam is the boot's own realization --
the activation probes, the verify-width probe, the drafter pre-sweep -- and
a private cuMem pool never gives a freed block back, so all of it is
physical the KV grow is never offered. The serving floor's arena row prices
that residency as the cover for what a served step re-maps
(:func:`~arbi_serve.engine.inprocess_capture.forward_arena_regrow_plan`),
and that is a CLAIM about the boot mark, not a measurement of what serving
needs.

This is the measurement. The pool is handed back to the driver, and then
the realization drives the widest step the scheduler can issue through
``build_batch`` at the pressure the grow will leave -- so what it re-maps is
the requirement, in the pool state the release ships, rather than the boot's
shape mix inherited. ``serving_bytes`` is the high-water over every pass;
``boot_bytes - serving_bytes`` is boot residue the card was holding.

NOT A RESERVE AND NOT A FORECAST. Nothing here is carried to another
process or another boot, and no margin is added in either direction: the
convergence loop re-drives at each proposal's own pressure, so a page count
only comes back at a size some drive in this process stood at.

What the boot realization took, and whether it ran at all.

``bytes_`` is free VRAM that is GONE -- spent by the realization and still
resident when the grow reads free. It is a receipt, not a reserve: the
floor does not hold it, because it is already outside the free VRAM the
floor divides up.

Why the realization must NOT run here, or ``""``.

A realization is a real serving step. Under TP>1 every one of its linears,
its vocab-parallel embed and its lm_head all_reduce with peer ranks that
are inside their own ``build`` and are not receiving mirrored batches, so
issuing one is a half-issued collective at boot -- an NCCL wedge, not an
exception. The same refusal :func:`~arbi_serve.engine.inprocess_capture.
_dflash_draft_step_collective_refusal` makes for the draft-step probe, for
the same reason and with the same consequence: the floor keeps the
pre-measurement bound there and the row says so.

The widest step shape the scheduler can issue on this configuration.

Read from the config the scheduler itself reads, so the realization
presents the step admission will admit rather than one chosen here.

Distinct in-vocab prompt ids for one realization row.

DISTINCT PER ROW: identical prompts share a radix prefix, and a slate whose
rows dedupe to one prefix presents a narrower step than ``max_batch``
independent requests do. The ids are a fixed arithmetic walk of the vocab
rather than a sample, so two boots at one configuration realize the same
step; content does not decide the shape, and no realization row is ever
read by anyone.

Bracket one drive of the widest step between two free-VRAM readings.

Called with whatever pressure the caller has already put on the card, so
the bracket is a reading of the drive and of nothing else: the ballast is
already held when ``free_before`` is taken and is still held when
``free_after`` is.

Hold ``nbytes`` of device physical the way the KV grow is about to.

Yields the bytes actually held, which is 0 when nothing was asked for and 0
when the driver could not give it -- the caller must treat a short hold as
a pass that did not reproduce anything, because a drive that ran on a
roomier card than it claimed is the exact defect this module exists to
remove.

THE SAME PHYSICAL, NOT A STAND-IN. The bytes come from a cuMem
:class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion` -- the primitive
the paged-KV slab itself is built on
(:mod:`arbi_serve.cache.paged_kv_pool`) -- mapped at its own reserved VA and
outside every torch allocator. So the card the drive runs against differs
from the card serving runs against in which VA the pages sit at, and in
nothing the caching allocator can see: a segment it cannot get here is a
segment it will not get there.

Reversible, which is what makes it usable as a PROBE. Growing the real slab
would put the card into the same state, but the slab has no shrink -- a
proposal that turned out not to hold could not be taken back, so the boot
could only ever discover an over-reach by serving it. The region is unmapped
and closed on the way out of the ``with``, including on the raising path.

``(region, held_bytes)``; ``(None, 0)`` when nothing was or could be held.

A partial hold releases through the caller's ``finally`` like a whole one:
the region is returned even when the mapping stopped short, and ``held`` is
reported as 0 so no pass mistakes it for the card it asked for.

``scratch.forward_arena``'s reserved physical, or ``-1`` when unreadable.

The same reading :func:`~arbi_serve.engine.inprocess_capture.
arena_resident_reserved_bytes` takes, through the same pool, and it is
deliberately the SAME quantity: the mark this module measures during a
drive and the mark the floor prices the arena row against have to be
comparable or neither says anything about the other.

``-1`` rather than 0 because a boot with no arena pool and a boot whose
arena holds nothing are different facts, and only the second is a bound.

Drive the widest serving step through the real path; keep what it takes.

Returns what free VRAM the realization spent and did not give back. Runs
between :func:`~arbi_serve.engine.inprocess_capture.
settle_vram_for_kv_sizing` and the KV grow's single ``mem_get_info``, so
the bytes are outside the free reading the grow divides up.

``ballast_bytes`` is device physical held for the drive's whole duration,
so the free VRAM it runs against is the free VRAM the KV grow is about to
leave (:func:`_pressure_ballast`). At 0 the drive has whatever the card
has, which for a post-capture boot is most of it -- and an allocator with
most of a card free answers a new block size with a fresh segment instead
of re-using one, so that reading is an UPPER BOUND on the requirement
rather than the requirement. :func:`realize_under_serving_pressure` is what
turns the bound into the reading.

Never raises. A realization that cannot complete is a realization that did
not measure, which the floor must be told about rather than have hidden
behind a partial reading: the return carries the reason and the floor
holds the pre-measurement bound. THAT INCLUDES AN ALLOCATOR REFUSAL under a
held card: a drive that could not fit is the proposal being refuted, and
the caller keeps the last size a drive stood at.

Run one step coroutine to completion, blocking this thread on it.

On its OWN thread with its OWN loop, because this drive is reached two
ways. At boot the calling thread has no loop. On a LIVE member rebuild it
is the engine's own loop thread, blocked on the synchronous build — and
that block is load-bearing: it is what keeps the engine's step loop off the
scheduler while this drive admits and runs rows of its own. So the step
cannot re-enter the outer loop (asyncio refuses) and must not be handed
back to it (it would drive the scheduler alongside this drive). A thread we
immediately join gives the coroutine a loop to run on and resumes nothing:
the outer loop stays exactly as blocked as it was.

Run the realization's rows through the run loop's own step sequence.

THE LOOP BODY IS THE RUN LOOP'S, down to the dispatch. The scheduler builds
the slate and :func:`~arbi_serve.engine.run_step.run_step_async` routes it,
which is what picks the spec-decode strategy and, for an MTP-opted slate,
the verify pass and its offloaded recurrent rollback. The synchronous
:func:`~arbi_serve.engine.run_step.step` reaches ``build_batch`` too but
never reaches those: it is the seed / legacy sub-slate body, and a
realization driven through it presents a step the engine does not serve —
GPU-measured, by the served step then OOMing in the rollback's masked
commit on memory the realization had told the grow was not needed.

Nothing about the step body changes: :func:`_run_to_completion` only gives
the coroutine somewhere to run, and it works from a booting thread and from
inside a live member rebuild alike.

Count batch materializations through the shipped builder for the drive.

NULL CONTROL. The realization's whole claim is that the batch it presents
is the batch serving presents, and the only evidence for that is that it
went through the same materializer. A drive whose count comes back 0
realized something else, and :func:`realize_serving_step` refuses to hand
the KV sizing its reading.

Reads :data:`~arbi_serve.runtime._batch_build_helpers.BATCH_MATERIALIZED`,
which BOTH entries fire: ``build_batch`` for a slate and
``reconstruct_verify_batch`` for a verify plan. A counter on the slate
entry alone reads 0 for every MTP verify step and would call the drive
unmaterialized on exactly the configuration this floor is sized for.

Give the realization's rows their pages back and leave no trace.

``scheduler="remove"`` rather than ``"finished"``: a finish commits the
row's full pages into the radix prefix cache, and the realization's prompts
must not be reachable by a later request.

What one pressure pass settled about the page count it was given.

The pass is a PROOF OBLIGATION, not a measurement that happens to be
smaller: it holds the physical ``candidate_pages`` would hold and drives
the widest step against what is left.

``"REFUTED"`` — the drive did not complete on that card, so the proposal is
withdrawn and the caller keeps the last page count a drive did stand at.
``"CONVERGED"`` — it completed, and its reading buys no more pages than the
count just proven; that reading can be kept, because the card it leaves is
at least as roomy as the one it was measured on.
``"GROW"`` — it completed and its reading buys MORE. The larger count is a
proposal like the one before it, and the next pass is its proof; nothing
hands it back unproven.

Hand ``scratch.forward_arena`` back to the driver before the first drive.

Returns the mark with its ``boot_bytes`` read and the release recorded. The
serving high-water is filled in by :func:`_close_arena_mark` once the passes
have run.

The release is :func:`~arbi_serve.engine.inprocess_capture.
release_idle_forward_arena` with ``keep_bytes=0`` -- the SAME
destroy-and-recreate the serving re-tightening uses, with the same live-set
refusal, so a pool holding anything a captured graph may have baked is left
alone here exactly as it is there. A refusal is recorded and the boot then
behaves as it does today: the pool keeps every byte, and the mark reports
itself unmeasured rather than reporting the boot's own number as if a drive
had produced it.

Fill the serving high-water in from the passes that ran, and say it.

The high-water is over EVERY pass, not over the one whose reading is kept:
a pass that was later withdrawn still mapped what it mapped, and the pool
still holds it.

Converge on what the widest step takes from the card the grow will leave.

The first pass has the whole post-capture card. It is an UPPER BOUND and
the loop treats it as one: an allocator that never has to re-use answers
every new block size with a fresh segment, so what it takes is what an
empty card permits. Serving never runs there -- it runs against free VRAM
equal to the grow floor, where the same step re-uses instead.

So each pass works out the KV size its reading buys, puts the card into
exactly that state (:func:`_pressure_ballast` holds the physical those
extra pages would hold) and drives again. The reading that comes back is a
reading of the state that will actually exist. NO PAGE COUNT IS RETURNED
THAT A DRIVE DID NOT STAND AT: a pass whose drive cannot complete refutes
its own proposal and the last proven reading is handed back instead
(:func:`pressure_verdict`).

``max_passes`` is a BOOT-COST BUDGET, not a correctness knob. Every value
returns a page count some drive stood at -- a budget that runs out while
the size is still growing hands back the last proven one and says so, so
lowering it costs pages and never safety.

Returns the reading the serving floor should hold. The page arithmetic is
:func:`~arbi_serve.engine.memory_budget.kv_budget.kv_grow_target_pages` and
the floor is :func:`~arbi_serve.engine.inprocess_capture.
serving_floor_for_grow` -- the same two the grow itself uses, called here
rather than re-derived, so a proposal and the grow that follows it cannot
be two different numbers.

Release what the drive left, through the sizing seam's own release.

A second, partial copy of "every release that can happen" here is how
the settle and the drive come to disagree about what was released.

The pool's own price list for a page target, or None if it has none.

Read from the pool on every use rather than captured once: the same
seam the grow prices its target with (``kv_mapped_bytes_for_pages``),
so a proposal and the grow that follows it cannot be two arithmetics.

Physical a drive must hold to stand where a ``target_pages`` pool leaves
the card. The grow's own cost, never ``pages x per_page_bytes`` — the
per-layer granularity round-up makes those two different numbers, and a
ballast short of the real cost measures a step on a roomier card than
serving gets.

Admin-requested process exit, and whether anything will restart it.

A process cannot restart itself; it can only exit and be restarted by
something else. Nothing inside the container can observe its own restart
policy -- that is daemon-side metadata -- so the deployment DECLARES it
(``ARBI_SERVE_SUPERVISED``) and the answer travels with the response. The
server never claims a restart it cannot guarantee; the caller decides what to
call the action.

The exit is deliberately hard. Once the drain has completed there is no engine
state worth unwinding, and the graceful path is the riskier one: teardown frees
every pool, which takes the caching allocator's mutex while a concurrent
``/metrics`` scrape holds the GIL -- the inversion that deadlocked boot
teardown. A dead process cannot deadlock.

The exit is also DEFERRED: this runs in the engine process, which is the
container's main process and the parent of the API child, so exiting here kills
the API before it can write the response. The delay lets the reply reach the
caller first.

Arm a deferred hard exit of THIS (engine) process.

Returns what the caller needs to describe the action honestly: that the
exit is armed, and whether a restart is expected of it.

``exit_code`` distinguishes the two reasons to leave. ``on-failure`` restarts it rather than reading a failure as a
normal stop. The timer thread is the ONLY mechanism used either way: it
touches no CUDA and takes no allocator lock, so it still fires when the
reason for exiting is that the GPU state is unusable.

``reason`` names WHO is leaving in the armed-exit line. An operator
reading the log must not be told an unrecoverable engine fault was an
"admin shutdown" -- the two have opposite implications for whether
anything is wrong.

Sleep mode (phase2 only): drop / restore GPU state behind stable VAs.

cuMem-backed named-pool mode (``cfg.cumem_pools=True``, the default):

- Every named pool (``weights_pool``, ``kv_pool``, ``graph_pool``,
  ``activation_arena``, ``cublas_workspace``, … — the full
  ``engine.CUMEM_BACKED_POOLS`` set) is backed by the cuMem pluggable
  allocator at boot; every allocation in those pools is individually
  cuMemAddressReserve-d so its VA is stable across sleep cycles.
- :func:`release_memory_occupation` calls
  :meth:`~arbi_serve.runtime.named_pool.NamedPoolRegistry.sleep_all`
  which issues one :meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.sleep`
  to offload / unmap every pool's physical pages, then calls
  ``eng.sleep_pool.release()`` for any residual per-tensor
  registrations, and finally de-instantiates every captured graph's
  ``cudaGraphExec_t`` (driver-side staging outside the pools).
  Captured CUDAGraphs survive because every ``data_ptr`` they
  reference is a stable VA and the captured ``cudaGraph_t`` topology
  is kept (``SleepableCUDAGraph``, ``keep_graph=True``).
- :func:`resume_memory_occupation` calls
  :meth:`~arbi_serve.runtime.named_pool.NamedPoolRegistry.wake_all`
  to remap fresh physical pages at the same VAs, re-instantiates the
  graph execs (not a recapture), then calls ``eng.sleep_pool.resume()``
  to restore any per-tensor state.

Sleep-to-floor invariant: with full pool coverage the post-sleep NVML
framebuffer is the CUDA context floor (driver context + loaded cubins +
a small native remnant), further reduced by the thread-stack/local-memory
shrink (:func:`shrink_thread_stack`). The remainder is the bare CUDA
context plus torch/native fatbin kernel images, which have no public
unload API.

The ``phase2`` stable-VA mode is the only supported path: it
preserves cudagraphs across release/resume by construction (a
drop-and-realloc approach would force full re-capture on every wake).

Drop the CUDA per-thread stack limit to free the context's
local-memory pool (sleep-floor shaving).

Only safe while no kernel can launch or replay on the device: eager
launches of kernels whose frame exceeds the limit make the driver
re-grow the pool, but graph replays skip that auto-grow — the
limit must be restored (:func:`restore_thread_stack`) before graph
execs are re-instantiated or any new member is built. Returns the
previous limit in bytes (0 ⇒ not shrunk: flag off / already shrunk /
cuda-python missing / driver error).

Restore the pre-shrink stack limit (no-op when not shrunk).

Must run before graph-exec re-instantiation and before building a new
member while another sits asleep — replays do not auto-grow the stack
and a capture under the shrunken limit would bake it in.

Build the cooperative idle gate for the background flat-weight dump.

Returns a predicate :func:`arbi_serve.loader.flat_dump.dump_flat_weights`
calls between tensors. It blocks the dump while any request is in flight
— so the ~GB D2H stream never contends with live decode through the
shared CUDA context (the warm dump is best-effort latency tuning, never
allowed to stall serving) — and returns ``False`` to request a clean
abort once the engine is terminating (so a paused worker drops the dump
promptly instead of touching GPU state the teardown is about to free).

Engine-wide idleness is read via :func:`inflight_count` — the same
``eng.requests`` accounting the swap drain (``has_inflight``) uses — so
the dump and the swap machinery never disagree on "idle". The read is
lock-free (a ``len`` under the GIL); a stale read at worst defers a
tensor by one poll. Runs on the ``asyncio.to_thread`` worker, so
``time.sleep`` here parks only the dump, not the event loop or the
engine thread.

Dump every model parameter as a single safetensors file.

Run after the engine is built (ideally after warmup). The flat-dump
file at :func:`arbi_serve.loader.flat_dump.default_flat_cache_dir` is
consumed by the boot's warm-load branch for a faster cold boot.

D2H + disk write run in a worker thread so the event loop is not
blocked, and an idle gate (:func:`_make_flat_dump_gate`) keeps the
stream from ever contending with live decode — the dump pauses the
instant a request is admitted and resumes when the engine drains.

Yield every ``torch.cuda.CUDAGraph`` the engine owns, deduped.

Walks every cudagraph pool kind (``eng.cudagraph_pools``: decode /
layer / drafter — plus any future kind added to
the dict) AND every multi-group pool
(``eng.captured_graphs_multi``), introspecting each captured-entry
dataclass for fields holding a ``torch.cuda.CUDAGraph`` (covers
``graph``, ``pre_graph``/``post_graph`` split captures, and any
future graph-holding field without editing this walker).

A field holding a LIST of graphs (the drafter chain's per-step segment
captures) is walked element-wise. A graph missed by this walk keeps its
instantiated exec — driver bytes outside every torch pool — across a
sleep, which is exactly the floor this module exists to lower.

Then asks each owner in :data:`_GRAPH_OWNER_ATTRS` for the graphs it holds
outside any such pool. The structural walk cannot find those: they live in
plain dicts of plain objects, so the owner has to declare them.

Report every live ``SleepableCUDAGraph`` still holding an exec.

The sweeps above are scoped to the RESIDENT model by construction, which is
what keeps a member park from destroying another member's execs. The cost
of that scoping is that an owner the walk does not reach keeps its execs
through the sleep — driver bytes the floor cannot reclaim, and silently,
because a graph nobody visits raises nothing.

The live set is the complete population (construction is the one point
every sleepable graph passes through), so what is still instantiated after
a full sleep is exactly the unreachable remainder. Diagnostic only: it
destroys nothing, because a process-wide destroy is precisely what the
per-model scoping exists to prevent. Returns the count.

Destroy every owned captured graph's ``cudaGraphExec_t`` (sleep).

The captured ``cudaGraph_t`` topology survives — wake re-instantiates
an identical exec (:func:`_reinstantiate_graph_execs`), which is not a
recapture. Graphs that are not :class:`SleepableCUDAGraph` (no
``deinstantiate_exec``) are skipped: torch owns their exec and offers
no destroy-only API, so they keep it across the cycle (logged once).
Returns the number of execs destroyed.

Re-instantiate every owned captured graph's exec at wake.

Eager (not lazy-on-first-replay) so wake latency is paid here,
predictably, instead of on the first post-wake request — and so a
failed re-instantiation surfaces as a wake error, not a serving
error. Returns the number of execs created.

Drives the SAME walk :func:`_deinstantiate_graph_execs` drives, which is
what makes the pair symmetric under a member switch: the controller
repoints the engine's per-model attributes before waking, so this restores
the resident model's execs and leaves a parked member's destroyed.

The active stable-VA record's cuMem tag namespace, or ``None``.

Under default residency every build-time pool allocation is tagged
``"<namespace>/<pool>"``; the sleep/wake paths must target those tags
(raw-name-only targeting silently skips them). ``None`` when residency
is not engaged — allocations are raw-tagged and the legacy targeting
is exact.

Clear every :data:`PARK_BLIND_RUNNER_CACHES` entry. Returns those dropped.

Runs at the head of the park, before anything is unmapped, so the reference
goes while the block's pool still owns it rather than being returned to a
slept pool's free list from the woken member's first step.

Coordinated sync release of one engine's GPU state behind stable VAs.

The synchronous core shared by :func:`release_memory_occupation` (the
single-engine sleep) and the multi-model residency controller (which parks a
prepared member engine — a separate :class:`Engine` whose own
``sleep_pool`` / ``named_pools`` / growable-KV slab must be released
coordinately, not just the raw cuMem namespace). Releases, in the proven
order: registered per-tensor sleep entries (``sleep_pool.release()``) →
cuMem-backed named pools (``named_pools.sleep_all()`` — every named pool
since sleep-to-floor, see ``engine.CUMEM_BACKED_POOLS``) → the B2
growable KV slab (its own VA reservation) → captured-graph exec
de-instantiation (:func:`_deinstantiate_graph_execs`) → ``empty_cache``.

``exclude_pools`` names cuMem-backed pools that must stay mapped — the
residency park passes ``{"model.weights"}`` (a parked member's weights stay
VRAM-mapped: the donor-share invariant, and an active sharer reads them on
every forward).

``offload_sequence_state`` (default ``True``) copies the member's
PER-SEQUENCE state D→pinned-host so a woken member resumes committed
sequences. It governs the two places that state lives behind its own VA
reservations — the B2 growable-KV slab and the recurrent sentinel-alias
arenas — as ONE decision, because they hold one thing between them and
preserving half of it preserves nothing. (The ``state.*`` NAMED pools are
the third place; a light park leaves those mapped and
``StableVaResidencyController.evict_parked_to_host`` applies the same
decision to them.)

A DRAINED park — the caller drains the scheduler first, so no committed
sequence must survive the swap, and it flushes the prefix cache and the
recurrent savepoints — passes ``False`` to discard that state instead.
Discarding is not an optimization: the offload destination is unswappable
pinned host RAM, the state fills most of the card at a high
``gpu_memory_utilization``, and the woken member re-prefills from empty
anyway, so the copy buys nothing and can OOM-kill the host.

``namespace`` is this engine's cuMem tag namespace (stable-VA residency:
every build-time pool allocation is tagged ``"<namespace>/<pool>"``).
Without it the raw-name pool sleep matches ~nothing of a namespaced
member, the park silently keeps GiBs mapped, and a side-by-side member
build (the backend-swap path) OOMs.

Does not drain / take a critical section / mark ``_memory_released`` — the
caller owns ordering + the section. Returns a small info dict.

Bytes the growable-KV re-provision may consume, from the SAME closed KV
ledger the boot deferred resize uses. ``None`` when it can't be computed
(fall back to a full remap). Used only on a discard-parked / cross-model
member's wake.

:func:`arbi_serve.engine.memory_budget.wake_kv_budget` assembles the ledger
from the boot path's named reserve terms — serving floor + the
``cudaGraphInstantiate`` exec reserve + margin — PLUS the two the wake
itself maps after the grow (the sentinel-alias state arenas and the
per-tensor sleep set; see below) and binds the page budget
through :func:`~arbi_serve.engine.memory_budget.member_kv_available_bytes`
(free ∧ ``gpu_memory_utilization``), exactly as the boot member build does,
so a wake can never silently miss a reserve the boot ledger accounts for. It
emits a "wake KV capacity ledger" INFO line mirroring the boot "KV capacity
ledger" (every term named, residual line) so a future missing term surfaces
here instead of as a mystery wake-time OOM.

Settles the allocator first (gc + synchronize + empty_cache): at the switch
seam the outgoing member's just-dropped allocations can linger in the caching
allocator until GC runs, so a raw ``mem_get_info`` under-reads free VRAM and
the re-provision budget collapses to a sliver of what the card actually has.

Rank-symmetric: the free-VRAM read is collapsed to the all-rank MIN so every
TP rank sizes the IDENTICAL page count — a per-rank budget maps a divergent
page count, leaving another rank's block-table index unbacked and desyncing
the NCCL collectives on the first post-wake forward. Admin switch path only —
never the serving hot path.

Carve-out, mirroring :func:`~arbi_serve.engine.build_memory_sizing.assert_kv_pages_floor`,
:func:`~arbi_serve.engine.build_memory_sizing.assert_context_fits_kv_ceiling`
and :func:`~arbi_serve.engine.build_memory_sizing.assert_swap_kv_capacity`:
an explicit ``--num-pages`` pin is the operator's deliberate pool size, so
the re-provision restores the full pre-park slab (``None``) instead of a
``gpu_memory_utilization``-derived budget. Without it this function is the
FIRST of the four sites a contended card reaches, and a co-tenant holding
more than ``1 - gpu_memory_utilization`` of the device collapses the ledger
to 0 B — refusing the wake below the 2-page minimum before any of the three
gates that honour the pin ever run. A driver refusal on the pinned size
still raises loudly out of ``map_range``.

Coordinated sync resume of one engine's GPU state at its stable VAs.

Inverse of :func:`release_engine_state_sync`; the sync core shared by
:func:`resume_memory_occupation` and the residency controller's member
wake. Remaps in the proven order: cuMem-backed named pools first
(``named_pools.wake_all()`` → ``graph_pool`` physical at the same VAs) →
the B2 growable KV slab → captured-graph exec re-instantiation
(:func:`_reinstantiate_graph_execs`, not a recapture) → registered
per-tensor state (``sleep_pool.resume()``), so the captured graphs are
callable before any replay. ``namespace`` scopes the pool wake to this
engine's raw + namespaced tags so waking one member never remaps another
parked member's slept allocations. Does not take a section / clear
``_memory_released`` — caller owns that. Returns a small info dict.

``cap_growable_kv_to_budget`` marks the growable-KV slab as discard-parked
(no content to restore): the wake then re-provisions the slab's physical
from scratch, capped to the live budget — the free VRAM left once the named
pools are back, minus the serving floor AND minus everything the remaining
steps of this same resume still have to map (state arenas, graph-exec
staging, the per-tensor restore; :func:`_growable_kv_wake_budget` names each
one) — and to the pre-park size. A
residency swap-back can re-activate a member whose slab was grown larger
(built first, more free) than the sibling that just freed its VRAM, and
re-provision-to-fit is what makes that work instead of OOMing
``map_range``. The pool rebuilds its page bookkeeping from the
re-provisioned size, zeroes the pages, and fails loud if the driver
mapping does not match. An offloaded (content-preserving) wake must not
pass it — the pool refuses to re-provision content-bearing chunks
regardless, but the intent is the guard.

Drain in-flight, drop GPU state behind stable VAs, mark released.

Stable-VA (``phase2``) mode only. Captured cudagraphs survive —
every ``data_ptr`` they reference is stable across the
release/resume cycle.

Idempotent. The drain runs through the CriticalSection; the
GPU-state drop runs in a worker thread so the event loop stays
responsive for ``/health`` etc.

``cuMemMap`` fresh physical pages at the same VAs + HtoD restore.

The scheduler / attn_ops / metadata_builders / captured cudagraphs
are untouched because every ``data_ptr`` they reference is stable.

``flat_dump_dir`` is deleted below: the remap restores the same physical
weights at the same VAs, so there is nothing for a flat dump to reload.
Kept on the signature for client/API compatibility only.

Engine wiring for multi-model stable-VA residency.

This is the integration of the GPU-proven mechanism in
``runtime/stable_va_residency.py`` into the live :class:`~arbi_serve.engine.engine.Engine`.
It generalises what ``engine/sleep.py`` already does for one model
(``release_memory_occupation`` / ``resume_memory_occupation`` — unmap physical
behind stable VAs, remap at the same VAs, captured graphs survive) to N resident
records, so a request for a different model performs a recapture-free switch.

Design (``docs/memory-accounting.md``, "Release, wake, swap")
------------------------------------------------------------
The load-bearing invariant is that a captured CUDA graph bakes virtual
addresses, not physical pages. ``engine/sleep.py`` exploits this for a single
model. For N models we need each model's allocations to live in a disjoint VA
range so their graphs never alias.

We achieve disjointness by construction, without a fixed-base sub-allocator,
using a property of the cuMem allocator already present in the codebase: a
park unmaps a model's physical pages but keeps its VA reservations alive. So
while model A is parked, the driver will not hand A's still-reserved VAs to
model B — B's allocations land at fresh, disjoint addresses automatically. The
:meth:`CuMemPoolAllocator.tag_namespace` primitive stamps each model's pool
allocations with a per-model tag prefix so park/wake can target exactly one
model's physical pages.

``ResidentModelRecord``
-----------------------
Per prepared model we retain (in ``record.extra``) the engine's full per-model
attribute set — the same objects ``sleep.py`` and the hot-swap reload path
release/resume, but a snapshot per record:

  * ``model`` (the ``nn.Module``)
  * ``captured_graphs`` / ``layer_captured_graphs`` / ``drafter_graphs``
  * ``named_pools`` (the per-model ``NamedPoolRegistry``)
  * ``pool`` / ``page_table`` / ``scheduler`` / ``sampler``
  * ``attn_ops`` / ``metadata_builders`` / ``active_backends`` / ``active_specs``
  * the runner's graph-dispatch state (``piecewise_buffers`` etc.)

``park(record)`` snapshots+unmaps; ``wake(record)`` repoints+remaps. A switch
runs strict park-then-wake under ``engine.critical_section(ALL_KINDS)`` so the
scheduler is paused for the unmap->map window and no forward runs against
half-mapped physical.

The default single-model path never constructs this controller — it is engaged
only when ``cfg.stable_va_residency`` is set or more than one model is
registered. With one model the controller holds one record, never parks, and
behaves exactly like today.

The park-owned dump directory for ONE member beside ``dense_dir``.

The dense flat cache is keyed on the CHECKPOINT (a content fingerprint of
the model directory), which is right for it: every boot of that checkpoint
wants those bytes. A park-owned dump holds a member's LIVE tensors, which
the checkpoint does not describe — its o_proj may be folded with this
member's calibration, and its head-quant delta is this member's quantiser
output. Two members over one checkpoint therefore have two different
answers, and ``<dense_dir>.live`` had room for one: the second park
overwrote the first's bytes, and the first's wake DMA'd them back into
tensors whose ``_tkv_oproj_folded`` latch still suppressed
``rotate_output``. Nothing about that is loud — the shapes, names and
dtypes all match; only the numbers are another member's.

So the member's identity joins the path. The full key is hashed rather
than only sanitised, because two served names that differ solely in a
character the filesystem cannot take must not sanitise to one directory.
The ``.live`` / ``.headq`` suffix stays last so
``flat_cache._PARK_OWNED_SUFFIXES`` still recognises these as park-owned
(the LRU sweep must not evict them).

Repoint the process-global active-member RuntimeFlags overlay at ``rec``.

The live-config-override seam: a prepared member carries its own flag
overlay in ``rec.extra["flag_overlay"]`` (the variant's runtime-flag
deltas). Installing it on wake makes the fresh-read flags
(recurrent_prefill_chunk / prefix_grouping[_window] /
mtp_spec_disable_batch) member-aware. An empty overlay restores
the env-derived defaults (the boot model / default deployment).
``gdn_decode_num_warps`` also lives in this overlay but is scope="backend"
(capture-affecting) in :mod:`arbi_serve.config_overrides`, not fresh-read
— see ``engine/config_variant.py:_FRESH_READ_FLAGS`` for why. Best-effort:
never blocks a switch.

Drop the prefix cache + recurrent savepoints of a member whose KV content
was discarded (offload=False) on park.

Discarding frees the KV slab's physical without preserving content, so every
cached prefix page + recurrent savepoint now points at stale/garbage (and,
after a budget-capped wake, unmapped) VRAM. Left in place, a later request
whose prompt matches a cached prefix hits it and the forward gathers a freed
page → CUDA segfault. Flushing returns the (refcount-0, scheduler-drained)
pages to the free list and forces the next request to re-prefill cleanly.
Best-effort — a flush failure must not abort the park.

Re-point the process-global FLA / cuBLAS workspace-pool pointers at the
now-active model's named pools.

Two module-globals are repointed here:

  * ``_fla_persistent_cache._GDN_WORKSPACE_POOL`` — the eager GDN
    ``chunk_gated_delta_rule`` chunk-fwd workspace pool. Under ``split_attn``
    (the production default) :func:`route_gdn_workspace_for_serving` clears
    it to the default allocator for the woken member, so the process-global
    pointer never carries a stale reference to a previous member's pool. On
    the monolithic (``split_attn=off``) path the workspace is capture-pinned
    to ``scratch.gdn`` by the wake-time recapture sweep, so the helper leaves
    it alone. (Routing this workspace in-arena via
    ``scratch.forward_arena`` so a park could unmap it is a confirmed
    non-fix: GPU verify showed the
    serving eager GDN op runs while a foreign pool ``begin`` is already
    active, so begin-based routing cannot nest and degrades to
    unrouted under load. The actual park-residual fix migrates
    the AWQ-INT4 head-quant buffers into ``model.weights``, guarded by
    ``ARBI_SWAP_CORESIDENCY_ASSERT``.)

  * ``cublas_warmup._CUBLAS_WORKSPACE_POOL`` — a boot-only concern (capture
    side-stream seeding); a stale value at serving only mis-attributes a
    workspace in metrics. Repointed for consistency.

  * ``exl3.kernel_scratch._KERNEL_SCRATCH_POOL`` — the woken member's
    ``capture.io_buffers``, which an EXL3 descriptor's bsz-1 input row is
    allocated in. Only a lazy descriptor REBUILD allocates through it
    (``_build_inner`` after a release hook), so unlike the sampler's it is
    not on a per-token path — but a rebuild against the parked member's
    pool writes unmapped VA, and the repoint costs one global write.

Best-effort — never blocks a switch (the globals are missing only on
CPU/test paths).

Re-arm the sampler's process-global scratch-pool hooks at the woken
member's pools.

``sampler/topk_topp_triton`` and ``sampler/penalty_accumulator`` each hold
a module-level ``_POOL_ALLOCATE`` closure that a member's build installs
over THAT member's ``scratch.sampler_triton`` / ``scratch.penalty_accum``
pool (the penalty one over a carve of its persistent slab). The globals are
process-wide and nothing but a member BUILD ever re-installed them, so
after a park/wake switch they still reach into the member the swap parked —
whose pools the park unmapped (both carry the default ``offload`` sleep
strategy). Every sampled token allocates or re-uses scratch through them,
so the woken member's first request writes unmapped VA:
``cudaErrorIllegalAddress``, from a pool that is at full size and a page
count that is exactly right.

Same seam and same contract as :func:`_reset_cublas_workspaces_on_wake` and
:func:`_reset_head_quant_marlin_workspaces_on_wake` — re-establish on the
woken member what the wake's remap alone does not. The hook is carried on
the record's own ``_sampler_scratch_hook`` / ``_penalty_accum_hook``
attributes (``ENGINE_MODEL_ATTRS``), already repointed by ``state.restore``
when this runs.

Runs BEFORE :func:`~arbi_serve.engine.member_scratch_retire.
restore_device_memos`, because installing a hook drops the scratch memos it
feeds — the restore then puts the woken member's own tensors back at the
addresses its build allocated them at. A member with no hook recorded (CPU
bookkeeping, a boot that never installed one) UNINSTALLS instead of leaving
the previous member's armed: disengaged falls back to the default
allocator, which is the cold-boot behaviour, whereas armed-at-a-parked-pool
is the fault. Best-effort — a switch must not die on the sampler's
accounting.

Logs its OUTCOME unconditionally, at INFO, naming both arms and whether
each came back armed or disengaged. A wake seam that leaves no trace cannot
be told apart from one that did not run, and that difference is the whole
diagnosis when a switch still faults.

``corpses`` less the pools this member's own registry retired.

``NamedPoolRegistry.release_empty_pool`` is a boot-time reclaim of an EMPTY
pool: it hands the physical back and re-registers a fresh pool under the
same name, so the allocation record it leaves behind is a released one that
no tensor of this member's ever reads again. It is indistinguishable, in
the allocator table, from a record another member's build released — which
is why the owner has to be asked instead of the table.

Report any allocation under the woken member's tags that a wake CANNOT map.

:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.wake` maps a
sleeper (``asleep=True``) and skips a corpse (``asleep=False``) — the
distinction :meth:`release_tags` draws so a wake does not hand fresh
physical to a destroyed pool's dead VA. A corpse under a tag this member
still owns is neither: the member's tensors still hold those addresses and
the resume just declined to back them, so the first kernel that touches one
reports an illegal access from inside whatever op happened to run, with
nothing in the allocator raising.

Raw pool names are the way this happens. ``_release_pool_physical`` targets
``{name, "<ns>/<name>"}``, and a boot member whose build ran without a tag
namespace (``cfg.stable_va_residency`` unset — residency engaged later, on
the first live config override) carries RAW tags for everything. A pool
reclaim inside a VARIANT's build then matches the boot member's parked
allocations by raw name and buries them.

A pool the member destroyed ITSELF is not that, and is excluded. The boot
fold retires ``scratch.rope`` once its tables move to the persistent slab
(``engine.persistent_fold``), which leaves a permanent released record
under the member's own tag with nothing alive pointing at it. Reporting it
named a release that never happened, on the one path an operator reads to
find the release that did — so the member's own registry is asked which
names it retired (``NamedPoolRegistry.released_pool_ids``) and those are
dropped before the verdict.

Reports rather than raises: the member is already repointed and refusing
here would leave nothing serving, which is worse than serving and faulting
with the cause named. Cheap — an allocator-table scan, no device work.

Drop torch's cached cuBLAS/cuBLASLt workspaces so the woken member's
first eager GEMM re-establishes a fresh workspace against its own mapped
memory.

Torch caches one cuBLAS/cuBLASLt workspace per ``(device, stream, handle)``
in a process-global map (``torch._C._cuda_clearCublasWorkspaces``), lazily
allocated on the first GEMM each stream/handle runs — for the serving-stream
workspace, that first GEMM is a member's own forward, so the workspace lands
in that member's per-model activation/scratch pool. The cache is not keyed by
stable-VA member, so after a residency switch it still points at the outgoing
member's workspace — whose pool the park unmapped. The base decode stack
replays a captured graph (its workspace VA was baked at capture into the woken
member's own pool and remapped by the scoped wake), so it survives — but the
MTP / DFlash seed + verify forwards run eager, and their eager
``cublasGemmEx`` reads the stale cached workspace pointing at the unmapped
pool → ``CUBLAS_STATUS_INTERNAL_ERROR`` / ``cudaErrorIllegalAddress`` on the
woken member's first forward.

Clearing the cache here (after the pools are remapped) is only half the
seam: torch would otherwise re-allocate the workspace lazily, on whichever
eager GEMM runs first and therefore in whatever allocator context that GEMM
happens to carry. After the Phase-2 freeze that re-allocation has no home —
it is the ``residency.wake_reallocation`` class
(:mod:`arbi_serve.engine.post_freeze_budget`). So the drop is paired here
with an immediate re-seed into the woken member's own
``scratch.cublas_workspace`` pool
(:func:`~arbi_serve.engine.cublas_warmup.reestablish_cublas_workspace`),
which the member's build already seeded with a workspace-sized block. The
workspace is handle-fixed, a cuMem pool keeps its physical mapped across a
free, and the re-seed's operands stay outside the pool — so the pair drops
and re-takes the same block and the wake maps no new physical. The
tag's mapped bytes are read either side of the pair to check exactly that.

Safe for captured graphs: replays carry their own baked workspace VA (in the
woken member's remapped pool) and never consult this runtime cache; the
cleared entries are the eager-path workspaces. Best-effort — a torch without
the private API must not break a switch, so the miss is logged, not silent.

Physical bytes mapped under ``pool``'s cuMem tag, or 0 when unreadable.

Reads the allocator rather than the caching allocator: a free returns a
block to the pool's free list without unmapping it, so mapped bytes are the
quantity a re-seed must leave unchanged.

Drop the TP/EP coordinators' boot-stable collective staging buffers
while the outgoing member is still mapped.

``GroupCoordinator._staging`` holds one buffer per
``(dtype, device, numel)``, allocated on first use into the attached
``comms.nccl`` :class:`NamedMemPool` and then held for the process life.
The cache is a process singleton — not keyed by stable-VA member — but the
pool it allocates from is strictly per-model, so the buffers belong to
whichever member ran the first live-eager collective. After a switch the
cache still points into the outgoing member's pool, which the park
unmapped: the woken member's first live-eager ``all_gather`` stages through
it and segfaults in ``cuMemcpyDtoDAsync``.

Same class as the cuBLAS and INT4-Marlin wake resets, with one ordering
difference: those re-establish after the wake remaps, whereas these
buffers must be released before the park unmaps, so the frees land on
live pages rather than on a pool whose physical is already gone.

Safe for captured graphs: an address-pinned collective (compiled or raw
capture) never touches these buffers — it gathers straight from the
activation whose address the graph pins. Only live eager collectives stage,
and they re-allocate on demand into whatever pool is attached at the time.

Point the TP/EP coordinators' staging allocations at the woken member's
``comms.nccl`` pool.

The counterpart to :func:`_release_collective_staging_before_park`: the
wake has just repointed ``named_pools`` at this member, so re-attach before
the first live-eager collective allocates — otherwise it lands in another
member's pool and the next park unmaps it out from under this one.

Reset then attach, in that order. ``attach_staging_pool`` refuses to
re-point while buffers exist (it would orphan addresses), and buffers do
exist on the build-required variant route: the member built beside this one
allocated its own during capture, and this wake is what un-binds it. Reset
first — the same order the hot-swap reload path uses, and safe for the same
reason: only live-eager collectives touch these buffers, so dropping them
orphans no captured graph, and the drop is allocator bookkeeping that does
not read the pages (which is why the reload path can do it after a
teardown has already freed them).

Re-establish the drafter/head-quant Marlin cross-CTA lock workspace on
the woken member (drafter-int4 / full-int4 wake decode hang).

The INT4 head-quant linears (drafter ``lm_head`` override, ``mtp.fc``,
full-int4 ``lm_head``) carry a non-persistent ``_marlin_workspace`` — an
SM-count int32 array of Marlin's cross-CTA reduce lock slots, which the
kernel requires zeroed on entry (it resets them to zero at kernel end for
reuse). At boot these buffers are migrated into the member's
``model.weights`` cuMem arena (the single-VRAM-resident co-residency
invariant). A residency park releases that arena's physical; the scoped wake
remaps it at the same VA but restores only the persistent per-tensor sleep
entries — the non-persistent workspace bytes come back stale. The base
decode stack is unaffected (it replays a captured graph whose workspace VA +
lock lifecycle are baked at capture), but the eager MTP drafter-seed forward
reads ``self._marlin_workspace`` live: its vocab-projection GEMM (N = padded
vocab ≫ 2048 ⇒ ``use_atomic_add`` forced False ⇒ the lock-based fp32
cross-CTA reduce) reads the stale lock slots and spins forever — a wedged
(no CUDA error / IMA) ``marlin::Marlin<…>`` kernel that never drains,
blocking the seed forward's D2H (``drafts[:req_k].cpu()``).

Zero the lock slots in place here (after the pools are remapped) so the
woken member's first eager drafter GEMM re-establishes a clean reduce. This
is the "re-establish what the wake failed to restore" seam — the same
contract as :func:`_reset_cublas_workspaces_on_wake` and the sampler-cache
clear, not a serialize / retry / skip-drafter-on-wake. In-place zero keeps
each buffer in its arena (co-residency invariant intact) at its
already-valid remapped VA. Best-effort — head-quant off yields no such
linears, so it is a no-op there.

Record ``pages`` as this member's KV high-water; return the new value.

Monotone by construction, and that is the whole point. A park also records
``kv_pages_at_park`` — what the member happened to be serving at THIS park —
and a wake that came back short lowers it. Capping the next restoration at
that number turns a transient shortage (a sibling still resident, physical
the outgoing member had not returned) into the member's permanent size, and
every further swap can only lower it again. The high-water never falls, so
the pages come back when the card does.

The most pages a wake may restore this member to.

A CAP, never a target: :func:`~arbi_serve.engine.config_variant._regrow_active_kv`
is free-bounded by ``serving_floor_for_grow`` and reaches this only when the
VRAM is genuinely there. It exists because "as much as fits" is a different
and unsafe behaviour at a drained wake, where the free VRAM IS the next
serving step's working set (``tests/test_kv_regrow_respects_the_serving_step.py``
carries the GPU-measured incident). ``0`` means no cap is known — the caller
passes ``None`` and the free bound governs alone.

Engine-side controller for recapture-free multi-model switching.

Wraps :class:`StableVaResidencyRegistry` (the pure bookkeeping +
park/wake state machine) and supplies the GPU mechanics (cuMem
namespace sleep/wake + engine attribute repointing) as the injected
``park_fn`` / ``wake_fn``.

Opt-in. Construct with the booted engine + the boot model's key; the
boot model is registered as the active record (already mapped). Prepare
additional models with :meth:`register_prepared` (each must have been
built under its own tag namespace). Drive switches with
:meth:`switch_to` (sync, called from inside the engine critical section
by :meth:`aswitch_to`).

Scope in which a park discards the member's PER-SEQUENCE state.

A residency swap parks the active member only to free its VRAM for a
sibling member. The scheduler is drained (no committed sequences), the
park flushes the prefix cache and the recurrent savepoints, and the
woken member re-prefills from empty — so nothing the member's sequence
state holds is ever read again. Copying it to host is pure waste, and
not merely wasteful: the destination is UNSWAPPABLE pinned host RAM,
and at ``gpu_memory_utilization`` 0.99 that state is most of the card.
An admin config-override on a large model could therefore demand tens
of GiB of host RAM to park a member whose state it was about to throw
away, and reach the host OOM killer.

The scope covers the WHOLE park sequence, not just the park call: the
growable-KV slab, the sentinel-alias recurrent arenas, and the
``state.*`` named pools that :meth:`evict_parked_to_host` reclaims
afterwards. Those three hold one thing between them — the member's
sequence state — so one contract governs all three; a caller that
opened this scope for the park and closed it before the evict would
discard the slab and then pin the rest, which is the bug this scope
exists to make unrepresentable. The flag auto-restores.

Everything OUTSIDE the sequence state (weights, the captured-graph
pool and its I/O buffers, the scratch arenas, the drafter) is still
offloaded content-preserving: those hold bytes a wake cannot
regenerate, or that captured graphs read at addresses this scope must
not disturb.

Scope in which a park DROPS the member's weights from

A CROSS-MODEL swap parks the outgoing member to free its VRAM for a
DIFFERENT model whose weights cannot co-reside (the stable-VA
donor-share invariant keeps every member's weights mapped, so two big
distinct models' weights sum past VRAM). Within this context the park,
after the light release, discards the member's weights tag (physical
freed, VA kept) — the on-disk flat dump backs them and the next wake
reloads them. Only ever entered for a cross-model park; same-model
config variants never enter it (their weights are donor-shared and must
stay mapped). The flag auto-restores.

The underlying ``model.path`` a record was built from, or None.

Config variants (``cfg-<sig>``) and the boot member of the same model
carry the SAME ``model.path``; only distinct models differ. Delegates to
:func:`arbi_serve.engine.residency_view.resident_model_path`, the one
reader of that snapshot field — the request router resolves a
checkpoint-path request through the same function, so the weight-drop
decision here and the routing decision there can never disagree about
which records share a model. Returns None on any missing snapshot so a
caller falls back to the conservative branch.

True iff both records were built from the SAME ``model.path``.

Same path ⇒ same-model config variants that donor-share one weight set
(never drop). Different path (or either unknown) ⇒ CROSS-MODEL, so a
swap between them must drop the outgoing member's weights to fit the
incoming model's.

True when ``rec`` owns its weights and no other member aliases them.

The authoritative donor-share guard: a member that donor-bound another
record's weights (``weights_donor``) or whose weights a live sharer
aliases (``weight_sharers``) must NEVER have them dropped — the storage
is shared. Only a sole owner's weights are safe to discard.

The cuMem tags a member's weights pool may carry (raw + namespaced).

A member built under ``tag_namespace(ns)`` stamps its weights allocs
``"<ns>/model.weights"``; any runtime growth outside a namespace is raw
``"model.weights"``. Target BOTH forms so a discard/wake covers the pool
regardless of when each alloc was made (mirrors
``named_pool_registry._effective_tags``).

Return ``rec``'s bulk flat-dump dir, publishing whatever a wake needs.

Derives the dense dir with the SAME ``default_flat_cache_dir`` the
cold-boot loader uses (per-model content fingerprint + per-rank suffix).
The cross-model wake reloads this dump IN PLACE into the already-built,
live model, so the fill must cover the live tensor set EXACTLY —
including any in-memory-quantized drafter buffers (``mtp_head.fc`` marlin
qweight / scales / format marker) the head-quant transform produced AFTER
the head-quant-INDEPENDENT warm-boot dense cache was published.

The bulk dump is NEVER re-serialized: the warm dense cache holds it.
When that cache cannot fill the live model only because it LACKS a few
live tensors (the head-quant delta), write just those tensors to a tiny
park-owned SUPPLEMENT (a few MB) recorded in
``rec.extra["flat_supplement_dir"]``; the wake loads bulk-from-dense +
delta-from-supplement. Only a deeper divergence (the dense cache is
absent, or its shared tensors mismatch) forces a full live dump. A
member is live on the engine when parked, so its weights are still
mapped for the D2H. Caches the bulk dir in ``rec.extra["flat_dir"]``.

A member that folded its o_proj weights in place
(:func:`~arbi_serve.engine.stable_va_model_state.
oproj_fold_generation` > 0) never reuses the dense cache: that blob
holds the checkpoint's UNFOLDED bytes, and the header-only compatibility
check cannot see the difference. Its own dump goes to a PER-MEMBER path
(:func:`park_owned_dir`), because ``<dense_dir>.live`` is one directory
and two members over one checkpoint have two different weight sets.

The cached answer is reused only while the member's o_proj fold count is
the one the dump was written under. That count rises on the boot fold and
again on every calibration-reload refold, each of which rewrites
``o_proj.weight`` in place — leaving the resolved directory holding the
PREVIOUS calibration's bytes, which the wake would DMA back with the
``rotate_output`` skip still latched.

Discard ``rec``'s weights from VRAM (VA kept), backed by its flat dump.

The cross-model park mechanism: ensure the on-disk flat dump exists
(while the weights are still mapped), then cuMem-discard the weights tag
stay put (cudagraphs' baked addresses survive) and the next wake remaps
+ reloads from the dump. Refuses if the weights are donor-shared. No-op
on CPU / no-driver bookkeeping.

Remap + refill

The inverse of :meth:`_drop_member_weights_to_flat`, run on wake BEFORE
:func:`resume_engine_state_sync` (whose ``wake_all`` would otherwise
remap the weights ZEROED — discarded pages carry no host backup): remap
the weights VAs (zeroed pages come back at the same addresses), DMA the
flat dump into the live tensors IN PLACE (state_dict refs — cudagraph-
pinned addresses preserved), then fire each quant module's
``rebind_after_compaction`` so its kernel descriptor re-derives over the
filled buffers. Clears ``weights_on_flat``. No-op unless the flag is set.

Fan out a resident-set / active-key mutation to the engine's
residency-changed hooks (process mode: the bindings emit the
``residency_changed`` control event so the API child's cached
routing view follows). Best-effort by contract — the hooks own
derived state only; a test double without the method is fine.

Register the already-booted, GPU-resident model as the active record.

Snapshots the engine's current per-model state into the record and
marks it mapped (no park/wake). The boot model was built by the normal
path; we wrap it so the first real :meth:`switch_to` parks it
correctly.

Register a model that was prepared (built + captured) under its own
tag namespace but is not currently the engine's active model.

The caller builds the model's per-model state into the primary engine
under ``with alloc.tag_namespace(namespace):`` (so its pools' physical is
tagged for per-model park/wake), snapshots that state with
:meth:`_EngineModelState.capture`, then hands the snapshot here.

``mapped`` records whether the prepared model's physical is still mapped
at registration. The single-resident prepare orchestrator
(:func:`prepare_model_pool`) builds each member as the sole VRAM
resident, snapshots it, registers with ``mapped=True``, then immediately
:meth:`registry.park`s it (offload to host) before building the next —
so only one member's physical is ever mapped, and a member's capture
never runs while another member is parked-in-VRAM. The default
``mapped=False`` is for a caller that already offloaded the physical
itself (the record is parked from the start).

Refuse a new record if the resident set is over its declared cap, or
if the graph memory this registration still owes cannot be paid.

``max_records`` is the operator's declared ceiling on the resident set
and is enforced unconditionally.

The free-VRAM half asks
:func:`~arbi_serve.engine.build_graph_pool.unpaid_graph_exec_bytes`:
driver-side ``cudaGraphExec`` memory lives outside every cuMem arena, so
free VRAM is the only thing that can pay for it — but only for the
execs that have NOT been instantiated yet. By the time a prepared member
registers, its capture sweep has run and the driver has already taken
those bytes (they are the ``driver.cudagraph_exec`` ledger row), so the
term is paid and this gate demands nothing.

It does not become owed again while the record is parked: a park
DESTROYS every owned exec (:func:`~arbi_serve.engine.sleep.
_deinstantiate_graph_execs`) and the wake re-instantiates them
(:func:`~arbi_serve.engine.sleep._reinstantiate_graph_execs`) out of a
budget that holds the instantiate reserve explicitly
(``_growable_kv_wake_budget``). So a parked record carries no graph
memory for this gate to reserve, and the moment it re-acquires some is
already sized against the free VRAM live at that moment. Charging every
resident record for bytes only the ACTIVE one holds — and holding them
free on top of the card they are already on — is what used to make this
gate refuse a build it had itself sized for.

Restoration cap for the member currently resident, or ``0``.

The drop path grows the ACTIVE member's slab into the VRAM a teardown
just freed, on the same drained engine and against the same free
reading a wake uses — so it needs the same cap, and it is the same
member's history that supplies it (:func:`kv_restoration_cap`). ``0``
for a member that has never parked: nothing is known to restore TO, and
the free bound governs alone, which is what that path always did.

Offload ``rec``'s physical VRAM to host; keep its VAs + graphs.

Runs the proven coordinated single-engine release
(:func:`arbi_serve.engine.sleep.release_engine_state_sync`) against the
one primary engine — releasing ``rec``'s per-tensor sleep entries
(``sleep_pool``, which pin the captured-graph input buffers), its
cuMem-backed named pools (``graph_pool`` …) and its growable-KV slab to
pinned host RAM. The captured graphs survive (their data_ptrs are stable
VAs); only the physical pages move to host.

Precondition: ``rec`` must be the model currently installed on the
engine (the registry's :meth:`switch_to` parks the active record while
its state is still live; :func:`prepare_model_pool` parks the member it
just built into the engine). We assert that so a misordered park can
never release the wrong model's pools.

Repoint the engine at ``rec`` + remap its physical from host.

Restores ``rec``'s snapshot onto the one primary engine (so
``eng.sleep_pool`` / ``named_pools`` / ``pool`` point at this model's
released state), then runs the proven coordinated resume
(:func:`arbi_serve.engine.sleep.resume_engine_state_sync`) — remap the
named pools' physical at their same VAs, the growable-KV slab, and
restore the per-tensor sleep entries from host. After resume the
captured graphs are immediately callable (their baked VAs never moved):
no recapture. The engine now serves ``rec`` through the shared run loop.

Release the just-parked member ``key`` entirely. Its key, or ``None``.

``None`` means the member must be kept, and the common reason is the
right one: the boot model and every declared pool member are refused,
which is exactly the record a switch BACK to the baseline is heading
for. Only a variant built through ``config_override`` is ephemeral
enough to destroy, and only when the caller asked
(``drop_previous``) — a sweep alternating between two configurations
wants the parked member kept so re-selecting it stays instant.

THE EXPOSURE, stated because it is the caller's to accept: once this
returns a key, that member is gone and the wake that follows has
nothing to fall back to. It is not a regression on the alternative —
the wake this frees memory for is the one that fails without it — and
the park has already discarded the member's sequence state, so nothing
was recoverable but the build. Re-issuing the override rebuilds it.

Never raises: a reclaim that cannot run leaves the member parked, which
is the pre-existing behaviour, and the swap continues.

Reclaim a parked record's residual cuMem pool physical.

A park keeps the member's namespaced pools VRAM-mapped (the ~10 ms
switch); a member build beside it needs that VRAM — the swap to a
bigger-KV backend OOM'd its prepared-member build against the parked
member's still-mapped graph pool / arena / KV. Reclaims every
sleepable pool except ``weights_pool`` (donor-share invariant: the
new member may donor-bind the parked weights). The record's tag
namespace targets the member's ``"<key>/<pool>"``-tagged allocations;
the next wake's scoped resume remaps them at the same VAs (no
recapture). Returns bytes reclaimed; 0 on CPU / bookkeeping states.

HOW each pool is reclaimed is the caller's contract, not this method's
choice. Under :meth:`discard_sequence_state_on_park` the ``state.*``
pools — the member's per-sequence state, which the park has just
flushed the bookkeeping for and which the woken member re-prefills
from empty — are DISCARDED: no D->host copy, pages come back zeroed.
Every other pool is offloaded content-preserving as before.

That distinction is the whole host-RAM budget of a swap. At
``gpu_memory_utilization`` 0.99 the ``state.*`` pools are most of the
card, the offload destination is UNSWAPPABLE pinned host RAM, and a
member is parked for as long as it stays parked — so offloading state
the swap has already declared dead is how an admin config-override on
a large model reaches the host OOM killer. Called OUTSIDE the contract
it reverts to offloading everything, which is why the orchestrators
keep the scope open across both calls.

Drop a just-parked DISTINCT-model member's weights to its flat dump.

The multi-DISTINCT-model boot path (:func:`prepare_model_pool`) calls
this after it parks + evicts each member so the NEXT member builds as a
true sole VRAM resident — a pool of two big distinct models cannot keep
both weight sets mapped (their sum overruns VRAM, the boot The final :meth:`registry.wake` of the boot member then reloads
its weights from its dump (via :meth:`_wake_mechanics`). No-op when the
weights are donor-shared (a same-model config-variant pool never reaches
this path) or on CPU / no-driver bookkeeping.

Park the active record, wake ``key``, repoint the engine. Sync.

Must run inside :meth:`Engine.critical_section` (scheduler paused) —
use :meth:`aswitch_to` from async call sites, which holds the section.

After a completed switch, runs the engine's registered model-switch
hooks (``Engine.register_model_switch_hook``) — the server layer's
refresh of model-derived state that lives outside the repointed
per-model engine attributes (e.g. the app-state ``TokenizerPool``:
a cross-architecture switch replaces ``eng.tokenizer`` with a
different-vocabulary tokenizer, and a boot-frozen pool would encode
the new member's prompts with the wrong vocabulary). Hook failures
are logged, never raised — the engine state is already consistent.

Memory-safe park. Rather than the registry's light park — which frees
the outgoing member's growable-KV slab but keeps its namespaced pools
(graph buffers / arena) VRAM-mapped for a ~10 ms wake-back — we fully
evict the outgoing member's residual pool physical to host before
waking the target (:meth:`evict_parked_to_host`, the same park+evict
:func:`prepare_model_pool` pairs). At a tight gpu-memory-utilization
two same-model members (a boot member + a config variant) do not
co-reside: the target's KV/graph pools must remap into the space the
outgoing member's leftover pools would otherwise squat, so the
light-park wake OOMs its KV ``map_range``. The switch is H2D-bound
regardless (parked members are host-offloaded, so the wake already
restores from host), so the extra evict copy is consistent with the
residency model, not a new cost class.

Park+evict the active record, then wake ``key`` (memory-safe switch).

The registry's own :meth:`StableVaResidencyRegistry.switch_to` does a
light park (pools stay mapped) then wake — fine when both members fit
co-resident, but not at a tight gmu where two same-model members can't.
This mirrors that state machine (noop / park-then-wake, same result
dict) but inserts an :meth:`evict_parked_to_host` between the park and
the wake so the target has the VRAM. Evict is a no-op on the CPU
bookkeeping path (returns 0), so this is behaviour-identical there.

Async switch: hold the engine critical section across park->wake.

No forward runs against half-mapped physical because the scheduler is
drained + paused for the whole unmap->map->repoint window (the same
guard ``release/resume_memory_occupation`` use).

Shared per-model state definitions for multi-model stable-VA residency.

The dependency-light core the controller
(:mod:`arbi_serve.engine.stable_va_controller`) and the single-resident
prepare/build orchestrator (:mod:`arbi_serve.engine.stable_va_pool_builder`)
both build on: the ``PoolMember`` type, the ``ENGINE_MODEL_ATTRS`` roster of
per-model engine attributes, and the ``_EngineModelState`` snapshot dataclass. Kept in one place
so both siblings import it without a cycle. Every name here is re-exported from
``stable_va_controller`` for backward-compatible imports.

How many times ``state``'s member has rewritten its ``o_proj`` weights.

The engine's ``_oproj_folds_applied`` counter, read through one expression
so no caller reaches for the attribute name itself. It rises on the boot
fold and again on each calibration-reload refold — every event that leaves
the member's live weights different from the bytes any earlier artifact
recorded — which is what makes it usable as the generation a park-owned
flat dump is keyed on (:meth:`~arbi_serve.engine.stable_va_controller.
StableVaResidencyController._ensure_member_flat_dump`), not just as a
yes/no.

Accepts an :class:`_EngineModelState`, its raw ``attrs`` mapping, or the
``None``/sentinel states CPU bookkeeping records carry (→ ``0``).

True when ``state``'s member rewrote its ``o_proj`` weights in place.

The TKV o_proj fold (:func:`arbi_serve.engine.tkv_oproj_fold.
fold_o_proj_for_engine`) multiplies ``o_proj.weight`` by the member's OWN
calibration-derived Every artifact keyed on the CHECKPOINT rather than on the live
tensors therefore describes a different weight set than a folded member
holds:

  * the same-model donor weight share, keyed on ``(model.path, dtype)``
    (:func:`arbi_serve.engine.build_helpers.find_weight_share_donor`) —
    a folded donor would impose its calibration on every sharer;
  * the content-fingerprinted dense flat cache a cross-model park reuses to
    back its weight drop
    (:meth:`~arbi_serve.engine.stable_va_controller.
    StableVaResidencyController._ensure_member_flat_dump`) — that blob is
    published at the pre-fold dense seam, and ``flat_dump_compatible`` is
    header-only (name/dtype/shape), so a folded member passes it and the
    wake would refill UNFOLDED weights while the module-level
    ``_tkv_oproj_folded`` latch still suppresses ``rotate_output``. That
    caller reads :func:`oproj_fold_generation` rather than this boolean,
    because it must also notice a SECOND fold (a calibration reload's
    refold) invalidating a dump it already published.

Both consumers reach the ``_oproj_folds_applied`` count through this one
pair of functions, so the fold cannot be visible to one and invisible to
the other.

Accepts an :class:`_EngineModelState`, its raw ``attrs`` mapping, or the
``None``/sentinel states CPU bookkeeping records carry (→ ``False``).

Single-VRAM-resident prepare/build orchestration for stable-VA residency.

The startup + config-variant build side of multi-model stable-VA residency:
:func:`build_member_into_engine` constructs one member's per-model state into
the shared primary engine under its cuMem tag namespace, and
:func:`prepare_model_pool` sequences that across a member list (build each as
the sole VRAM resident, offload to host, proceed) and engages the
:class:`~arbi_serve.engine.stable_va_controller.StableVaResidencyController`.
Both names stay importable from ``stable_va_controller`` (re-exported).

Build one member's per-model state into the primary engine, in place.

The load-bearing piece of the single-VRAM-resident design. Instead of
booting the member as a separate :class:`Engine` (whose run-loop seams —
``requests`` / ``_wakeup`` / ``_forward_executor`` / ``model_runner`` —
would not transfer on a cross-engine repoint and would stall a woken
member's live serving), we:

  1. Reset the primary's per-model containers (fresh named pools, cudagraph
     pools, sleep registry, backend/metadata registries; cleared handles)
     via :meth:`Engine.reset_model_state_for_build`, so the build populates
     the member's state cleanly without mutating the objects the
     already-built members' snapshots still reference.
  2. Run the member ``builder`` (which sets ``primary.cfg`` to the member's
     config and calls ``primary.build()``) under the member's
     ``tag_namespace`` so every cuMem allocation this member MAKES carries
     the member's tag (per-model park/wake). The other members' VA
     reservations stay alive while this one allocates, so the driver hands
     it disjoint addresses. Storage this member does NOT allocate but BINDS
     — a same-``model.path`` sibling's weights, via
     ``find_weight_share_donor`` — is deliberately the donor's VAs under the
     DONOR's tag, and stays mapped for as long as a sharer is live.
  3. Snapshot the resulting per-model state and return it.

VRAM
resident during its capture. That is what avoids a capture-while-parked
IMA: no second model's pools are mapped while this member's cudagraphs
are captured.

Build the ``[(served_name, builder), …]`` list for a stable-VA pool.

The FIRST entry is the boot model (already built by ``abuild`` / the
driver's ``engine.build`` — its builder is a no-op the orchestrator skips
with ``build_primary=False``). Each EXTRA member gets a builder that derives
a per-member :class:`~arbi_serve.config.ServerConfig` from the boot config
(inheriting everything except the member's own ``served_name`` / ``path`` /
``max_context`` / ``calibration_path``), installs it on the primary engine,
and calls ``eng.build()`` — so ``prepare_model_pool`` builds + captures the
member as the sole VRAM resident, then offloads it to host.

Shared by BOTH boot call sites: the single-process lifespan
(``server.engine_boot._engage_stable_va_residency``) and the per-rank
distributed driver (``DistributedEngineDriver._engage_model_pool``). It
derives purely from ``cfg`` (which every rank holds identically), so every
rank builds the same member list + overlays and stays in lockstep.

Returns ``(members, overlays)``; raises ``ValueError`` on a duplicate
served name (the boot name plus every member name must be unique — they are
the routing keys).

Discard a failed/aborted boot-pool member so the abort is clean on every
rank.

Runs on a NEGATIVE cross-rank barrier verdict (this rank's build failed, or
a peer's did). The boot orchestration aborts the whole boot afterwards
(the exception propagates out of ``build`` and every rank dies together), so
there is no ``prev_active`` to re-wake — this only drops whatever the
orphaned member mapped/registered so the abort does not strand a
half-built record. Mirrors :func:`config_variant._rollback_member_build`'s
orphan-teardown half, minus the prev re-wake.

Build ONE non-boot member into the primary + offload it to host.

``confirm_build is None`` (single-process / TP=1): build, register, make
active, park, evict — the unchanged single-resident step.

``confirm_build`` provided (TP>1): build this rank's shard, then gate the
commit on the cross-rank barrier so an asymmetric per-rank build failure
aborts the member on EVERY rank together (no NCCL-watchdog desync). Both
the success and failure paths reach ``confirm_build`` exactly once so the
barrier collective stays in lockstep; a negative verdict rolls the member
back symmetrically and re-raises.

Single-VRAM-resident prepare: build each member as the sole resident,
offload it to host, proceed to the next; leave the first member active.

This is the startup entry point for a multi-model stable-VA residency set in
one primary engine, under the deliberate relaxation that only one model
is VRAM-resident at a time (the rest parked in host RAM). The supported
phasing: capture each model's cudagraphs while it is the sole VRAM resident
— never while another model's physical is mapped — which avoids a
capture-while-parked IMA.

``members`` is ``[(key, builder), ...]`` in priority order. Each ``builder``
takes the primary engine, sets ``eng.cfg`` to its model's config, and calls
``eng.build()``; :func:`build_member_into_engine` owns the per-model reset +
namespace tag + snapshot, and this orchestrator owns the offload-to-host.

Flow (single-resident):

  1. Build the first member (the boot model) into the primary if
     ``build_primary`` — or take it as already built. Engage the controller
     wrapping it as the active record, then offload it to host.
  2. For each remaining member, while all prior members are offloaded:
     build it into the primary as the sole VRAM resident (recursive reset →
     build under its namespace → snapshot), register it, then offload it to
     host. At no point are two members' physical mapped together.
  3. Wake the first member from host so it is the active resident again.

A later request for a parked member auto-switches recapture-free (park the
active → wake the target, both via the proven coordinated single-engine
sleep). Single-member pool ⇒ the default single-model path: one active
record, nothing parked. Returns the engaged controller.

``confirm_build`` is the TP>1 hook: a callable ``(local_ok) -> cluster_ok``
(the driver passes ``confirm_build_across_ranks`` bound to its engine). When
provided, EACH member's build commit is gated on it so an asymmetric
per-rank build failure aborts that member on every rank together (see
:func:`_prepare_one_member`). Default ``None`` ⇒ the single-process path
commits every build unconditionally (byte-identical to before).

Find live device tensors that point into UNMAPPED cuMem memory.

A residency park unmaps the parked member's pool physical behind stable VAs.
Anything that still holds one of those tensors and touches it after the switch
reads unbacked memory: ``cudaErrorIllegalAddress``, reported asynchronously at
whatever synchronising call comes next, which is almost never the code that
caused it.

The seams that re-establish per-member state at a wake are enumerated one by
one (:mod:`arbi_serve.engine.member_scratch_retire`,
``stable_va_controller._reset_cublas_workspaces_on_wake`` and its siblings), so
each is only as complete as the roster behind it. This module is the general
check that does not depend on that roster being right: after a wake has
finished, walk every live CUDA tensor in the process and report any whose
storage lies inside an allocation the cuMem allocator currently has unmapped.
A finding is a bug by construction — nothing reachable may point at unbacked VA
once a wake has returned.

The walk is a full ``gc.get_objects()`` sweep, so it is opt-in
(``ARBI_STALE_PTR_AUDIT``) and belongs on a swap-debug run, not in the serving
path.

``(start, end, tag)`` for every allocation the allocator has unmapped.

Read straight off :class:`~arbi_serve.runtime.cumem_allocator.
CuMemPoolAllocator`'s live allocation table, so it covers every named pool
of every member — the parked member's included — without this module
knowing which pools exist. Empty when the driver is absent.

A short, bounded description of what still references ``obj``.

Names the container type and — for the dict-of-tensors shape that every
one of these memos takes — the key it is filed under, which is what turns
a finding into a file to open. Never raises and never walks deep: a
referrer graph on a live engine is unbounded.

Report every live CUDA tensor whose storage is currently unmapped.

Returns the number of findings (0 = clean). ``phase`` names the seam the
audit ran at, so a log reader can tell a pre-wake sample from a post-wake
one. No-op returning 0 when ``ARBI_STALE_PTR_AUDIT`` is off, when the cuMem
driver is absent, or when nothing is unmapped.

Enable it on a swap-debug run: a park/wake that still faults has a holder
this process can name, and naming it here is cheaper than inferring it from
an asynchronous illegal address several steps later.

ONE per-step outcome seam, for every run loop this repo has.

There are three engine run loops — the single-process
:func:`~arbi_serve.engine.run_step.loop._run_forever_inner`, the rank-symmetric
:meth:`~arbi_serve.engine.distributed_driver_spmd._SpmdLoopMixin._run_loop_spmd`,
and the legacy rank-0 driver loop — and until this module existed the outcome of
a step was handled in the first one only. Everything a step outcome is supposed
to trigger therefore had EXACTLY ONE CALL SITE IN THE TREE, on the loop a TP>1
boot does not run:

* ``note_infra_failure`` / ``note_step_succeeded`` — so at TP>1 an OOM latched
  nothing. ``/health/ready`` stayed 200 while every request error-finished, and
  the twenty-line comment stating the opposite contract sat on a branch that
  could not be reached.
* ``note_step_served`` — so at TP>1 the serving forward-arena reading (and, on
  the same edge, the driver-growth reading) fell back to its only other trigger,
  the metrics gauge callback. That is verbatim the "the reading happens only if
  somebody scraped ``/metrics``" defect the watch was written to remove: with no
  scraper a pool growing past what the boot realized went unreported for the
  whole life of the process.

Both defects are the same defect. A fix that lands on one of two copies of a
branch is this repo's most repeated bug class, so the answer is not to copy the
branch into the other loops — it is for the branch to exist once and for every
loop to call it. That is :func:`note_step_failed` and :func:`note_step_ok`.

The loops keep what is genuinely theirs: finishing the slate, rolling back a
duplex tick, emitting the worker tick-done, deciding what "re-raise" means for
their own teardown. What they no longer own is the CLASSIFICATION or the
LATCHES, because those are a property of the step, not of the loop that ran it.

What the run loop must do with a step that failed.

``fail_loud`` means the loop finishes the slate with an error and then
RE-RAISES: the relevant sticky latch is already set, so ``/health/ready`` is
already 503, and the loop's exit is what turns that into a restart.
``False`` means the failure was absorbed — finish the slate, keep looping.

``kind`` names which arm of the classification decided, for the loops' own
logs and for tests that need to assert the classification rather than infer
it from a side effect.

Classify one failed engine step and apply the reaction.

``multi_rank`` is the caller's own answer to
:func:`~arbi_serve.engine.run_step.engine_is_multi_rank` — passed rather
than re-derived because the SPMD driver knows it structurally while the
single-rank loop reads it from the config, and because a seam that
re-derives the caller's own premise can disagree with it.

The order is the precedence of the facts, from the one no reaction can
survive to the one every reaction is for:

1. **memory exhausted** — the ladder in
   :mod:`arbi_serve.engine.memory_pressure` already ran out on some other
   path and latched. Nothing to classify: propagate.
2. **context poisoned** — an illegal address / device-side assert. No
   forward on this context can be trusted again; latch sticky and exit.
3. **out of memory** — the allocation did not fit. At multi-rank the ranks
   are desynced and it is fatal (the contract this seam finally makes run);
   single-rank it latches the self-healing infra window AND goes to
   :func:`~arbi_serve.engine.memory_pressure.note_oom`, which narrows the
   next step, closes the front door when there is nothing left to narrow,
   and escalates only when that did not help.
4. **other infra** — a cuBLAS handle that will not initialise, a generic
   CUDA error. Multi-rank fatal for the same desync reason; single-rank the
   self-healing infra latch, unchanged.
5. **per-request** — logged and finished, server stays healthy, unchanged.

A drafter circuit-breaker trip is NOT classified here: it is the drafter's
own verdict about the drafter, it carries no fact about the card, and each
loop already re-raises it before reaching this seam.

A full engine step SUCCEEDED. The other half of the same seam.

Three observations that all key off one causal edge — a step ran — and had
one call site between them:

* :func:`~arbi_serve.engine.infra_health.note_step_succeeded` clears the
  self-healing infra latch;
* :func:`~arbi_serve.engine.memory_pressure.note_progress` clears the
  memory-pressure state and, on the recovery edge only, re-reads the card;
* :func:`~arbi_serve.engine.arena_watch.note_step_served` takes the serving
  forward-arena and driver-growth readings.

Cheap on the healthy path by construction: each is one attribute or dict
read that early-outs before doing anything. Nothing here allocates, syncs,
or calls the driver on a step that had no failure to recover from and no
pool that mapped new physical.

The multi-rank never-wedge verdict, in one place for both arms.

A CUDA failure inside a forward that carries cross-rank collectives almost
certainly DESYNCED the ranks: the peers already finished the collective, or
are blocked in it, or are waiting on the SPMD shm control bridge. The
self-healing infra latch would "recover" and run the NEXT step, which then
hangs the shm writer ("waited 60s") or trips the NCCL watchdog (600s →
SIGABRT on both ranks) — the silent wedge this repo refuses. The rank state
cannot be reconciled in-band, so the caller latches the sticky fatal fault
and tears the loop down for an orchestration restart. Same never-wedge
contract as the mid-collective ``DrafterCollectiveFault``.

Per-step HOST phase probe — where the engine thread spends a step.

``ARBI_STEP_PHASE_PROBE=1`` (default OFF) turns the engine step into a
sequence of named :func:`mark` points and aggregates the wall-clock
between CONSECUTIVE marks. The output answers the one question a Kineto
trace answers only indirectly: *while the GPU was idle, what was the
host doing?*

It answers it more ACCURATELY than the trace, too, and that is the point
of using a probe rather than reading a profile. Kineto/CUPTI records one
kernel activity per captured-graph NODE per launch, on the host, INSIDE
``cudaGraphLaunch`` — so that call's measured duration scales with the
graph's node count and is instrumentation cost, not served cost. On a
1190-node decode graph that is ~0.6 us/node, i.e. ~590 us of purely
observer-induced host time per step, which a trace reports as GPU idle.
This probe attaches no CUPTI: it is two ``perf_counter_ns`` reads and an
int add per boundary, so the gaps it reports are the gaps production
actually pays. A phase total here that is far BELOW the same window in a
trace is the expected result, not a contradiction.

Why pairs and not fixed phases: the mark sequence differs per step shape
(MTP verify / MTP seed / plain K=1 / empty slate), and a fixed phase
table would silently mis-attribute a step it was not written for.
Aggregating by the ORDERED PAIR ``(previous_mark, this_mark)`` makes
every step shape self-describing, and the recorded SEQUENCES
(:func:`snapshot` ``["sequences"]``) name which shapes actually ran.

Reading a report:

  * ``step_begin -> sched_done`` — scheduler host cost (GPU idle, unless
    the previous step's kernels are still in flight).
  * ``drain_begin -> verify_drain_synced`` — the host BLOCKED on the
    parked verify step's D2H event. Host waiting, GPU busy.
  * ``verify_drain_synced -> drain_done`` — the deferred commit:
    truncation, draft-slot frees, ``post_token`` (detok + SSE) for up to
    ``(K+1) * B`` tokens. GPU IDLE.
  * ``drain_done -> verify_plan_built`` — verify plan build. GPU IDLE.
  * ``verify_plan_built -> verify_forward_done`` — the main-model
    forward. GPU busy (host blocked only when the offload path is on).
  * ``verify_drafts_host`` closes the drafter chain. With the host pull
    of the drafts (``ARBI_MTP_DEVICE_DRAFTS=0``) everything from there to
    the NEXT ``verify_plan_built`` is host-only work with an empty GPU
    queue; with device-resident drafts (the default) the chain is still
    running through that window and the next plan build overlaps it.
  * ``verify_drafts_host -> verify_rollback_done`` — the hybrid
    recurrent reconcile, which carries its own ``bool(any_partial)``
    D2H gate and two host pulls. ~0 on a pure-attention model.
  * ``drain_begin -> verify_flush_done`` — the staged-output handover,
    kept separate so a slow flush is not read as a GPU wait.

The "did not run" tell is unambiguous: a probe that never ran logs
nothing and :func:`snapshot` reports ``{"enabled": False, "steps": 0}``.
An enabled probe that saw steps but no MTP marks reports the step count
with only the generic pairs — "the MTP path did not execute", not
"nothing to find".

Cost when OFF: one module-global bool read per mark site.

Read ``ARBI_STEP_PHASE_PROBE`` once and latch it.

Defensive: a probe must never be the reason a boot fails, so a
runtime-flags read that raises (partially-built config in a unit
harness) latches OFF rather than propagating.

Record a phase boundary; accumulate the gap since the previous one.

``name == "step_begin"`` also closes the previous step's sequence and
advances the step counter (so the report is per-step, not per-mark).

Count a per-step boolean observation (e.g. "was the event ready?").

A pair that reports a host WAIT cannot say whether the GPU was busy or
the producer was merely late; a ``query()`` taken before the wait can.
Both outcomes are counted, so all-True and all-False are equally
legible and neither can be confused with the probe not running.

JSON-safe view of the accumulated pairs + step shapes.

``enabled`` is the "did the instrument run" discriminator: a caller
reading ``{"enabled": False}`` learns the probe was never armed, and
``{"enabled": True, "steps": 0}`` that it was armed and no step
reached a mark — neither is "nothing to find".

StepPlan + StepResult — public engine step contract.

Unifies the three batch-building paths (main K=1 ``ScheduledBatch`` in
:mod:`runtime.model_runner`, MTP verify in
:mod:`spec_decode.mtp`, external drafter
in :mod:`engine.build_external_drafter`) into one immutable plan dataclass
the scheduler emits and the model runner consumes.

The K dimension is always present:

  - ``K=1`` is the no-spec-decode path.
  - ``K>=2`` is MTP / ngram / external-drafter verify.

Whether the step is a verify pass is carried explicitly on
``is_verify`` so the run loop can branch without inspecting K.

Immutable, K-aware step plan.

All token tensors are flat (concat-along-batch) and laid out the
way :class:`AttnPagedKVMeta` and friends already expect today; the
``cu_seqlens`` cumulative-length tensor partitions the flat axis
back into per-request rows.

``state_meta`` is the per-:class:`StateKind` opaque dict the
backend's :class:`MetadataBuilder` consumes. The shape of each
value is backend-specific; the plan transports it without
inspection. The verify path (:meth:`from_verify_view`) populates
the full schema the per-step :class:`ScheduledBatch` reconstruction
in :meth:`EagerModelRunner.forward_plan` consumes:

  - ``StateKind.PAGED_KV`` → ``{block_table, slot_mapping,
    seq_lens, cu_seqlens_k, max_seq_len, max_query_len, causal,
    sliding_window, is_prefill}`` — the paged-KV builder reads
    the first three; the scalar attn fields rebuild the
    :class:`ScheduledBatch` routing header.
  - recurrent kind (``StateKind.GDN`` / ``StateKind.MAMBA`` /
    ``StateKind.SHORT_CONV``) → ``{state_indices}`` — the per-row
    recurrent slab-row mapping the recurrent builders consume
    (``None`` value on pure-attention models, where the entry is
    absent entirely).

The K=1 main path keeps ``state_meta`` empty today — the runner's
``_build_batch`` owns that path's metadata; only the verify path
threads metadata through ``state_meta``.

Frozen so it can be passed across the rank-0 → worker broadcast
seam (:class:`DistributedEngineDriver`) without surprise mutation
en route.

The plan is built lazily from a per-step slate
(``[(req, n_tokens), ...]``), with the runner / spec strategies
consuming the slate-shaped ``ScheduledBatch`` internally. ``slate``
is retained as a private carry-through field; consumers should
treat it as opaque.

Per-step output: sampled tokens + acceptance lengths + finish flags.

``sampled`` is a ``(K, B)`` tensor on device — for ``K=1`` paths the
``K`` dim is just 1; for spec-decode verify the row is the per-K
accepted sequence (untruncated; ``accepted_lengths`` carries the
per-row cutoff).

``accepted_lengths`` is ``(B,)`` (int32). Always-1 for the ``K=1``
path; for verify it is the number of K-step tokens the verifier
accepted before the first rejection (so the engine can commit
``accepted_lengths[i]`` tokens and roll back recurrent state past
that point).

``finished`` is the list of request IDs whose terminal condition
fired this step (EOS / max_tokens / stop string / cancellation).

Return the ``(kind, payload)`` for the recurrent entry, if any.

The verify plan stores the per-row recurrent slab-row mapping
under the model's recurrent :class:`StateKind`
(``GDN`` / ``MAMBA`` / ``SHORT_CONV``); pure-attention plans
have no such entry. Both the verify-forward
:class:`ScheduledBatch` reconstruction and the post-forward
partial-accept rollback resolve the slab rows through here so
they read the same mapping the forward wrote.

Build a :class:`StepPlan` from a legacy ``[(req, n_tokens)]`` slate.

Convenience constructor: the scheduler emits a slate; the
runner builds tensors from it. The plan carries the slate
through as an opaque payload and exposes the K / verify /
spec_mode top-line so spec strategies dispatch on metadata
without reaching back into the slate.

The actual flat tensors (input_ids / positions / cu_seqlens)
are not built here — the runner's ``_build_batch`` path owns
that. We populate placeholder zero-length tensors so the
frozen invariant holds; the strategy / runner path resolves
them from ``slate``.

Build a :class:`StepPlan` from a staged ``VerifyBuffers`` view.

The MTP verify path stages flat tensors into a persistent
:class:`VerifyBuffers` set; this constructor wraps that view
as a :class:`StepPlan` so the public engine boundary sees the
same K-aware shape MTP-on and MTP-off paths emit.
``state_meta[StateKind.PAGED_KV]`` carries the
``block_table`` / ``slot_mapping`` / ``seq_lens`` triple the
:class:`PagedKVMetadataBuilder` consumes plus the scalar attn
routing fields (``cu_seqlens_k`` / ``max_seq_len`` /
``max_query_len`` / ``causal`` / ``sliding_window`` /
``is_prefill``) the verify-forward
:class:`ScheduledBatch` reconstruction
(:meth:`EagerModelRunner.forward_plan`) needs — so no staged
``ScheduledBatch`` payload rides along. Verify-only fields
(per-row K, draft mask, draft slots) live on the strategy-
side :class:`MtpStepPlan` adjacent to this plan.

``recurrent_state_indices`` (when the model has a recurrent
kind) carries the per-row real pool slab-row mapping resolved
at plan-build time; it is stored under ``recurrent_kind`` so the
recurrent metadata builder + the post-forward rollback read the
same slab cells the verify forward wrote. ``None`` on pure-
attention models — the recurrent entry is then absent.

``is_verify`` defaults to True for the spec-decode path; the
cold-path branch (drafter cache miss / draft-slot allocation
failure → every row at K=0) passes ``is_verify=False, K=1``
so the dataclass invariant holds.

API-side submission: build a :class:`SubmitMsg`, await the SubmitAck.

P3 intake split (``docs/engine_core_process.md`` §3/§5.1): ``asubmit``
always builds a typed :class:`~arbi_serve.engine.proc.messages.SubmitMsg`
and the engine-side :func:`~arbi_serve.engine.request_factory.build_and_publish`
(RequestFactory) reconstructs the :class:`Request` from it on the engine
loop. This module is the API half:

* **request-id allocation** happens here (API side, ``eng._next_req_id``
  — the per-process allocator; in process mode (P4b) the API child owns
  its own counter). ``SubmitMsg.request_id`` is client-set; the factory
  never allocates.
* the consumer half (:class:`ClientRequest`) is created api-side and
  registered with the OutputApplier before the submission is marshaled,
  so the engine can never emit for an id the applier doesn't know. The
  SSE chat handler pre-creates it via :func:`new_client_request` to
  attach its tool stream race-free; every other caller lets ``asubmit``
  create it.
* the **SubmitAck**: thread mode marshals the factory call onto the
  engine loop through the existing coalesced intake
  (``LoopBridge.run_on_engine_coalesced`` — a burst of N submits still
  collapses to ~1 engine wakeup) and awaits an asyncio future keyed by
  ``request_id`` (:class:`~arbi_serve.engine.loop_bridge.SubmitAckRegistry`
  — transport-agnostic: P4b resolves the same future from the wire-ack
  recv path). Inline mode awaits the factory directly. A validation
  failure resolves the ack with
  :class:`~arbi_serve.engine.client_request.SubmitRejected`
  (a ``ValueError`` subclass carrying the HTTP status class).

In-proc by-reference payloads (thread/inline): the pre-made
``ClientRequest``, the multimodal features dict and the rich
``TenantContext`` cannot ride a msgspec struct, so the bridge intake
closure captures them alongside the message — see
:mod:`arbi_serve.engine.request_factory` (module docstring) for the
contract and the process-mode (P4b) replacements.

Register the consumer half with the engine's OutputApplier.

The applier owns ``{request_id -> ClientRequest}``; every emitted
``TokenOut`` / ``AudioOut`` / ``FinishOut`` resolves through it, so
a request must be registered before it is submitted — this runs on
the API side, strictly before the marshaled publish. Tolerates a
partial / stub engine (no applier attribute) for direct-construction
test harnesses.

Tokenize (API-side, §4). HTTP routes already pass ids; the str
path serves in-proc callers (bench / calibration / scripts). In-proc
this is the engine's tokenizer instance; the process-mode API child
(P4b) loads its own from the model dir.

Prompt-id ownership rule. The list returned here is the ONE prompt
buffer of the submission: the tokenizer's fresh list, or the caller's
ids copied exactly once (so a caller that reuses its own list cannot
reach into the request). From here it is handed by reference to the
:class:`ClientRequest`, the :class:`SubmitMsg` and the engine
:class:`Request` — in process mode the wire encoding is the copy the
engine side receives. No holder mutates it in place: the engine
rebinds ``Request.prompt_token_ids`` to a new list when it needs a
different prompt (preemption re-admit), which leaves every other
holder's reference untouched.

Allocate a request id + build the consumer half (API side).

For callers that need the :class:`ClientRequest` before submitting —
the streaming chat handler attaches its ``ChatToolStream`` to it so
the extractor exists before the engine can emit a single token
(attach is fully API-side per P3). Pass the result to ``asubmit(...,
client=...)``; registration with the applier still happens inside
``asubmit``.

Mirror the per-request intake onto the wire struct.

``kind`` (single source of truth for the task mode) comes from
``sampling.task``; ``SamplingMsg.task`` is still written for wire
compat but the factory reads only ``kind``. ``media`` stays empty
in-proc (the features dict crosses by reference; the process path
passes the :mod:`arbi_serve.engine.proc.media` MediaRef list).
``lora`` / ``mtp_k`` ride in ``sampling`` (their SubmitMsg mirrors
stay unset — one source).

``client_request_id`` is read HERE rather than passed in by each caller:
this is the one API-side seam every submission crosses, and it runs inside
the HTTP handler's context, where
:func:`~arbi_serve.server.middleware.get_request_id` resolves the
``X-Request-ID`` the middleware bound for the call. Reading it per caller
would make the correlation depend on which handler remembered to pass it.
``None`` outside a request scope — an in-proc submission has no caller to
correlate to, and a synthesised id would point at nothing.

Tokenize + enqueue a plain request (sync compat entry).

In-proc only (bench / calibration / scripts / distributed driver).
Rejects grammar-bearing requests — compilation is CPU-bound and must
run off the event loop; use :func:`asubmit`. Builds the same
``SubmitMsg`` + ``ClientRequest`` as ``asubmit`` and runs the
factory's sync core on the caller thread, marshaling only the
publish body onto the engine loop.

Async submission through the message-shaped intake (P3).

Builds the :class:`SubmitMsg` (allocating the request id API-side),
registers the ``ClientRequest`` with the applier, marshals the
engine-side RequestFactory onto the engine loop, and awaits the
SubmitAck. Returns the engine :class:`Request` handle (in-proc
consumers use ``handle.client`` / ``new_token_event``).

Raises :class:`SubmitRejected` (a ``ValueError``) on any pre-publish
validation failure — prompt length, vocab range, mtp_k gates,
unknown LoRA, grammar compile.

``target_resident`` (stable-VA residency): the resident key
``_maybe_switch_model`` resolved — the factory re-validates it against
the actual active resident and switches when a concurrent/stale
routing decision skipped it (P6 stale-cache guard). ``auth_token`` /
``tenant_id`` /
``tenant`` bind the request's identity before publish (P3 fixes the
post-submit ``_bind_auth`` race). ``client`` lets the SSE chat
handler pass a pre-made consumer half (tool stream already
attached); default ``None`` creates one here.

Thread-mode ack round-trip: coalesced intake -> engine factory ->
ack future resolution back on this (HTTP) loop.

The intake closure runs on the engine loop inside the coalesced
drain. The common factory has zero awaits, so the closure drives
the coroutine to completion synchronously — the publish happens
inside the single drain hop (burst coalescing + FIFO preserved). A
request the factory will await for (grammar compile; a P6 pending
resident switch) is wrapped in an engine-loop task instead; its ack
resolves when the await finishes. ``submit_will_await`` decides —
computed inside the closure, on the engine loop, so the resident
check can never race a concurrent switch — and if the awaitless
drive ever suspends anyway (someone added an await without
updating the gate), the ack resolves with a loud RuntimeError,
never a hang.

Process-mode submit handle: the client half only (design §3).

Routes consume ``handle.client`` (the applier-fed ``ClientRequest``)
and ``handle.request_id`` (the cancellation wrapper) — nothing else.
The engine ``Request`` lives exclusively in the engine process and
never crosses (§5.3), so this is the entire handle surface.

API-child submission: SubmitMsg over the wire, await the ack (P4b).

``api`` is the process facade
(:class:`arbi_serve.engine.proc.api_engine.ProcEngine`) — it owns the
child's request-id counter, OutputApplier, SubmitAckRegistry, and the
:class:`~arbi_serve.engine.proc.client.EngineClient`. Signature
mirrors :func:`asubmit` so routes call ``eng.asubmit(...)`` unchanged;
the by-reference kwargs translate per the request-factory contract:

* ``multimodal`` (features dict) → ``SubmitMsg.media`` MediaRefs
  (:func:`arbi_serve.engine.proc.media.features_to_media_refs`);
* ``tenant`` (rich context) does not cross — the engine reconstructs
  a minimal :class:`TenantContext` from ``tenant_id``;
* ``target_resident`` crosses on the message (P6): the routing
  decision here was made against the child's cached resident view;
  the engine-side factory re-validates against the actual active
  resident and switches itself if the cache was stale (a request
  naming member B is never silently served by member A).

Raises :class:`SubmitRejected` on a negative ack (engine-side
validation) and on an already-dead engine (503).

Hot-swap admin operations: per-StateKind backend, calibration, full model.

Each operation goes through the shared drain / mutate / resume
:class:`arbi_serve.engine.critical.CriticalSection`; this module is
a thin layer of admin-shaped helpers over a single ``async with``.

  - :func:`aswap_attention_backend` — kind-scoped backend swap.
  - :func:`areload_calibration` — TKV calibration JSON re-bind, no
    pool teardown (named-buffer ``copy_()`` semantics; sub-50ms).
  - :func:`areload_model` — whole-engine reload via ``ALL_KINDS``.

The refusal both capacity gates raise, built in ONE place.

SUSPENDED IS NOT EXITED, and the message has to say so. The fail-safe stops
serving but keeps the process alive holding every byte of VRAM it had, so
an operator who reads "suspended" and starts a replacement engine against
the same card gets a CUDA OOM whose text points at the memory budget --
a plausible, wrong diagnosis that costs a debugging session. Naming the
residency here is the only place it can be read at the moment it matters.

One builder, two call sites: the same sentence living twice is how one copy
gets corrected and the other keeps telling the older story.

The refusal for a swap that failed WITHOUT costing the engine a member.

Distinct from :func:`_swap_refused` because the operator question it
answers is the opposite one: nothing is suspended, nothing needs
correcting before serving resumes, and the only thing that did not happen
is the swap. Naming the backend that is still active is the whole message —
a caller that reads "refused" and assumes the target took effect anyway is
the one failure mode this text exists to prevent.

Kind-scoped backend swap. ``spec`` is a ``"kind:name"`` spec.

The target need not be pre-declared at boot — a runtime-introduced backend
is built on demand as a prepared variant (and validated then: an
unbuildable backend, e.g. uncalibrated TKV, fails the prepared-member
build and leaves the active member serving). Backend is a runtime variable,
not a boot-fixed option.

``calibration_path`` attaches a TKV calibration bundle in the same config
variant as the backend switch — the ``calibration_path`` override alias sets
``cfg.cache.calibration_path`` so the prepared-member build loads it via
``build_phases_load`` (exactly as a boot with ``--calibration`` would). This
is how the in-process boot swaps a freshly-calibrated bypass engine onto the
TKV codec: without it the variant rebuild would re-derive an empty
calibration and the TKV build would refuse as uncalibrated.

``(free_pages, available_pages)`` for the paged-KV pool.

``free_pages`` is the pool's free list. ``available_pages`` adds the
prefix-cache pages the evictor reclaims on demand, read off the page
table -- the same view the admission gate uses. A page table that does
not publish one (a flat table, a test double) leaves the two equal.

Drain in-flight, re-bind TKV calibration buffers in place, resume.

Named-buffer ``copy_()`` semantics: existing per-layer centroid /
boundary / per-channel-scale tensors are rewritten in place; no
pool teardown, no recompile, no graph invalidation. Sub-50 ms in
steady state — drain dominates only when many requests are
in-flight at the moment of the swap.

The bit-width fail-fast guard inside
:func:`apply_calibration_to_tkv_ops` rejects a mismatched
single-bit JSON loudly. ``min_improvement_pct=0`` accepts any
bundle; higher values run the same gate as ``--autotune``. A live
``/v1/admin/config`` ``min_improvement_pct`` raises the FLOOR: the
stricter of the two applies, so a caller cannot weaken the gate the
operator set.

Re-arm stable-VA residency on whatever model the engine now serves.

One writer, so the successful reload and the rollback cannot re-engage
differently — the residency key is read off ``eng.cfg`` at call time, which
is the config that actually built.

Step 1 of the reload releases the old weights before the new ones load
(single GPU, no double-buffer), so a failed build leaves NO model, no pool
and no scheduler. ``/health/ready`` 503s, which withholds traffic but is not
a recovery: nothing in the process would ever build a model again, and
nothing would ever exit. This function is the recovery, and it always
raises — every exit from a failed reload states what the engine is now.

Three outcomes, in the order they are checked:

* The build POISONED the CUDA context (illegal access / device-side assert).
  Rolling back would re-fault or hang on the dead context — the same reason
  :func:`~arbi_serve.engine.config_variant._rollback_member_build` is skipped
  there. Latch the fatal fault and arm the process exit; a poisoned context
  is only recoverable by restart, and the timer that exits takes no CUDA lock.
* The previous model REBUILDS. Its weights are still on disk and ``prev_cfg``
  is the config that demonstrably built them, so this is the reload's own
  bring-up re-driven at the old config. The engine serves again exactly as
  the member/variant seam leaves the previous member serving, and the caller
  gets the original failure as a 4xx.
* The rollback ALSO fails. The process cannot serve and has no way back, so
  it exits with :data:`~arbi_serve.engine.shutdown_request.EXIT_CODE_UNRECOVERABLE`
  instead of lingering as a permanently-503 shell. The exit is deferred so
  the reply carrying the reason reaches the operator first.

Whole-engine reload — drain everything, free, replace cfg, rebuild.

Single-GPU constraint: cannot double-buffer the model, so old
weights must come down before the new ones are loaded. Worst-case
the engine sits with no model for a beat between teardown and the
new :func:`load_model` call so health probes can observe the
transition.

``backend_specs`` optional override — when None we keep the
existing registered backends. ``calibration`` similarly nullable;
an empty string clears the global default. Static config
(parallel topology, dtype, device, port) is not reconfigurable
via this path; restart the process to change those.

Single-VRAM-resident invariant: enforce co-residency==1 at every residency
park.

The stable-VA residency model requires that only the active member's physical
is mapped: a park unmaps the parked member's named-pool (VaArena) physical so
the serving member sees its full boot-time free-VRAM budget (see
``docs/instant_config_swap_design.md``). The one thing a park cannot reclaim is
memory the default caching allocator reserved outside any cuMem named pool: the
allocator has no unmap-physical-keep-VA primitive, and ``empty_cache`` only
returns fully-free segments — an expandable segment whose high-water was driven
up by a transient (then freed, but pinned below the water line by a longer-lived
block) stays reserved. Such bytes are out of arena and survive park.

The GiB-scale instance of this leak class was the AWQ-INT4 head-quant buffers
stranding in the default allocator; the fix migrates them into the
``model.weights`` cuMem VaArena (``_apply_head_quant_phase``, see
``build_phases_load.py``). Routing the eager GDN chunk-fwd
workspace through an arena pool instead is a confirmed non-fix: GPU
verify showed such routing degrades to unrouted under the serving
allocator-context nesting.

Two seams:

  * :func:`route_gdn_workspace_for_serving` — clears the eager GDN chunk-fwd
    workspace pool pointer to the default allocator on every wake, so the
    process-global ``_GDN_WORKSPACE_POOL`` never carries a stale reference to
    a previous member's pool.

  * :func:`assert_single_member_resident` — a fail-loud co-residency check at
    park: after the park unmaps + settles, the out-of-arena default-allocator
    residual must be within a tight tolerance. Any future out-of-arena growth
    trips here, naming the offending bytes, rather than as a wake-time IMA.

:func:`probe_swap_coresidency` is the off-by-default attribution instrument
(per-pool + out-of-arena residual + free VRAM), free when off.

Yield ``(owner, named_pool_registry)`` for every resident member.

The active member is ``eng.named_pools`` (owner ``"active"``). Each other
residency record carries its member's registry in the captured engine-state
snapshot (``rec.extra["state"].attrs["named_pools"]``). A single-model boot
(no residency controller) yields just the active registry.

Return the free-VRAM / reserved / out-of-arena residual breakdown, or
``None`` on a CPU / no-CUDA path.

``out_of_arena_reserved`` = total torch-reserved bytes minus the sum the
cuMem-backed named pools account for = what the default caching allocator
reserved outside every member VaArena. This is exactly the term a residency
park cannot unmap (mirrors the ``residual_reserved_bytes`` reconciliation in
``memory_live_snapshot``). Best-effort: any probe failure returns ``None``
rather than perturbing the park.

Log the per-pool + out-of-arena residual + free-VRAM breakdown at a swap
seam. No-op (and no snapshot taken) unless ``ARBI_SWAP_CORESIDENCY_PROBE``.

Called at the residency park (after unmap+settle) and at the growable-KV
wake-budget seam — the two points that decide whether the woken member gets
its full slab back. Free when off (the guard short-circuits before any CUDA
call), so it never touches the serving hot path.

When ``ARBI_SWAP_CORESIDENCY_SNAPSHOT_DIR`` is also set (attribution-only —
never the serving hot path, this runs at the drained park seam), a
full segment-level breakdown (per-pool allocated-vs-reserved, the
out-of-arena residual's allocated-vs-reserved split, top live allocations
with call-site frames when recording is on, and the allocator config) is
written to ``<dir>/<phase>.<ts>.json`` — the definitive proof of what the
out-of-arena bytes are (live persistent tensor vs freed-but-unreclaimed
expandable-segment high-water).

Write the attribution artifact for this park seam, or no-op.

No-op unless ``ARBI_SWAP_CORESIDENCY_SNAPSHOT_DIR`` is set. Best-effort:
any failure is swallowed (an attribution dump must never perturb a park).
Uses the same ``live_snapshot`` / ``top_allocations`` machinery the
``/v1/admin/memory/live`` endpoint does, so the numbers reconcile with it.

Drive the out-of-arena residual down toward ``tol`` before the guard trips.

The residual right after a workload is dominated by RECLAIMABLE transients —
just-finished request activations + the eager kernel workspace — that "read 0
at idle" but have not settled by the time a park measures: their tensors are
freed but a lingering Python ref keeps the caching-allocator segment live, so
a single ``empty_cache`` cannot return it to the driver. ``gc.collect()``
drops those refs; ``synchronize`` completes any async frees; ``empty_cache``
then returns the now-free segments. Retry a few times (park is not the hot
path) until the residual is within tolerance or the attempts run out — a
genuinely STUCK (live) residual survives every pass and still trips the guard.
Returns the final measurement (or ``None`` on a no-CUDA path).

MAX ``value`` across every WORLD rank (int64 all_reduce); identity at
world_size 1 or with no live process group (CPU tests).

The co-residency verdict MUST be rank-symmetric. The guard runs inside a
park, which at TP>1 happens BEFORE the swap's cross-rank commit/abort
barrier (``confirm_build_across_ranks``). A per-rank raise (one rank over
tolerance, another under — asymmetric load transients) makes the tripping
rank leave while its peers block in the next collective → NCCL watchdog →
SIGABRT (whole engine dead). All-reducing the residual so every rank sees
the SAME max makes them raise (or proceed) TOGETHER — a clean symmetric
refusal that cannot desync. Gated on world_size (not tp_size): every world
rank runs the same broadcast swap, so agreement must span the world.

Fail loud if a park left out-of-arena (default-allocator) bytes resident.

The single-VRAM-resident invariant: after a park unmaps the parked member's
named-pool physical (and ``empty_cache`` settles the default allocator),
nothing should be reserved outside the member VaArenas beyond the small,
bounded driver/cuBLAS floor. Anything larger is memory a park cannot reclaim
— it will steal the slab the waking member needs. Trip here, naming the
bytes, instead of as a wake-time decode IMA.

A residual OVER tolerance first gets a bounded reclaim-retry
(:func:`_reclaim_out_of_arena`): the transients from a just-finished workload
are reclaimable but may not have settled at measure time (they "read 0 at
idle"), so gc + sync + empty_cache a few times before tripping. This is what
makes a rapid A/B flip robust under load without depending on a client-side
settle guess. A genuinely stuck (live) residual survives the retry and trips.

No-op unless ``ARBI_SWAP_CORESIDENCY_ASSERT`` (default on). Best-effort
measurement: a probe that cannot read CUDA does not raise (nothing to
assert). Park is not the hot path, so the retry + ``mem_get_info`` + pool sum
this costs is free of any serving-perf concern.

Clear the eager GDN chunk-fwd workspace to the default allocator for
GDN + ``split_attn`` models.

The single DRY seam for split_attn GDN-workspace routing, called from the
boot/build post-profile path and every residency wake, so the process-global
``_GDN_WORKSPACE_POOL`` never carries a stale pointer at a previous member.

Routing this workspace through the active member's
``scratch.forward_arena`` named pool (arena-backed) so a residency park
could unmap it is a confirmed non-fix: GPU verify (27B-AWQ RTX 4090) showed the serving eager GDN op
runs while a FOREIGN pool ``begin`` is already active, so begin-based
routing degrades to unrouted (it cannot nest) — it never reclaims
the park-time out-of-arena residual.
The real fix migrates the AWQ-INT4 head-quant buffers into ``model.weights``
(see ``ARBI_SWAP_CORESIDENCY_ASSERT``).

Contract / ordering:
  * Only for GDN + ``split_attn`` models. On the monolithic
    (``split_attn=off``) path the workspace is capture-referenced and pinned
    to ``scratch.gdn`` at the capture site — this seam leaves that alone.
  * Best-effort — never blocks a boot or a switch.

TKV decode-autotune as a prebaked, introspectable artifact.

The decode split-K autotune persists one JSON table per (gpu, H_kv, head,
page, num_sms, k_bw, v_bw, ...) fingerprint under ``$TKV_CACHE_DIR/
splits-tune``. Each ~90 s sweep is paid once; a matching table on disk turns
every later boot's warm into a cache hit (0 fresh sweeps). This module turns
that cache into an operable artifact:

  * :func:`read_autotune_tables` — load + structure every persisted table;
  * :func:`format_autotune_report` — operator-readable introspection of what
    split config each served fingerprint resolved to;
  * :func:`prebake_summary` — count fingerprints already baked vs needed, so a
    CI/image-bake step can warm them once and ship the tables in the image,
    eliminating the first-serve sweep on every fresh ``docker compose up``.

Parse every persisted splits table into a structured record.

Each record: ``{file, gpu, k_bw, v_bw, h_kv, head_dim, num_sms,
picks: {(batch, bucket): (kind, splits, tile_tokens)}}``. Corrupt /
unreadable files are skipped (never raises).

Which (k_bw, v_bw) fingerprints are already baked on this host.

``needed_bits`` is the distinct per-layer (k,v) set a deployment will
serve (e.g. from a smart-mix bundle). Returns baked / missing so a bake
step knows whether the shipped image already covers the model.

Warm turbo-attn's decode split-K autotune before cudagraph capture.

Because the decode kernel is captured into a cudagraph, an unwarmed autotune
table makes the dispatcher fall to ``_static_splits_heuristic`` — a
context-blind, fixed split count — and that config is then frozen into the
captured graph for the life of the server. Warming the table before capture
lets the graph bake the autotuned-optimal split config instead, so every
later decode is a warm-table cache hit.

This module assembles a :class:`tkv.runtime.autotune.WarmLayerSpec` from each
per-layer :class:`~arbi_serve.backends.tkv_backend.TkvAttnOp` (its installed
codec + byte layout) and the pool's live per-layer slab, dedups by
``(k_bw, v_bw)`` (so a smart-mix bundle's N combos each sweep once), and calls
the shared engine-agnostic ``warm_autotune_layers``. That also force-compiles
the per-(gqa, bits) decode/.compress-store ``.so`` so smart-mix layers never
JIT under capture (which can deadlock).

``mtp_block_ms`` carries the verify widths that reach the split-K register
kernel: the dispatchable set (:func:`_reachable_mtp_block_ms`) narrowed by
the route tkv RESOLVED (:func:`_plan_splitk_family`). The module is keyed by
``block_m`` and is compiled by no other boot path, so a width on that route
must be listed or the first verify batch JITs it under the request; a width
the Turbo prefill verify route owns must NOT be, because listing it spends
per-block_m nvcc builds and a measured split sweep on a kernel the process
never launches. A skipped width is guarded, not merely omitted — see
:mod:`arbi_serve.spec_decode.mtp_splitk_family`.

``warm_autotune_layers`` (turbo-attn >= 0.21.0) raises on a combo it cannot
pre-compile instead of warning and deferring it to the first request, and
returns per-combo coverage. So this module asserts ``loaded == reachable``
rather than reporting a reachable count it cannot confirm was loaded.

Wire-up: called once from the build path right before the cudagraph capture
sweep, after codecs are installed and the KV pool exists.

True for errors that already name the variant and the operator action.

turbo-attn's :class:`~tkv.runtime.autotune.WarmCompileError`, our own
coverage mismatch and the split-K family refusals all already state
exactly which combo or width failed and what to do about it. Re-wrapping
any of them in the generic "unwarmed table" diagnosis below would bury
that name behind the wrong cause.

The distinct ``(k_bw, v_bw, vq_k, vq_v)`` combos this boot can dispatch.

``warm_autotune_layers`` dedups its pre-compile by exactly this key, so
this set is the ``reachable`` side of the coverage assertion: every combo
in it must come back in ``WarmResult.combos`` (the ``loaded`` side).

The per-side quantizer family belongs in the key, not just the widths: a
group-2 VQ side compiles a DIFFERENT kernel than the scalar Lloyd one at
the same width, so a bundle that mixes them (K=vq2 everywhere, V=vq2 on
one layer) is two reachable combos that a widths-only key would report as
one — and a pre-load that built only the scalar of the pair would pass a
coverage check it should have failed, leaving the miss to surface much
later as a cache-miss abort under ``TKV_NO_JIT=1``.

Fail loud unless every reachable combo was actually pre-loaded.

Returns ``(n_loaded, n_variants)``. ``warm_autotune_layers`` raises on a
combo it cannot pre-compile (turbo-attn >= 0.21.0), so a short
``result.combos`` is not a compile failure — it means the primitive
silently skipped a combo serving can reach, i.e. the two sides disagree
about what is reachable. That is the drift the counter exists to catch, so
it must crash the boot rather than log a coverage number nobody checked.

Human-readable per-``(k_bw, v_bw, vq_k, vq_v)`` fingerprint lines.

``warm_autotune_layers`` dedups the sweep by exactly this key, so the
number of distinct fingerprints — not the layer count — is what drives
the (potentially minutes-long) nvcc JIT. Each fingerprint compiles the
per-(gqa, bits, family) decode + compress-store kernel variant set, so
naming them up-front makes a long compile obvious in the boot log.

MTP-verify ``block_m`` (= K+1) widths this boot can dispatch.

Derived from the same inputs the verify cudagraph capture enumerates —
:func:`arbi_serve.engine.cudagraph_admin._verify_s_set` over the MTP
driver's ``max_k``, the tkv served ceiling, and
``cfg.mtp.capture_full_k_ladder`` — so the pre-compiled variant set and
the captured verify set cannot drift. ``()`` when MTP is off.

The MTP split-K module is block_m-specific and is not covered by the
plain-decode variants: without these widths the first verify batch JITs
the kernel on the request path. ``compile_decode_and_cs`` clamps each
width to the split-K register ceiling — widths served by the Turbo prefill high-K
route carry no split-K ``.so`` and are skipped there.

This is the DISPATCHABLE set, not the set to compile: which of these
widths actually reaches the split-K kernel is the routing question
:func:`_plan_splitk_family` answers from the route tkv resolved.

The BUILT per-layer decode attend objects, if any core exists yet.

``TkvAttnOp`` builds its :class:`TKVCore` eagerly at boot, but a test
double (and a layer whose codec is not installed) has none — those
layers simply contribute no evidence, and the caller falls back to the
flag-derived floor rather than assuming an answer nobody gave.

Raise ``floor`` above any width a BUILT layer still keeps on split-K.

``tkv_prefill_verify_min_block_m`` answers for the process; this asks
the resolved per-layer objects — ``DecodeAttend.verify_routes_to_prefill``
is the same predicate the dispatch reads, so it also carries the
conditions the process-level floor cannot see (a core with no prefill
engine bound, a subclass that owns the verify body). A layer that
declines at width ``w`` pushes the floor to ``w + 1``, which puts ``w``
and everything under it back in the build set. Layers that cannot be
asked leave the floor alone.

Which of ``block_ms`` still needs the block_m-keyed split-K family.

The floor comes from tkv — :func:`tkv_prefill_verify_min_block_m` reads
``tkv.runtime.attention``'s own resolved rule — and is then confirmed
against the built layers. Never from the environment: this side must
not re-derive a routing decision the attention engine already made.

Best-effort path where tkv persists the compiled kernel .so + the
autotune splits table (for the loud start line). The .so land in the
torch_extensions build dir; the splits table under the tkv cache dir.

Assemble ``WarmLayerSpec`` per TKV attention layer.

``attn_ops`` is ``eng.attn_ops`` (mixed: only :class:`TkvAttnOp`
entries with an installed codec are used). ``slab_for_layer(layer_idx)``
maps a global layer index to its live packed-KV buffer (or ``None`` to
skip) — the same global index the pool's ``layer_view`` keys on. Pure /
engine-free so it can be unit-tested without standing up a full engine.

Return ``slab_for_layer(global_layer_idx)`` backed by the engine
pool's ``layer_view`` — the same aliasing-aware accessor the runtime
and cudagraph capture use to hand a layer its live per-layer KV view
(honours Gemma-4 ``kv_source_layer`` sharing). Returns ``None`` for a
layer the pool can't resolve (best-effort).

Warm the tkv decode autotune for every served layer.

Returns the number of distinct ``(k_bw, v_bw)`` fingerprints swept.
Returns 0 — an explicit, by-design skip — only when there is nothing
to warm: tkv not present, ``autotune_splits`` disabled, no TKV
attention layers, or no warmable specs. In those cases the static
heuristic is the correct (and only) config, so the no-op is silent.

If the sweep is required (there are warmable TKV specs) but the
autotune fails, this boot-fails (raises). An unwarmed table freezes
the context-blind ``_static_splits_heuristic`` into the captured
decode graph for the life of the process — a permanent throughput
loss invisible to the API caller. That degradation must crash the
boot (``/health/ready`` stays 503), never warn-and-serve.

Coverage is asserted, not assumed: every reachable ``(k_bw, v_bw)``
combo must come back in ``WarmResult.combos``, or this raises
:class:`TkvWarmCoverageError`. turbo-attn >= 0.21.0 raises on a combo
it cannot pre-compile rather than deferring it to the first request,
which is what makes ``loaded`` observable here at all — before that,
this line could only report ``reachable`` and hope.

Idempotent per-backend codec install: TKV via :func:`install_tkv_codecs`,
MLA via :func:`arbi_serve.engine.mla_install.install_mla_codecs`.

Guarded by a per-engine flag (``eng._tkv_codecs_installed_for``)
keyed on ``id(eng.attn_ops)`` so a redundant call against the
same ops list becomes a no-op. Re-runs whenever a fresh attn_ops
list replaces the prior one (e.g.
:func:`arbi_serve.engine.active.build_active` rebuilds the
production ops after the profile pass).

Returns:
    True when at least one TKV backend installed; False on
    short-circuit or when no TKV backend is registered.

Construct + install per-layer codec ops on each TKV AttnOp.

Source layers get a fresh :class:`TkvCodec`; calibration (if
any) is applied below. Gemma 4 ``num_kv_shared_layers`` (YOCO)
layers — those whose :attr:`LayerSpec.kv_source_layer` is set —
inherit the source layer's codec ops so that decode-side
centroid + rotation math matches the way the source layer
encoded the (aliased) cache slots.

Run this core's FIRST compress-store now, in the pool the core was built in.

``TKVCore.compress_store`` builds its
:class:`~tkv.runtime.compress_store.CompressStoreState` on the first store,
and the fused kernel derives a 4-bit vq2 side's grid LUT on the first launch
that indexes one. Both are model-lifetime, not per-step scratch: the launch
takes them by device pointer, every later step reads them, and a captured
decode/verify graph bakes their addresses. Left lazy they are minted by the
first FORWARD, in whatever pool is ambient there — ``scratch.forward_arena``,
where they are exactly the class of tenant
:func:`~arbi_serve.engine.inprocess_capture.release_idle_forward_arena`
refuses on, so the arena is never handed back, the serving mark goes
unmeasured, and KV is sized against the boot high-water instead. Nor can they
be declared in :mod:`~arbi_serve.engine.arena_step_caches`: evicting a table
a replay has baked is the one thing that seam must never do.

So the first store happens HERE, inside ``scratch.attn_codec``, beside the
codebooks the tables derive from — cuMem-backed with a stable VA and an
offload (not discard) sleep, which is the address stability a captured
launch already requires. It is tkv's own store path rather than a mirror of
it, so what the serving step finds memoised is what it would itself have
built, and the layer's own slab is what the kernel is handed. Every slot is
:data:`_COMPRESS_STORE_SKIP_SLOT`, which the kernel returns on before it
addresses a page, so nothing is written.

Returns True when a store ran: False for a KV-sharing layer, which never
stores, and for a build with no paged pool to store into.

Eagerly build every :class:`TkvAttnOp`'s ``TKVCore`` at boot.

Under ``torch.compile`` the model forward is traced and the
``arbi_serve::tkv_attention`` custom op is emitted directly into
the compiled graph; the Python ``TkvAttnOp.forward`` wrapper that
lazily built ``_core`` on first call is never executed at steady
state. Inductor invokes the op's real-impl straight from generated
code, which asserts ``_core is not None`` — so without an eager
boot-time build the very first TKV forward dies with an
``AssertionError`` (``_core is None``). Build the cores here, after
:func:`ensure_tkv_codecs_installed` (codecs + calibration must be
in place) and before the o_proj fold (so the fold can flip
``core._o_proj_folded`` on an already-built core).

Built inside ``scratch.attn_codec``, the pool that already holds the
codec's calibration tables and the MLA path's codecs. What a core
allocates — the migrated rotation matrices and WHT signs, the derived
rotation-state folds — is model-lifetime, read every step, and read by the
captured forward, so it belongs beside the tensors it is derived from
rather than in the default allocator, where it pins segments no reclaim can
release. The pool is cuMem-backed with a stable VA and offload (not
discard) sleep, so an address a captured graph bakes survives a sleep
cycle.

Each core's FIRST compress-store runs here too
(:func:`_prime_compress_store`), inside the same pool: the state and the
vq2 grid LUTs it builds are model-lifetime and would otherwise be minted by
the first forward, in the per-step arena.

``shared`` collapses the host-side aliasing that a per-attribute
``.to(device)`` would otherwise fork: the codec derives all four rotation
matrices from one process-global Hadamard, so without it every layer pays
four device copies of one matrix. Scoped to this call, so a device tensor
can never outlive the pool it was allocated in.

Returns the number of cores built (idempotent skips return 0).

MEASURE what an OSCAR basis made resident in ``scratch.attn_codec``.

The core build above materializes each layer's rotation state eagerly, and
under an OSCAR rotation that state is PER-KV-HEAD: up to five fp32
``(D_pad, D_pad)`` matrices per layer instead of the shipped Hadamard
path's shared ones. Which of them exist is path-dependent — the k side and
the v side are independent entries in the rotation file, ``_oscar_out`` and
``_oscar_inv_out`` are selected by ``_o_proj_folded`` and alias each other
when the source is already contiguous, and the per-(slot, dtype) casts
expand K/V heads to Q heads only for the slots a path actually reaches. So
the footprint has no closed form to predict, and
``predict_tkv_scratch_pool_bytes`` prices this pool's TQ buffers and
nothing else: the plan books the basis at zero and ``compute_kv_budget``
hands the difference to KV.

This reads the pool with the boot ledger's own meter
(:func:`~arbi_serve.engine.memory_budget.pool_residency.
named_pool_resident_bytes`) so the number is comparable, byte for byte,
against the row the freeze reconcile prints for it.

WHY THE WHOLE POOL AND NOT AN INCREMENT. The reading is the pool's
residency, not a delta across this call: ``build_active`` runs twice on a
boot and the second pass rebuilds every core into the same cuMem pool,
whose mapped bytes never shrink, so a delta would read ~0 on the pass that
matters. MONOTONE for the same reason — a later pass may only raise it.

ZERO, AND NO READING TAKEN, when no layer carries an OSCAR rotation. That
is what keeps a boot on the shipped Hadamard path byte-identical: its
rotation state is inside every measurement the current plan was validated
against, while the basis is a component nothing has ever held bytes for.

Returns the bytes recorded (0 when there is no basis to measure).

Drop every TKV core's retained fresh-diagonal chunk buffer.

The buffer is allocated by ``compress_store(emit_rotated=True)`` — inside
whatever pool is routing the forward, which at serving is
``scratch.forward_arena`` — and kept per layer so the next continuation-
prefill chunk re-uses it instead of allocating. Its CONTENTS are one step
deep: the hybrid attend consumes the stash and clears it, and a step that
declines re-derives the rotation rather than reading a stale one.

So the buffer is a memo, and this is the seam that gives it up: the arena
release (:func:`~arbi_serve.engine.inprocess_capture.
release_idle_forward_arena`), which destroys the pool the buffer lives in.
Keeping it there would forfeit every wholly-free segment in that pool —
the block is small next to what it holds hostage.

Sound at this seam and only at this seam. tkv declines the fresh-diagonal
path outright while a CUDA graph is capturing
(``_fresh_diagonal_eligible``), so no captured graph can have baked this
address; and the next allocation is one the post-release step was going to
make anyway, out of a pool that no longer exists.

Returns the number of cores that gave up a buffer.

Refuse a LIVE calibration reload that changes the vq2 codec config.

``vq_mode_k`` / ``vq_mode_v`` are compile-time defines (``VQ_MODE_K`` /
``VQ_MODE_V``) that select a different compiled decode / compress /
prefill kernel, and the KV pages already resident in the pool were
encoded by the codec that is live right now: the same four bits per
channel mean two scalar Lloyd indices under ``lloyd`` and one group-2
codebook index per channel PAIR under ``vq2``. Flipping the mode under
a populated pool would make every cached token dequantize to garbage,
and installing the codebook buffer for the first time reallocates it
(its ``data_ptr`` moves, which the reload path's whole named-buffer
contract exists to prevent).

No amount of graph-dropping fixes the resident-cache half, so this
transition is refused rather than implemented: it needs a full engine
rebuild, which the model-reload and backend-swap admin paths already
do (they re-enter this function with ``live_reload=False``).

A same-config codebook REFRESH (vq2 stays on, same shape, new
codewords) is allowed — that is the same class of change as new scalar
centroids, which live reload has always accepted.

Copy centroids from a calibration JSON onto the per-layer ops.

Deployable schema only; multi-bit bundle preferred; per_channel_scales
required. Bit-width validation is fail-fast — a mismatched single-
bit JSON would corrupt the codec because centroid counts differ
per bit width (K4 has 16/channel, K8 has 256/channel).

Bundle PARSING is not arbi-local: every layer goes through
:func:`tkv.runtime.calibration.parse_layer_calibration`, the same
validated parser the vLLM and SGLang integrations reach via
``apply_calibration_to_ops``. What stays here is the part that is
genuinely engine-specific — writing the parsed tensors into
pre-registered named buffers with :meth:`torch.Tensor.copy_` so a
live recalibration keeps every ``data_ptr`` stable (see the
calibration named-buffer contract on
:class:`arbi_serve.backends.tkv_attn_op.TkvAttnOp`). Reassigning the
codec's tensors the way the shared applier does would break that
contract.

``live_reload`` marks the ``/v1/admin/calibration`` path (via
:meth:`arbi_serve.engine.critical.CriticalHandle.apply_calibration`).
It rejects a bundle that would change the vq2 codec configuration on
a live engine; boot installs (the default) may introduce it freely.

TKV o_proj fold — eliminates the per-layer ``rotate_output`` cuBLAS GEMM.

Each TKV decode step ends with a per-layer ``rotate_output``
matmul that applies ``folded_v_R_inv = v_R_inv * v_per_channel_scale``
to the attention output before ``o_proj``. The math is associative:

    output = (attn_out @ folded_v_R_inv) @ W_o.T            (unfolded)
           = attn_out @ (folded_v_R_inv @ W_o.T) ... per head
           = attn_out @ W_o_folded.T                         (folded)

Per-head, the equivalent in-place rewrite of ``W_o`` (shape
``(hidden, num_heads * head_dim)``) is:

    W_o_folded[:, h*D:(h+1)*D] = W_o[:, h*D:(h+1)*D] @ R_v.T

where ``R_v = folded_v_R_inv = v_R_inv * v_per_channel_scale``. With the
fold applied at boot, the per-layer ``TKVCore.rotate_output`` is
skipped (gated by ``core._o_proj_folded``) — eliminating one cuBLAS
launch per layer per decode step.

**Calibration coupling.** ``v_per_channel_scale`` enters ``R_v``, so a
calibration reload that changes it stales the fold. We keep an
unfolded backup of every folded ``W_o`` (CPU, fp32) and re-fold from
the backup whenever ``TkvAttnOp.consume_capture_invalidation`` returns
True (the codec already raises that flag on a per-channel-scale
reload). The same hook the rotation-state cache uses.

**Skip conditions.** Architectures whose post-attention path is not a
pure linear (sigmoid gate, MoE-style mixing, etc.) cannot use the fold;
the per-head linearity is the load-bearing assumption. Today's only
gated arch in arbi-serve is Qwen 3.5 / 3.6
(:class:`_Qwen3_5AttentionBlock`). We detect by the parent attention
block's class name carrying a ``Qwen3_5`` prefix. Other skip reasons
(missing ``o_proj`` sibling, shape mismatch, codec missing required
ops) just log and leave the unfolded path live.

**Env gate.** ``ARBI_SERVE_TKV_O_PROJ_FOLD`` (default ``on``).

Locate the ``self_attn.o_proj`` sibling for ``layer_idx``.

Returns ``(o_proj_module, dotted_name, attention_block_module)`` or
``(None, "", None)`` if absent. The naming convention every
arbi-serve arch follows (per ``arbi_serve/loader/flat_loader.py``)
is ``...layers.<idx>.self_attn.o_proj``; we walk
:meth:`torch.nn.Module.named_modules` looking for that suffix.

Return ``(ok, reason)`` for whether ``attn_block``'s post-attention
path is pure linear (and hence can absorb the V-side rotation).

Blocks with ``attn_output_gate=True`` are refused: the per-head
sigmoid gate multiplied into the attention output before ``o_proj``
is non-linear and breaks the fold algebra. Non-gated
:class:`AttentionBlock` and the MLA / MoE blocks accept the fold.

Snapshot the pre-fold ``W_o`` onto the op as an fp32 CPU backup.

Used by the calibration-reload path to re-fold from the original
weight rather than chain corrections (each round-trip would compound
fp16 quantisation error). CPU storage keeps the backup off the
weight slab — one extra copy per attention layer at
``hidden * num_heads * head_dim * 4`` bytes (e.g. Qwen3.5-9B
relative to the model itself, and only allocated when the fold is
enabled).

Apply the o_proj fold for one :class:`TkvAttnOp` instance.

Idempotent: short-circuits when ``op._tkv_oproj_fold_applied`` is
True. Returns ``(folded, message)`` where ``folded`` is True iff
the weight was rewritten (False on idempotent skip OR any of the
skip reasons described in this module's docstring).

Apply the o_proj fold across every TKV layer on ``eng``.

Returns the count of layers actually folded. No-ops (returns 0) when:
  * :func:`env_fold_enabled` returns False; or
  * ``eng.attn_ops`` carries no :class:`TkvAttnOp` instances; or
  * every op short-circuits on an idempotent skip.

Called from :func:`arbi_serve.engine.active.build_active` after
:func:`ensure_tkv_codecs_installed` (codecs must be installed for
the ``v_R_inv`` / ``v_per_channel_scale`` lookups to find values).
Re-runs are safe — every per-op application is idempotent.

Re-apply the o_proj fold for every op whose calibration reload
invalidated the rotation state.

Reads :meth:`TkvAttnOp.consume_capture_invalidation` — same flag
the captured-graph drop path consumes. Called from
:func:`apply_calibration_to_tkv_ops` after the per-layer
``update_calibration`` loop, and from
:func:`arbi_serve.engine.swap_admin.areload_calibration` (via the
same path) so a live reload re-folds before the next forward step.

A True return from ``consume_capture_invalidation`` means
``v_per_channel_scale`` (or ``k_per_channel_scale`` for K-side, but
the same set of ops are affected because both scales are reloaded
together) changed; ``R_v = v_R_inv ``W_o`` is corrupted relative to
the new ``rotate_output`` math. Restore from the backup, then
re-fold with the fresh codec values.

Note: this reads but does not consume
``_capture_invalidation_pending`` — :func:`swap_admin.areload_calibration`
still owns the consume + captured-graph drop accounting. Reading
without consuming is intentional: a per-channel-scale change stales
both the codec rotation state's derived fp32 tensors and the folded
``W_o``; both invalidators must fire on the same reload.

Tool-call body parser registry + model auto-detection.

The active parser is **model-tied**: resolved at engine build from the
``--tool-call-parser`` config (default ``auto`` → sniff the chat template)
and stored on the engine, so a hotswap to a different model re-resolves it
through ``eng.build()``. See :mod:`arbi_serve.engine.tool_parsers.base`.

Return the :class:`ToolParser` for this model.

``override`` (the ``--tool-call-parser`` value) wins when it names a
registered parser. ``auto`` / ``None`` sniff the chat template for the
family's own call marker: NemotronVoiceChat's ``<TOOLCALL>`` →
``nemotron``; the Qwen3-Coder ``<function=`` XML surface → ``qwen3_xml``;
otherwise the JSON ``hermes`` default (whose ``<tool_call>`` marker is
also what a template with no tool markup at all is assumed to use).

The two sniffs are disjoint on every shipped template — NemotronVoiceChat's
renders no ``<function=``, and neither Qwen3 template renders ``<TOOLCALL>``
— so their order does not change any existing model's resolution.

Pluggable tool-call parsers: markers, body encoding, enforcement.

The marker framing — and its optimized, off-event-loop *incremental*
streaming scan — lives in :mod:`arbi_serve.engine.tool_parsing`. That scan
is format-agnostic machinery, but the marker STRINGS it scans for are
owned here: :attr:`ToolParser.tool_call_trigger` /
:attr:`ToolParser.tool_call_close` are read by the scanner (which caches a
compiled span regex per marker pair) as well as by the xgrammar
enforcement side, so a model family whose framing is not Qwen3's
``<tool_call>`` is parsed with its OWN markers rather than silently
mismatching.

What each family differs in:

  - Hermes / Qwen3 JSON:  ``<tool_call>{"name": "f", "arguments": {…}}</tool_call>``
  - Qwen3-Coder XML:      ``<tool_call><function=f>\n<parameter=k>\nv\n</parameter>\n</function></tool_call>``
  - NemotronVoiceChat:    ``<TOOLCALL>[{"name": "f", "arguments": {…}}, …]</TOOLCALL>``

The first two put exactly ONE call in each marker span; the third wraps a
JSON ARRAY of 1..N calls in a single span. :meth:`parse_calls` is the
general "one span → N calls" hook the scanner actually calls; its default
implementation wraps the single-call :meth:`parse_body` into a length-1
list, so one-call-per-span parsers implement only :meth:`parse_body` and
array-shaped families override :meth:`parse_calls`.

The enforcement side mirrors the same split: :meth:`build_tool_tags`
builds the whole tag SET for a request's tools, defaulting to one
:meth:`build_tool_body_tag` per tool (Hermes / Qwen3), and is overridden
by families whose wire format packs every tool into one tag.

Parse one tool-call marker span for a given model family.

Stateless and model-tied (selected from the chat template at engine
build, re-resolved on hotswap). ``tools`` is the per-request tool list
(engine ``ToolSpec`` objects or OpenAI-shaped dicts) — used for
schema-aware value typing where the wire format is untyped (XML).

Recursively drop xgrammar-unsupported keywords from a JSON schema.

Returns a cleaned copy; the input is not mutated. ``$defs`` / ``$ref`` are
kept (xgrammar resolves local refs); only unsupported validation keywords
are removed.

Look up a parameter's JSON-schema ``type`` for value coercion.

Tolerant of both engine ``ToolSpec`` objects (``.function.parameters``)
and OpenAI-shaped dicts (``{"function": {"parameters": ...}}``). Returns
the type string (first non-``null`` member if the schema lists several),
or ``None`` when unknown — callers then fall back to a safe default.

Parse one tool-call body into ``(function_name, arguments_json_str)``.

``arguments_json_str`` is a JSON *string* (OpenAI shape), even when
the source encoded an object. Returns ``None`` when the body is not
a complete/valid call (caller skips it — e.g. a partial stream).

Families whose marker span carries SEVERAL calls override
:meth:`parse_calls` instead; this method then reports the first of
them, so a direct single-call caller still gets a sane answer.

Parse one marker span's body into ALL the calls it names.

This is what :mod:`arbi_serve.engine.tool_parsing` calls per matched
span. The default is the one-call-per-span shape Hermes / Qwen3-Coder
use: :meth:`parse_body`, wrapped into a length-1 list (empty when the
body did not parse). NemotronVoiceChat's ``<TOOLCALL>[…]</TOOLCALL>``
array overrides it to explode the array into N entries.

Build the xgrammar ``TagFormat`` that forces one call into THIS
family's wire format — the enforcement inverse of :meth:`parse_body`.

``st`` is the ``xgrammar.structural_tag`` module, passed in so this
package never imports xgrammar (parser tests stay xgrammar-free).

Build the whole tag SET for a request's ``(name, parameters)`` list.

Default: one :meth:`build_tool_body_tag` per tool, which is the
Hermes / Qwen3-Coder shape (each call is its own marker span, so
each tool is its own alternative). A family that packs every tool
into ONE span — NemotronVoiceChat's single ``<TOOLCALL>`` wrapping a
JSON array — overrides this and returns a single combined tag.

Hermes / Qwen3 JSON tool-call body parser.

The body between ``<tool_call>`` and ``</tool_call>`` is a JSON object
``{"name": "f", "arguments": {...}}`` — the default for models whose chat
template does not use the Qwen3-Coder ``<function=...>`` XML surface.

Enforce one call as ``<tool_call>\n{"name": "NAME", "arguments": {...}}\n</tool_call>``.

The body is a JSON object whose ``name`` is pinned to this tool and
whose ``arguments`` follow the tool's parameter schema (``json``
style) — the exact surface :meth:`parse_body` reads.

NemotronVoiceChat ``<TOOLCALL>[…]</TOOLCALL>`` tool-call parser.

NVIDIA-NemotronLabs-VoiceChat-11B's chat template
(:data:`arbi_serve.chat_templates.NEMOTRON_VOICECHAT_CHAT_TEMPLATE`,
reproduced verbatim from the reference NIM container's
``s2s/prompt_template.jinja``) teaches the model exactly one call format::

    <TOOLCALL>[{"name": "tool_name1", "arguments": "tool_args1"}, {"name": "tool_name2", "arguments": "tool_args2"}]</TOOLCALL>

Two things separate it from every other family this package serves:

  * the markers are UPPERCASE and un-underscored (``<TOOLCALL>`` /
    ``</TOOLCALL>``), not Qwen3's ``<tool_call>`` — so a model served with
    the Hermes default would never have its calls detected at all; and
  * ONE marker span carries a JSON ARRAY of 1..N calls, where Hermes and
    Qwen3-Coder repeat the whole marker pair per call.

The realtime ``/v1/realtime`` path already reads this format off the
model's separate function-logits channel
(:class:`arbi_serve.engine.nemotron_function_channel.NemotronFunctionChannelExtractor`);
this parser is the same reading, expressed as a
:class:`~arbi_serve.engine.tool_parsers.base.ToolParser` so the
model-agnostic ``/v1/chat/completions`` path (which sees one interleaved
text stream, not a dedicated channel) gets it through the shared
incremental scanner in :mod:`arbi_serve.engine.tool_parsing`.

Argument shape. The template's own format line writes ``"arguments":
"tool_args1"`` — a placeholder, not a commitment to string-encoded
arguments; Nemotron Nano v2 (the underlying text LLM) emits an object.
:meth:`NemotronToolCallParser.parse_calls` therefore accepts BOTH (object
→ compact JSON string, string → passed through verbatim), matching the
realtime extractor's tolerance, while the enforcement grammar pins the
object form because that is the one a tool's parameter schema can
constrain.

OpenAI's ``arguments`` JSON *string* for one call's raw value.

Objects/arrays are re-encoded compactly, an already-string value passes
through verbatim (the template's own format line shows that shape), a
missing/``None`` value becomes ``"{}"``.

Explode one ``<TOOLCALL>`` body into every call it names.

The body is normally a JSON array; a bare object is accepted as a
one-element array (the model occasionally drops the brackets on a
single call, and the realtime extractor already tolerates it).
Entries without a usable ``name`` are skipped rather than failing
the whole span — the same "skip what doesn't parse" convention the
rest of this package uses. ``tools`` is unused: the JSON body is
self-describing, so no schema-driven value typing is needed.

One tag for the whole request: ``<TOOLCALL>`` + JSON array + ``</TOOLCALL>``.

Where Hermes / Qwen3-Coder get one tag per tool (each call is its own
marker span, so the tags are alternatives), this family's span holds an
ARRAY — so the enforced surface is a single tag whose content is an
array schema. Each element is pinned to one of the offered tools by a
``name`` const plus that tool's own parameter schema, so a 1-call and
an N-call block are both valid and both parse back through
:meth:`parse_calls`.

Enforce a single-tool block — :meth:`build_tool_tags` of one spec.

Kept so the base class's per-tool contract still holds for this
family (a caller that reaches for the single-tool hook gets a valid,
parseable one-element-array tag rather than a ``NotImplementedError``).

Qwen3-Coder XML tool-call body parser.

Qwen3.6's chat template (``chat_template.jinja``) renders tool calls as::

    <tool_call>
    <function=NAME>
    <parameter=K1>
    VALUE1
    </parameter>
    <parameter=K2>
    VALUE2
    </parameter>
    </function>
    </tool_call>

Crucially, the template writes each value as ``args_value | string`` for
string params and ``args_value | tojson`` otherwise, so a *string* value
appears RAW (no quotes) while ints/bools/arrays/objects appear as JSON.
Parsing therefore must be schema-type-driven — blindly
``json.loads``-ing a raw string value would corrupt it. This is the inverse
of that rendering, using each parameter's declared JSON-schema ``type`` to
coerce the captured text back to a typed value.

Coerce a captured XML parameter value back to its schema type.

Best-effort: any coercion failure falls back to the raw string so a
malformed value never crashes parsing (the downstream caller / model
sees a string argument rather than a dropped call).

Parse tool-call markers out of a model's output text.

Each model family frames its calls with its own marker pair, owned by the
family's :class:`~arbi_serve.engine.tool_parsers.base.ToolParser` and read
from it here — Qwen3 / Hermes::

    <tool_call>
    {"name": "func_name", "arguments": {"k": "v"}}
    </tool_call>

NemotronVoiceChat, whose single span carries a JSON ARRAY of 1..N calls::

    <TOOLCALL>[{"name": "f", "arguments": {"k": "v"}}, {"name": "g", "arguments": {}}]</TOOLCALL>

So both the marker STRINGS and the "how many calls does one span name"
question are parser-driven: the scanner asks the active parser for its
``tool_call_trigger`` / ``tool_call_close`` (caching one compiled span
regex per distinct marker pair, so a per-request parser costs no
recompilation) and hands each matched body to
:meth:`~arbi_serve.engine.tool_parsers.base.ToolParser.parse_calls`, whose
default wraps the family's single-call ``parse_body`` into a length-1 list.

This module is the OpenAI-shaped framing layer:

  - :func:`extract_tool_calls` consumes the final assistant text and
    returns ``(content_without_tool_calls, [openai_tool_call, ...])``.
  - :class:`StreamingToolCallExtractor` is a stateful *incremental*
    parser used by SSE streaming — it emits OpenAI-shaped tool-call
    delta dicts as they become available, only for *complete* tool
    calls (partial-args streaming is intentionally not supported).

OpenAI shape::

    {
        "id": "call_<uuid>",
        "type": "function",
        "function": {"name": "...", "arguments": "<JSON STRING>"},
    }

Note ``arguments`` is a JSON *string*, not an object — OpenAI compat
requires this even though it's awkward.

Performance note (the reason this parser is incremental). On the SSE
streaming path the extractor is driven once per decoded token. A naive
implementation that re-scans the entire accumulated ``output_text`` on
every call is O(n) per token = O(n^2) per stream, which would contend
the GIL with every other concurrent stream if run on the HTTP/ASGI event
loop. :class:`StreamingToolCallExtractor` instead runs on the engine
thread in the inline-detok path, staging results on the request for the
HTTP generator to pop (see :mod:`arbi_serve.engine.tool_stream`), and is
genuinely incremental: it tracks a scan offset and, while no open marker
has appeared (the common case — a stream that never calls a tool), takes
a zero-regex, zero-allocation fast path that only inspects the freshly-
appended suffix. Making the markers parser-driven does not cost that
property: the extractor resolves its family's marker strings ONCE in
``__init__`` and binds them to instance attributes, so the per-token fast
path does the same substring math on a local string it always did.

One family's marker pair, pre-derived for the scan.

``open`` / ``close`` are plain strings so the fast path can do cheap
substring math without touching the regex engine; ``open_len`` and
``span_re`` are the two derived values every scan needs, computed once
per distinct pair instead of per call.

Length of the longest suffix of ``text`` that is a *proper* prefix
of ``open_marker``.

Used to decide how many trailing characters must be held back from a
content delta: if the text ends with e.g. ``"<tool_"`` those chars
might still grow into a full ``<tool_call>`` marker on the next token,
so they cannot be emitted as content yet. Returns 0 when no trailing
fragment looks like the start of a marker.

Only proper prefixes (length 1 .. len(marker)-1) count — a *complete*
marker at the tail is handled by the marker-found path, not here.

Pull every complete tool call out of ``text``.

Returns ``(remaining_content, tool_calls)``. Tool calls are returned
in OpenAI-shaped dicts. Partial / unclosed open markers at the end
are kept in ``remaining_content`` (the model didn't finish; the
caller can decide whether to treat as "stop" or discard).

``parser`` selects both the marker pair and the body encoding (Hermes
JSON vs Qwen3-Coder XML vs NemotronVoiceChat's ``<TOOLCALL>`` array);
it defaults to the Hermes JSON parser for backward compatibility.
``tools`` is the request's tool list, used by schema-aware parsers
(XML) to type parameter values.

One marker span may name SEVERAL calls (the Nemotron array): each span
contributes as many OpenAI dicts as
:meth:`~arbi_serve.engine.tool_parsers.base.ToolParser.parse_calls`
returns, and is stripped from the content exactly once.

Incremental parser used by SSE streaming.

Multi-tenant safe: one instance per request (lives on the request
object, never shared). The parser tracks how much of the streamed
text it has already returned to the caller, plus any complete tool
calls already framed.

Partial-argument streaming is not implemented: the parser emits each
tool call in one chunk only when the family's closing marker arrives.
A family whose span names several calls (Nemotron's ``<TOOLCALL>``
array) therefore emits all of them on the token that closes the span,
with consecutive ``index`` values.

Returned chunks (one per call to :meth:`feed`):

  ``(content_chunk: str, tool_call_deltas: list[dict])``

where ``content_chunk`` is text outside any tool-call markers.

Incremental contract. :meth:`feed` is given the *full* accumulated
text each call (so it remains robust to the detokenizer retroactively
shortening ``output_text`` on a stop-string trim). It is engineered so
that, for the overwhelmingly common no-tool stream, it does O(new
chars) work with NO regex and NO allocation beyond the slice of brand
new content it returns:

  * ``_marker_seen`` flips True the first time a complete open marker
    is observed anywhere in the text. Until then the parser stays on
    the fast path.
  * Fast path (no marker yet): everything is content. We emit the newly
    appended suffix, holding back only a tiny tail that could still be
    the *start* of an open marker (a substring/prefix check, no
    regex). The held-back tail is re-emitted once the next token
    resolves it (either into content or into a real marker).
  * Careful path (marker seen): fall back to the exact full-text
    :func:`extract_tool_calls` framing over the "safe" region (text up
    to any final *unclosed* marker).

The careful path is byte-identical to the original full-rescan parser;
the fast path produces the identical *cumulative* content stream (same
bytes, possibly chunked at different per-token boundaries — and a stop
trim that shortens the text is handled by re-deriving from the full
text each call).

Quick check used by the engine to set ``finish_reason='tool_calls'``.

Scans for ``parser``'s own marker pair — the served model's, threaded
in from ``eng.tool_call_parser`` — so a family framing its calls as
``<TOOLCALL>…</TOOLCALL>`` is promoted just like a ``<tool_call>`` one.
Defaults to the Hermes markers when no parser is supplied.

No open marker anywhere: all text is content.

Emit the newly-appended characters, except a trailing fragment
that could still grow into this family's open marker (held back so
we never stream half a marker as content). Zero regex, zero
allocation beyond the returned slice.

Robust to a stop-string trim that *shortens* ``text_so_far``: we
clamp ``_content_emitted`` to the new length and only emit forward
progress, so a retroactive shorten never re-emits or mis-slices.

Full-text framing once a marker has appeared.

Byte-identical to the original full-rescan parser: find the start
of any unclosed marker span (content stops there), strip complete
tool calls out of the safe region, emit the cleaned-content delta
beyond what's already been sent, and frame each new tool call as
an OpenAI delta.

Consumer-side tool-call extraction + per-request SSE delta staging.

Tool parsing consumes DECODED TEXT — which already crosses the
engine->consumer boundary as ``TokenOut.text_delta`` — so it lives on
the consumer side, beside the SSE handler that drains it (proc design
§4). The :class:`~arbi_serve.engine.client_request.OutputApplier`, after
applying a batch's text deltas to a request's
:class:`~arbi_serve.engine.client_request.ClientRequest`, runs the
*incremental*
:class:`~arbi_serve.engine.tool_parsing.StreamingToolCallExtractor` and
STAGES the results — the cleaned content-delta plus any newly completed
tool-call deltas — onto a small per-request buffer
(:class:`ChatToolStream`). The SSE generator
(``arbi_serve.server.routes.openai_chat._chat_stream``) then simply POPS
the pre-computed deltas and emits frames, re-deriving nothing.

Ordering / correctness.

* The applier feeds the extractor only AFTER the batch's text deltas
  (including any ``trim_to`` stop-string retro-trim) are applied to the
  client ``output_text``, so a delta a stop-trim removed is never
  staged. The extractor is fed the full (already-trimmed) text each call
  and is itself stop-trim safe.
* Feed and drain both run on the CONSUMER loop (the applier applies
  there in every marshaler mode), so the ``deque`` hand-off is
  loop-local; the ``new_token_event`` the applier fires after the feed
  is the drain signal — staging adds no new wakeup.

Gating / cost. The buffer is attached to a ClientRequest only for
STREAMING CHAT requests (``attach_tool_stream`` at submission). Every
other request keeps ``tool_stream = None`` and the applier pays a single
``is None`` check per touched request per batch. Even when present, the
extractor fast-path is a substring check over the new suffix (no regex,
no allocation) until a real ``<tool_call>`` marker appears.

Per-request staging buffer for applier-driven tool extraction.

One instance per streaming chat request, created at admission by
:func:`attach_tool_stream` and read by the SSE generator. Holds the
incremental extractor plus a FIFO of staged SSE delta events.

Staged events are small dicts tagged by kind so the SSE generator
can frame them without re-deriving anything:

  ``{"kind": "reasoning", "text": "<delta>"}``
  ``{"kind": "content", "text": "<delta>"}``
  ``{"kind": "tool", "delta": {<openai tool_call delta>}}``

Producer = the OutputApplier (:meth:`feed`). Consumer = the SSE
generator (:meth:`drain`). Both run on the consumer loop.

True iff this request should run streaming tool extraction.

Any CHAT request (``is_chat``) — a
:class:`StreamingToolCallExtractor` runs for *every* chat stream,
including ones that offered no tools or set ``tool_choice="none"``,
so a model that emits a ``<tool_call>`` marker anyway is framed
identically. Streams that never produce a marker (the overwhelming
majority, tools or not) ride the extractor's zero-regex fast path,
so this broad gate costs only a per-feed substring check — the
"fast-path only" budget for non-tool streams.

Attach a :class:`ChatToolStream` to ``client`` when tool parsing is
in play. No-op otherwise (the request keeps ``tool_stream = None`` and
the applier pays a single attribute check per touched request).

``parser`` is the engine's model-tied tool-call parser
(``eng.tool_call_parser``); the request's own ``tools`` give the parser
its schema for value typing. ``reasoning_parser`` is the engine's
model-tied reasoning parser (``eng.reasoning_parser``), and
``opens_in_reasoning`` is the route's reading of the RENDERED prompt —
whether the model's own template left generation inside a think block.
Called at submission, before the request is published to the engine.
Idempotent.

Run the incremental extractor over ``output_text`` and stage any
resulting content / tool-call deltas. Applier side.

Idempotent for an unchanged ``output_text`` (skips the call when
the length is identical to the last feed), so a batch that only
touched the request's audio lane or finish costs one comparison.

Release any reasoning tail held back as a possible partial marker.

Called once when the engine signals finish. A block truncated at
``max_tokens`` can end mid-``</think>``-prefix; those chars are real
trace text and must be staged rather than dropped.

Pop and return all staged events in FIFO order. SSE side.

Returns an empty list when nothing is staged. Feed and drain run
on the same consumer loop, so the FIFO hand-off needs no lock.

ONE boot-time VRAM ledger whose terms sum to the card.

The bar this module exists to clear: *someone reading one boot log can answer

Every previous surface failed that bar in the same way — it reported a SUBSET
and left the operator to reconcile it against ``nvidia-smi`` by hand. The pool
gauges omit the driver-resident bytes; the freeze reconcile prints a signed
"gap"; the KV ledger prints a residual it explicitly cannot name. Three views,
three different totals, and the same physical bytes carrying different names in
each — which is how one row's bytes came to be called ``scratch.activation`` in
the pool gauges, "activation reserve" in the boot log and "overhang" in the
freeze triage. That row is now ``scratch.forward_arena`` everywhere, and the two
former synonyms are what they always were: the PREDICTION of it, and its
reserved-but-free DETAIL column.

So the ledger here is built on one rule: **rows are disjoint and they sum to the
card.** Concretely, the card decomposes with no overlap and no gap as

    NVML total
      = driver-reserved (never handed to CUDA at all)
      + CUDA total, and CUDA total
          = Σ named-pool mapped            (one row per cuMem tag)
          + growable KV mapped             (state.attn_kv.mapped)
          + torch segments outside pools   (unpooled.torch_default_pool)
          + other processes on the card     (driver.foreign_process)
          + driver-resident, no allocator  (driver.cuda_context, .local_memory,
                                            .modules_loaded, .cudagraph_exec)
          + driver free                    (transient.serving_step + unclaimed.no_owner)

``driver.residual`` is the REMAINDER of that identity — computed last, from the
card down, never estimated and never folded into a neighbour. It is the only
line permitted to absorb error: if a pool under-reports, or a configuration
legitimately loads more cubins than its held ``driver.modules_loaded``
baseline, the bytes land there and nowhere else.

It is NOT a leak signal, and it is NOT the same quantity as
``unpooled.unregistered_pool``. That name belongs to a different surface — the torch
caching-allocator row for a private ``MemPool`` that is neither registered nor
released (:func:`arbi_serve.engine.memory_budget.pool_residency
.residency_by_pool`, ``GET /v1/admin/memory/live``). That row names an owner
nobody declared and IS a leak signal. This one is the difference between two
measurements of the whole card, and legitimate driver-resident growth reaches
with ``driver.residual`` at ~0 is named residency held to a stale expectation,
not an unregistered allocation.

Every row also carries its PROVENANCE
(:data:`~arbi_serve.runtime.pool_taxonomy.PROVENANCE`) — whether the number was
observed on this boot, read back from a prior boot at this configuration key,
computed in closed form, seeded from a constant, or measured with no plan line
predicting it. A ``!`` marks the two that are not measurements of this boot's
own bytes, because a table that renders a guess and an observation identically
is a table that invites the guess to be trusted.

Rows are grouped by the ``category.role`` prefix, which already encodes lifetime
(:data:`~arbi_serve.runtime.pool_taxonomy.CATEGORY_LIFETIME`), so the primary
reading axis is "does this scale with my concurrency knobs, or is it fixed?".

Pure arithmetic + formatting: no CUDA, no engine state, fully CPU-testable. The
caller (:mod:`arbi_serve.engine.phase2_freeze`) does the measuring.

Assemble the ledger. Pure arithmetic — no CUDA, no engine.

Args:
    nvml_total_bytes: NVML framebuffer total (the number on the box). Pass
        ``cuda_total_bytes`` when NVML is unavailable; the driver-reserved
        row then reads 0 rather than a guess.
    cuda_total_bytes: ``cudaMemGetInfo`` total.
    driver_free_bytes: ``cudaMemGetInfo`` free.
    serving_floor_bytes: the runtime free-VRAM reserve one serving step
        which is the operator gate.
    serving_step_terms: that reserve ITEMISED, from
        :attr:`~arbi_serve.engine.boot_state.EngineBootState.serving_step_terms`.
        Empty yields one un-itemised ``transient.serving_step`` row rather
        than a fabricated breakdown.
    pool_mapped_bytes: ``{taxonomy name: cuMem-mapped bytes}`` per named
        pool. Zero-byte pools are dropped — a row of 0 is noise.
    pool_free_bytes: ``{name: reserved-but-free bytes}`` inside each pool,
        reported as a separate device-wide "torch caching overhang" belongs:
        those bytes are already inside the pool rows, so adding them as a
        row of their own double-books them (and did — see
        :mod:`arbi_serve.engine.phase2_freeze`).
    growable_kv_bytes: mapped bytes of the growable KV region →
        ``state.attn_kv.mapped``.
    external_pool_bytes: ``{taxonomy name: bytes}`` for physical a pool
        cuMem-maps in its OWN VA reservation rather than through the torch
        caching allocator (the state pools' sentinel-alias arenas). These
        get their own rows and are deliberately kept OUT of the
        ``unpooled.torch_default_pool`` subtraction below: they are not
        inside ``torch_reserved_bytes``, so subtracting them there would
        understate the default-pool residual by exactly their size.
    torch_reserved_bytes: ``torch.cuda.memory_reserved``. What it holds
        beyond ``Σ pool_mapped_bytes`` is the default-pool residual →
        ``unpooled.torch_default_pool``.
    foreign_process_bytes: device-wide used MINUS this process's NVML
        per-process used → ``driver.foreign_process``. 0 both when we are
        alone on the card and when NVML cannot report per-process usage;
        the row is dropped either way.
    driver_rows: the measured ``driver.*`` attribution rows from
        :meth:`~arbi_serve.engine.memory_budget.driver_residency.DriverAttribution.as_rows`.
        ``driver.residual`` in this list is IGNORED and recomputed here
        from the card down, so the ledger's residual is the residual of the
        WHOLE identity, not of the driver sub-identity alone.
    row_details: extra ``{name: detail}`` overrides/additions.
    notes: lines rendered under the table.

Split driver-free VRAM into named rows summing EXACTLY to ``free_bytes``.

One rule, shared by the boot ledger and the live card so the two cannot
describe the same free segment differently.

The itemised terms are emitted at their FULL declared reserve and any
difference goes on one signed row (:data:`STEP_IN_USE_ROW`) rather than
being clipped out of the terms: nothing measures WHICH term a live step took
its bytes from, and clipping would put a number on that guess. Those bytes
are resident on a pool row already, which is what makes subtracting them
here exact rather than a second booking.

Returns ``[(name, bytes, note, provenance)]``.

Pre-load TKV CUDA kernel modules during cold boot.

Each ``_get_module()`` call inside ``tkv.kernels.*`` does a
``torch.utils.cpp_extension.load_inline`` cache lookup, ``.so`` dlopen,
and Python-level module init even when the ``.so`` is already cached on
disk. On a fresh process this work lands on the first-request path.

This helper fires those ``_get_module()`` calls on a thread pool during
``Engine.build`` so they overlap with the weight-DMA + KV-pool phase.
By the time the first request lands the modules are already loaded.

Note this only pre-loads — it does NOT launch synthetic kernels (which
would also pay the per-kernel CUDA module init). Doing that requires
real input tensors of the right shape; the engine doesn't have them at
this point in build (no real request, no scheduled batch, no metadata).

The MTP verify (``block_m``-keyed) split-K module is not warmed here: it is
compiled pre-capture by :mod:`arbi_serve.engine.tkv_autotune_warm`, which
derives the reachable verify widths from the resolved config (the same set
the verify cudagraph capture enumerates), keeps the ones the resolved route
still sends to split-K, and passes those to turbo-attn's shared
``warm_autotune_layers`` / ``compile_decode_and_cs`` primitive. The
single-token (``block_m == 1``) decode family and the compress-store kernel
warmed here are unconditional — they are on the served path whichever kernel
owns verify.

Block until no background kernel warmup is in flight.

Returns ``True`` once idle, ``False`` if ``timeout_s`` elapsed first.
Call this BEFORE any GPU-state teardown (``empty_cache`` / cuMem
``free_all``) that may overlap a still-running
:func:`warm_tkv_kernel_modules` — freeing memory mid-``load_inline``
deadlocks the CUDA context.

Mirror :func:`tkv.runtime.attention._env_fuse_q_rot_default`.

Critical: ``turbo_attn_simt._get_unified_splitk_module`` builds a different
cache_name suffix ('_fq1' vs '') based on this flag. If we pre-load
with the wrong value, the cache_name doesn't match what precompile.py
produced — cache miss, ~30 s of nvcc compile per variant. Reading
the same env var keeps the warmup aligned with both precompile.py
and the runtime hot path (single source of truth).

Pre-load TKV decode + compress + Turbo-prefill kernel modules.

``layer_shapes`` is a list of unique ``(head_dim, num_q_heads,
num_kv_heads)`` triples — one per attention "family" the model
has. Most models are homogeneous (one entry); Gemma 4 / gpt-oss
have two (sliding vs full).

For each shape we submit, in parallel:

  - ``turbo_attn_simt._get_unified_splitk_module(...)`` for both ``emit_lse``
    variants (single-split fast path + multi-split reduce path).
  - ``cuda_compress_store._get_module(...)``.
  - The Turbo prefill adapter import chain (one-time cost regardless
    of shape — the module is layer-agnostic).
  - When ``has_sliding_window``, also the SWA split-K decode for
    each shape.

``vq_mode_k`` / ``vq_mode_v`` name the per-side codec family. A group-2
VQ side compiles a DIFFERENT kernel at the same bit width, under its own
cache name, so omitting them here pre-loads the scalar Lloyd module on a
vq2 boot — a module the serve path never asks for, while the one it does
ask for stays cold. That failure is silent: this pass reports success and
the JIT surfaces on the first decode.

Modules cache themselves in their own per-launcher dict, so a
subsequent runtime call is a hash lookup. The helper returns a
summary dict for logging.

Pre-load every TKV kernel module the built engine can dispatch to.

Runs as a BOOT PHASE (before the engine declares ready): the decode
split-K ``emit_lse`` variants, the compress-store codec kernel, the
SWA range kernels, and the Turbo prefill adapter import all materialize
here, so no first request (or first prefill) pays a ``load_inline``
module init / ninja build during serving. The pre-capture autotune
warm (``tkv_autotune_warm``) can legitimately skip its per-variant
compile on an autotune-TABLE cache hit — this pass is what guarantees
the module set is loaded regardless.

Shapes come from the constructed ``attn_ops``: those carry the
PER-RANK head counts and the per-layer smart-mix bits the kernels are
keyed on. The model's ``layer_specs`` carry GLOBAL (un-sharded) head
counts, which at TP>1 name a compress-store variant the runtime never
requests. No-op for non-TKV backends (tkv-bypass exposes no bit
widths and has no JIT compile step here).

Failure is logged loud and non-fatal: an unwarmed module costs the
first matching call its JIT latency, which the detector reports.

Return True iff at least one core has its MTP wrapper installed.

``cores`` is an iterable of ``TKVCore``-shaped objects (we only
read ``_mtp_attend``). Returns False on an empty iterable so the
caller can decline when no core has been built yet (cores are
lazy-constructed on first forward unless eagerly built at boot).

Name-attributed per-component model-weight VRAM accounting.

The boot path reports a SINGLE ``arbi_serve_gpu_weights_bytes`` total
(``Σ p.numel()*p.element_size()`` over every CUDA-resident parameter +
buffer). That tells you *how much* weight memory is resident but not
*where it goes*. This module splits that same total into named
components so the dashboard can show a stacked breakdown:

  ``text_layers``    — per-layer decoder blocks (``model.layers.N.*`` /
                       ``...decoder.layers.N.*``). The bulk of weights.
  ``lm_head``        — the output projection (untied only; a tied head
                       shares the embed storage and is attributed to
                       ``embed`` exactly once via data_ptr de-dup).
  ``embed``          — input token embeddings (``embed_tokens`` / any
                       ``*embed*`` tensor not inside a decoder layer).
  ``vision``         — vision tower (``vision_tower.*`` / ``visual.*`` /
                       any ``*vision*`` / ``*visual*`` path).
  ``mtp``            — multi-token-prediction / draft head
                       (``mtp.*`` / ``*mtp*`` / ``*draft*``).
  ``norm``           — final norm + any layernorm NOT inside a decoder
                       layer (``model.norm`` etc.).
  ``other``          — anything unclassified (catch-all so the sum is
                       always exact).

Design contract
---------------
The attribution is computed by walking the FULLY-LOADED model's
parameters and buffers — post-quant-swap (AWQ/EXL3 Marlin tensors are
final), post-compaction (flat-slab views). It therefore counts the
ACTUAL resident bytes, including INT4/AWQ packed qweights + per-group
scales + zero-points, exactly as they sit in VRAM. Because it is the
SAME walk the total gauge uses (``p.numel()*p.element_size()`` over
CUDA params + buffers, de-duped by ``data_ptr``), the components sum to
the total by construction — any allocator slack lives OUTSIDE this walk
(it is reserved-but-unallocated pool bytes, not tensor bytes).

Tied weights are attributed once: the de-dup keys on ``data_ptr()`` so
a tied lm_head/embed storage is counted under whichever name the walk
reaches first. The classifier is ordered so that an ``embed_tokens``
path wins over a generic ``lm_head`` alias.

The name classifier (:func:`classify_component`) is a pure string
function with no torch dependency, so it is unit-testable against a
synthetic name→bytes map without a GPU.

Bucket a fully-qualified tensor name into a weight component.

``qualified_name`` is the ``module_path.attr`` string from
:meth:`torch.nn.Module.named_parameters` /
:meth:`~torch.nn.Module.named_buffers` (e.g.
``model.layers.7.self_attn.q_proj.weight``,
``model.embed_tokens.weight``, ``lm_head.weight``,
``mtp_head.layers.0.mlp.gate_proj.weight``).

The classification is robust to architecture naming differences
(qwen / gemma / llama / deepseek / nemotron / …): it keys on
substrings rather than an exact arch-specific layout. Order
matters — the FIRST matching rule wins:

  1. vision tower (``vision`` / ``visual``) — checked first so a
     ``...vision_tower.encoder.layers.N...`` path is attributed to
     ``vision`` and not ``text_layers``.
  2. mtp / draft head (``mtp`` / ``draft`` / ``nextn``) — likewise
     checked before the generic per-layer rule so an MTP block's
     own ``.layers.0.`` doesn't read as a text layer.
  3. quant repack scratch (``marlin`` / ``quant_scratch`` /
     ``repack`` / ``workspace``) — resident Marlin workspace etc.
  4. per-layer decoder block (``.layers.<N>.`` / ``.blocks.<N>.`` /
     ``.h.<N>.``) → ``text_layers``.
  5. lm_head / output projection → ``lm_head``.
  6. token embedding (``embed``) → ``embed``.
  7. a remaining norm/layernorm → ``norm``.
  8. everything else → ``other``.

Returns one of :data:`COMPONENTS`.

True if ``lower_name`` contains an indexed decoder-block segment.

Matches ``layers.<N>.`` / ``blocks.<N>.`` / ``h.<N>.`` where ``<N>``
is one or more digits. The block keyword may appear at the START of
the name (``layers.0.weight`` — the bare ``nn.ModuleList`` path when
the model object itself is the inner module) or mid-name
(``model.layers.0.weight``). Pure string scan keeps this torch-free
and cheap; the trailing ``.`` requirement avoids matching a bare
``layers`` attribute that isn't an indexed list.

Bucket a ``{qualified_name: bytes}`` map into component totals.

Pure / torch-free — the unit-test seam. Every component in
:data:`COMPONENTS` is present in the result (zero when no tensor
mapped to it) so the emitted label set is stable across models.

The returned dict's values sum EXACTLY to ``sum(named_bytes.values())``
because :func:`classify_component` is total (always returns a
bucket) and the catch-all ``other`` absorbs anything unmatched.

Total CUDA-resident parameter + buffer bytes for a loaded model.

The ``arbi_serve_gpu_weights_bytes`` gauge's value, off the same cached
walk the per-component gauge uses — so the two agree by construction and
the tree is walked once per model rather than twice per scrape.

Walk a loaded model and return ``{qualified_name: cuda_bytes}``.

Counts every CUDA-resident parameter + buffer exactly once (de-dup
by ``data_ptr`` so a tied lm_head/embed storage is not
double-counted). Mirrors the byte accounting in
:meth:`arbi_serve.server.metrics.Metrics._observe_weights_bytes` so
the per-component buckets reconcile with the total gauge.

The qualified name is taken from the FIRST path that reaches a
shared storage (``remove_duplicate=False`` walk + manual data_ptr
de-dup). The classifier is ordered so an ``embed_tokens`` path is
reached and wins for tied weights.

Per-component CUDA weight bytes for a fully-loaded model.

Convenience composition of :func:`collect_model_named_bytes` +
:func:`attribute_named_bytes`. Returns a dict keyed by every name
in :data:`COMPONENTS`; values sum to the same total the
``arbi_serve_gpu_weights_bytes`` gauge reports.

Computed once per module tree (see :data:`_ATTRIBUTION_CACHE`) — the
answer cannot change while the model is loaded, and an observer must not
pay for re-deriving it. A copy is returned so a caller mutating its
result cannot corrupt the next reader's.

Env-truth: the machine's own account of what it was doing while we timed it.

Sibling of :mod:`arbi_serve.flag_truth`. Flag-truth answers *which code path
ran*; env-truth answers *what machine it ran on* — CPU run-queue depth, GPU
clocks, temperature, throttle reasons, and PCIe link state. A benchmark
number without the machine state that produced it is not a measurement, it
is an anecdote, so every result row carries its env-truth and a run on a
degraded host fails loud instead of quietly emitting a number.

Freeness contract: nothing in this module runs in the engine hot path and
nothing in it imports torch. The sampler is an out-of-process sidecar
(``python -m arbi_serve.env_truth``) that talks to a long-lived
``nvidia-smi --loop-ms`` pipe and reads ``/proc``, so it cannot perturb the
engine loop because it is not in it.

PCIe link state is only meaningful under load: an idle GPU down-trains its
PCIe link to save power and re-trains it when work arrives, so a boot-time or
idle PCIe sample proves nothing about link health. This is why
:class:`EnvWindow` evaluates the link only over samples whose GPU was
actually doing work (``util_gpu_pct >= _BUSY_UTIL_PCT``), and why a boot-time
check deliberately does not police the link gen at all.

Link *width*, by contrast, is negotiated at train time and does not fluctuate
with power state; a x8-of-x16 reading is normally slot bifurcation (two cards in
a board that splits one x16 into x8/x8), not a fault. We record it and compare it
to a per-host expectation rather than to the GPU's theoretical max.

The host's contention truth at one instant.

``procs_running`` is the run-queue depth straight from ``/proc/stat``. It is
the single most direct measure of "is something else eating the cores this
engine needs" -- more direct than loadavg, which is a decayed average and
lags a burst by tens of seconds. A bench that starts while the run queue
already exceeds the core count is measuring the neighbours, not the code.

Stream ticks from one long-lived nvidia-smi.

A fork-per-tick sampler would cost a process spawn every tick and, worse,
would take the driver lock repeatedly. ``--loop-ms`` keeps a single process
alive and streams CSV lines, so the marginal cost of a tick is one line read.

The env-truth of one timed window, reduced to what fits on a result row.

This is the object that makes a slow invocation self-diagnosing. Stamp it
on every row: when two rows disagree and their code is identical, the
field that differs here is the answer.

Thresholds for fitness-to-benchmark.

``runqueue_excess`` defaults to 0.5: the run queue may not exceed half the
core count before we boot. On a quiesced box this sits near 0. A bench
started while another process is saturating the run queue is timing the
neighbours, not the code.

Decide whether this host may be benchmarked on right now.

Deliberately does not police PCIe link gen: at preflight the GPUs are idle,
an idle link is down-trained by design, and policing it here would fire on
every healthy box. The link is judged over the run window instead, using only
the samples where the GPU was actually busy. See the module docstring.

Judge a completed timed window. This is what makes a row self-diagnosing.

Unlike :func:`preflight`, this may judge the PCIe link, because it looks only
at the samples taken while the GPU was busy.

Scoring a served endpoint against the published evals.

This package owns the whole protocol -- the catalogue (``catalogue.yaml``),
the harness tasks this project adds on top of lm-eval's built-ins
(``tasks/``), and the adapter that points the harness at any OpenAI-compatible
endpoint (``endpoint_lm``). It is self-contained by design: an install of
this package is everything a run needs, with no sibling checkout to find.

Run one from the command line::

    python -m arbi_serve.evals --model arbi-openai         --model_args model=<id>,api_base=http://localhost:8000/v1         --tasks aime26_chat

or from the console's Evals tab, which reads the same catalogue.

lm-eval itself is an EXTRA (``pip install 'arbi-serve[evals]'``) -- it pulls
torch and the dataset stack, which the serving path does not need. Everything
in this module that does not run a benchmark (the catalogue, its shape) is
importable without it, so the console can render the catalogue and say what is
missing rather than failing to import.

1-based line of each entry under ``evals:``, in document order.

``yaml.safe_load`` discards position, and the whole point of citing a
source line is that the reader can go look at it, so the document is
composed a second time for its node marks rather than the lines being
guessed from the parsed values.

``python -m arbi_serve.evals <lm-eval args>`` — lm-eval, with this
project's endpoint adapter and tasks already on the table.

Everything after the module name is lm-eval's own CLI, unchanged; the only
thing this entry point adds is importing :mod:`arbi_serve.evals.endpoint_lm`
first, so its ``@register_model("arbi-openai")`` has run by the time the
harness resolves ``--model``, and turning a missing harness into the install
line rather than a traceback.

lm-evaluation-harness adapter for any OpenAI-compatible endpoint.

Registers the harness model ``arbi-openai``, which speaks plain
``/v1/chat/completions`` and ``/v1/completions``. That is the whole reason a
run needs no bespoke harness per target: this server, a vLLM on another box
and a gateway behind mTLS are all the same ``api_base``.

Usage::

    python -m arbi_serve.evals         --model arbi-openai         --model_args model=<id>,api_base=http://localhost:8000/v1         --tasks aime26_chat

    # Behind a client-certificate gateway, pass the paths explicitly:
    --model_args model=<id>,api_base=https://gw/v1,client_cert=/p/client.pem,client_key=/p/client-key.pem,ca_cert=/p/ca-chain.pem

Supported lm-eval request types:
    - generate_until        -> POST /v1/chat/completions (stop=until)
    - loglikelihood         -> POST /v1/completions (echo=true, logprobs=1)
    - loglikelihood_rolling -> NotImplementedError (needs a local tokenizer)

The client-certificate context, when one was asked for.

No cert path is ever defaulted or discovered: a run either presents the
certificate the caller named or goes over ordinary TLS. Guessing a path
would make the same command mean different things on two machines.

POST /chat/completions with SSE streaming; return concatenated content.

Streaming keeps the TCP connection alive through long generations,
so a reasoning task whose single answer takes minutes does not trip
an intermediate gateway's idle timeout.

Chat-API-friendly process_results for lm-eval AIME task variant.

The upstream `aime26` task is written for raw text-completion APIs. Its
answer extractor (`\boxed{...}` / `$...$` precedence) misses the common
chat-API format "The answer is (B)" / "**Answer:** 391". This module
provides a lenient process_results used by the `aime26_chat` variant.

The flag-truth contract: every runtime flag declares whether it fires.

Part two of the flag-truth invariant (part one:
:mod:`arbi_serve.flag_truth` — the fire counters; part three: the gates in
``tests/test_flag_contracts.py`` (CPU) and
``tests/test_flag_truth_live_gpu.py`` (GPU)).

For the canonical served config (:data:`CANONICAL_CONFIG`), every field of
:class:`~arbi_serve.runtime_flags.RuntimeFlags` must be declared exactly one
of:

``MustFire(counter=...)``
    The flag's path is claimed live on the canonical config. Its
    :class:`~arbi_serve.flag_truth.PathCounter` must be ``fired > 0`` after
    the canonical smoke — a zero is a hard failure naming the flag, its
    declared state, its actual fire count, and any recorded refusal
    reasons. ``counter`` names the PathCounter registered at the path's
    execution site (several flags may share a counter when one executed
    path proves them all — e.g. a boot-time configuration flag whose
    effect is a property of the built engine).

``Inapplicable(reason)``
    The flag is explicitly declared not to apply on the canonical config,
    with the reason recorded here in code — reviewable, greppable, and
    wrong in public if it lies. This covers default-OFF opt-in levers,
    debug/diagnostic knobs, other-hardware / other-modality paths, and
    pure value knobs that parametrize a path owned by another flag's
    counter.

There is no third option. A new flag that declares neither fails the CPU
suite (``test_flag_contracts.py``), so a flag can never sit in an
undeclared state where "default ON" and "actually executes in production"
silently diverge — the class this system exists to kill: a path that is
structurally ineligible under production sampling, excluded by a blanket
gate elsewhere, pinned OFF by a deployment file, or chained behind another
flag that is OFF, while the registry still reports it ``True``.

Deployment pins (:data:`DEPLOY_PIN_CONTRACTS`): a compose file that pins an
``ARBI_*`` / ``TKV_*`` env var away from the code default deploys a
different engine than the one the code's defaults describe. Every such
divergence must be declared here with a reason, or the CPU gate fails
naming the file, the pin, and the code default.

The flag's path must actually execute on the canonical config.

``counter`` is the :func:`arbi_serve.flag_truth.path_counter` name the
path bumps at its execution site.

``phase`` is the scrape window and it is load-bearing — the gate reads a
different snapshot for each value, so a wrong ``phase`` is not a cosmetic
label on a failure message, it is a guaranteed-red cell that no amount of
re-running can turn green:

``"boot"``
    A boot/configuration path that runs once while the engine builds.
    Read off the pre-reset boot snapshot.
``"graph_record"``
    The increment is inside a cudagraph capture region: it fires at
    graph-record time and cannot fire at replay. Read off the pre-reset
    boot snapshot — a serving-window read of one of these is zero by
    construction.
``"serve"``
    A host-side path re-entered under traffic, outside any capture
    region. Read off the post-smoke snapshot.

``phase`` must agree with the window the counter declares at its own
registration site (:class:`~arbi_serve.flag_truth.FiringWindow`);
``tests/test_flag_firing_window.py`` fails on divergence, so the contract
cannot quietly claim a window the code contradicts.

The ``"serve"`` default is deliberate but narrow: it is the safe value
only because a mis-defaulted serve counter fails loud (a red cell someone
must explain), whereas the reverse would pass silently. It is not a
licence to skip the classification — declare the window on the counter.

The flag is declared not to apply on the canonical config.

``reason`` is mandatory and must say why (not just "off"): the shipped
default that keeps it dormant, the hardware/modality it needs, or the
owning flag whose counter carries its truth. This is the auditable
record — an Inapplicable with a false reason is a one-line reviewable
lie, not a silent state.

The one config the flag-truth contract is evaluated against.

This is the production serve: the model, parallelism, speculation,
KV codec, and sampling preset that production deployments actually
run. Every
``MustFire`` / ``Inapplicable`` declaration below is a claim about
this config — not about some config where the flag could fire.

Evaluate every contract against a counter snapshot.

``counters`` is a :func:`arbi_serve.flag_truth.flag_truth_report`
snapshot (possibly RPC-serialized). ``resolved`` is the live
:func:`arbi_serve.runtime_flags.runtime_flags` value per field (env +
active-member overlay). Returns ``(rows, failures)``:

``rows``
    One dict per flag — field, env var, declared ``default``, the
    ``resolved`` value the process actually booted with, contract kind,
    counter name / fired / refused / refusal reasons (MustFire) or the
    recorded reason (Inapplicable), plus ``suspect`` for the
    review-attention combination *default-enabled yet declared
    INAPPLICABLE* — a flag that ships ON while its own contract records
    that it cannot run on the canonical config (the lying-default shape
    this system exists to expose).

    ``resolved`` is the receipt's honest witness of served state: a
    request-default-policy flag like ``prefix_cache`` (declared
    INAPPLICABLE because an explicit per-request ``cache`` field
    overrides it) has no counter, so its ``default`` alone cannot tell a
    reader whether the box booted with prefix reuse ON or OFF. The
    resolved value can — a bench that boots ``ARBI_PREFIX_CACHE=0`` reads
    back ``resolved="False"`` here, never a phantom sourced from the code
    default. When ``resolved`` is not supplied the code default stands in
    (the best available witness), so every row always carries the field.

``failures``
    One formatted line per MUST_FIRE flag whose counter shows zero
    fires — naming the flag, its declared state, the counter, the
    actual count, and any recorded refusal reasons. Empty ⇒ the
    canonical config honors every MUST_FIRE claim.

Audit the process environment for unknown / renamed flag env vars.

Returns ``(errors, warnings)``:

* **errors** — vars under a renamed prefix (``TQKV_*``): dead names
  whose value is silently ignored while the engine resolves the
  default. Each message names the live replacement when the renamed
  name exists in a registry. These fail the boot.
* **warnings** — ``ARBI_*`` / ``TKV_*`` vars recognized by neither
  the RuntimeFlags registry, the tkv registry, nor
  :data:`KNOWN_TOOLING_ENVS`: a typo, a stale name, or a raw read
  dodging the registry. Logged prominently at boot.

Boot gate: warn on unknown flag envs, raise on dead/renamed ones.

Called from ``cli.bootstrap.main`` on every rank. A ``TQKV_*`` var in
the environment means the rig believes it pinned a codec knob while
the engine takes the default — the run would report numbers under a
config it is not actually serving, so the boot refuses instead.

Effect probes: prove a flag's counter witnesses a difference, not a branch.

A ``path_counter`` fires when its branch executes. Branch execution is not
evidence the flag did anything: a branch can run and return exactly what the
fall-through would have returned. On sm_89 ``gdn_conv_force_kernel``'s branch
executes and returns ``True`` — and the arch fall-through also returns
``True``. A counter there fires green while the flag changed nothing, and it
does so on the canonical lane, so it *looks* machine-verified. That is worse
than prose: the lie is now wearing a passing gate.

An :class:`EffectProbe` closes that gap. It drives the flag's real decision on
a config where the flag is *load-bearing* and asserts the observable diverges
when the flag flips. A flag whose decision is invariant to its own value on a
config has no honest counter placement there — the counter belongs at the site
where flipping the flag flips the output, and on the config where that site is
reached.

This is the null-control rule (a perf A/B needs identical arms differing only
in the lever) applied to telemetry: the counter earns its ``MustFire`` by a
measured differential, not by branch coverage. It is the dynamic complement to
``tools/flag_deadness.py``, whose ``COUNTER_LOCAL`` proves only that a counter
fires under *some* guard the flag controls — never that the guarded branch
changes the outcome, and never that the counter covers *every* site the flag
gates.

Import discipline: nothing heavy at module scope (no torch, no engine) so a
probe can be declared beside any decision without a cycle. A probe's
``observe`` closure may import what it needs when it runs.

A differential that proves one flag is load-bearing on one config.

``observe(flag_value)`` runs the flag's real decision — the same code the
counter guards — on the config named by ``on``, and returns the observable
the flag is meant to move: a returned decision, the kernel selected, and
(once a counter is wired) the pair ``(decision, counter.fired)``. The gate
asserts ``observe(True) != observe(False)``.

``on`` names the load-bearing config — one where the flag's value changes
the observable — not merely one where its branch runs. Naming a config
where the branch executes but the outcome is invariant is the failure this
type exists to reject: it relocates the lie from prose into an ``on=``
clause and dresses it in a green gate.

``on`` is a config surrogate, not a live boot: a pure host-side decision
(device compute-capability, a config field) is reachable on CPU by fixing
that surrogate, so most probes run without a GPU. A probe whose observable
is GPU-resident runs its differential on the ``on`` boot lane instead —
the same flag-on/flag-off null control that lane already owes.

True iff flipping the flag flips the observable on ``probe.on``.

False means the flag is inert on that config: no counter placed on the
branch it gates can honestly witness the flag there, because the same
observable is produced with the flag off.

BOOT_LANES: the configs that make ``on=<config>`` a booted fact.

The flag-truth contract (:mod:`arbi_serve.flag_contracts`) declares, per flag,
the config on which its counter must fire. That declaration is worth nothing
unless the named config is actually booted and its counters read — an
``on=<a config nobody boots>`` is ``Inapplicable`` wearing a costume. This
module is the anti-costume half: every config a contract can name is a
:class:`Lane` here, and the gate (:mod:`tests.test_flag_truth_live_gpu` /
:mod:`tools.flag_truth_gate`) boots each one and asserts the declared counters
fired.

Three config kinds, because today's live specimens need all three:

``BOOT`` — a server boot with an env/CLI overlay. Reads counters off the live
    ``/v1/admin/flag_truth`` snapshot. The canonical served config plus the
    off-canonical boots (diagnostics-on, audio, nvfp4) that carry the flags the
    canonical config cannot fire.
``FAULT`` — a boot with a fault injected (a malformed recipe on disk, a stale
    backend pin). The ~14 ``allow_*`` flags read only inside ``except`` /
    error branches; setting the flag ``True`` is not sufficient — the fault
    must be present. A clean boot lane can never fire them; claiming one would
    be a costume of a different colour.
``SCRIPT`` — a CLI diagnostic invocation, not a server boot.

Coverage is derived from the contract table, never hardcoded, so a flag
flipped ``Inapplicable -> MustFire`` by another slice is picked up
automatically: :func:`expected_counters_for_lane` reads ``MustFire.on`` when
that field exists, and otherwise falls back to "every ``MustFire`` counter is
owed by the canonical lane" — today's single-lane contract.

CPU-safe: imports nothing heavy, and reads no environment. The lane data is
inert — a lane stores a fixture name, never a resolved path. Path resolution
(which reads the harness env vars ``ARBI_LANE_LOCAL_MODELS_ROOT`` /
``ARBI_TEST_MODELS_ROOT`` / ``TKV_MODELS_ROOT``) lives in the gate
(``tools/flag_truth_gate.resolve_model_path``), outside ``arbi_serve/``, so the
single-source rule (only ``runtime_flags.py`` reads ``ARBI_*``/``TKV_*`` inside
the package — ``test_runtime_flags_single_source.py``) holds. These are harness
env vars, not runtime flags, so they belong beside ``tests/_fixture_paths.py``,
not on ``RuntimeFlags``.

One bootable (or runnable) config the flag-truth contract can name.

``name`` is the stable identifier a contract's ``on=`` refers to. ``model``
is a fixture directory name resolved under the models root at boot time
(never an absolute path baked here — the root differs per box).

``cli`` / ``env`` are the overlay that distinguishes this lane from a bare
boot. ``fault`` (``kind == LaneKind.FAULT``) and ``script_argv``
(``kind == LaneKind.SCRIPT``) carry the extra machinery those kinds need.
``canonical`` marks the one lane the
production serve actually runs — every ``MustFire`` is owed here until a
contract explicitly routes its counter to another lane via ``on=``.

``needs`` documents why this lane exists off-canonical (a modality, a
checkpoint, a fault) — human-readable, not load-bearing.

Counters this lane is contractually obligated to fire.

Derived from the contract table, never hardcoded — a flag flipped to
``MustFire`` by any slice is picked up here on the next call.

When ``MustFire`` carries an ``on`` field (a frozenset of lane names), a
counter is owed by lane L iff ``L.name in contract.on``; otherwise every
``MustFire`` counter is owed by the single canonical lane — today's
one-config contract.

MustFire counters no lane claims to fire — the static anti-costume check.

A non-empty result is a CPU-gate failure: a flag declared ``MustFire`` whose
counter is owed by no ``BOOT_LANES`` entry is ``on=<a config nobody boots>``
— the exact costume this module exists to forbid. Empty today, since every
``MustFire`` is owed by the canonical lane, which exists.

Can flipping this param change anything on THIS boot?

An A/B harness that flips every overridable param measures noise on the
ones whose path this configuration cannot reach, and noise on a
speculative arm is large enough to be reported as a win. The verdict is
three-way and the third value is load-bearing:

``REACHABLE``
    A live flip can change behaviour on this configuration. Benchmark it.

``INERT``
    A live flip cannot change behaviour on this configuration, for a
    reason that is derived from the configuration itself and carried in
    the verdict. Refuse the arm; a measured delta on it is measurement
    noise wearing the flag's name.

``UNKNOWN``
    A fact the verdict needs is not readable from the inputs supplied.
    Benchmark it, and label the result provisional. UNKNOWN never
    collapses into REACHABLE: "not shown to be dead" is not "shown to be
    live".

Two registries already in the tree carry most of the evidence, and this
module reads them rather than restating them:

:mod:`arbi_serve.config_overrides`
    ``PARAMS[name].scope`` says how a delta on the param is applied — a
    ``runtime`` scope installs a flag overlay on the ACTIVE captured
    member with no rebuild, every other scope builds a new member.

:mod:`arbi_serve.flag_contracts`
    ``FLAG_CONTRACTS[flag]`` names the ``PathCounter`` the flag's path
    bumps, and ``MustFire.phase`` says WHEN it bumps it: ``boot`` and
    ``graph_record`` fire while the engine builds, ``serve`` fires under
    traffic.

Crossing those two gives the sharpest generic rule here, and one no
single registry states on its own: a ``runtime``-scope param whose
contract fires in a build-time window is INERT under a live flip,
because the overlay install does not rebuild the model, recompile it, or
re-record a graph. The path it gates has already run.

``Inapplicable`` is NOT that oracle
-----------------------------------
:class:`~arbi_serve.flag_contracts.Inapplicable` marks a flag as not
applying to :data:`~arbi_serve.flag_contracts.CANONICAL_CONFIG` — one
specific model, parallelism and drafter. Its ``reason`` is prose written
against that config, so reading the marker as a reachability verdict for
an arbitrary boot inverts several entries: ``exl3_prefill_row_invariant``
is Inapplicable only because the canonical model carries no EXL3 linear,
and is perfectly reachable on a boot that does. The reasons are parsed by
a human once and encoded here as conditions over :class:`LiveConfig`;
the marker itself is never a verdict.

Reading the counters back
-------------------------
``GET /v1/admin/flag_truth`` serves the live
:func:`~arbi_serve.flag_truth.flag_truth_report` snapshot — every
registered counter's ``fired`` / ``refused`` / ``refusals``. Passing it
in as :attr:`LiveConfig.counters` turns a config-derived verdict into a
measured one: a flag that is engaged, whose declared counter shows zero
fires and a recorded refusal, did nothing, and the refusal text says why.
``POST /v1/admin/flag_truth/reset`` brackets one arm.

What a checkpoint's ``config.json`` says about its capabilities.

``extra`` carries the architecture-specific sub-config a capability
bit lives in — a DFlash drafter's ``dflash_config``, whose KEYS are
the capability set (a checkpoint with no ``markov_rank`` has no
markov head, so every markov-gated path is dead on it).

Which drafter this boot installed, and what it can do.

``kind`` mirrors ``engine.mtp_attach.build_mtp_driver``'s dispatch
exactly, in its order: a named external source wins over the
checkpoint's bundled head, so a boot that names one never builds the
bundled head at all and every bundled-head path is dead on it.

Everything the predicates read, and nothing they do not.

Every field that a predicate can fail to resolve is Optional, and a
predicate that needs a ``None`` field returns UNKNOWN rather than
guessing. That is the whole reason the third verdict exists: a
missing fact must be visible in the output, not absorbed by a
default.

One precondition, resolved against a live config.

``holds`` is tri-state on purpose: ``None`` means the input needed to
decide was not supplied, which is a different answer from "the
condition fails".

The path hangs off the checkpoint's own MTP head.

A boot that names an external source never builds that head
(``mtp_attach.maybe_append_mtp_spec`` returns False), so the seam the
path installs through does not exist.

The path is parametrized by another param that must be ON.

A key the catalogue carries with a ``None`` value is NOT the param
switched off — it is the param unreported, and reading it as off
refuses a flag on a fact the target never supplied.

The path needs another param set to one of ``values`` SPECIFICALLY.

The tri-state sibling of :func:`_needs_param_engaged`. That one asks "is
the owner on", which for a tri-state answers a different question than the
consumer has: ``exl3_int8_verify='1'`` is on and nonetheless leaves
``exl3_int8_verify_min_rows`` inert, because only ``auto`` consults a
threshold. A knob reported REACHABLE on the arm that never reads it is the
same false green this module exists to refuse.

An unreported key stays UNKNOWN for the same reason as there: reading a
fact the target never supplied would refuse a flag on an invention.

The flag has ONE consumer and it runs while the member is built.

The sibling of the generic live-overlay rule, for a flag whose
contract is :class:`~arbi_serve.flag_contracts.Inapplicable` and
therefore names no counter phase to read the window off. ``site`` is
the consuming code path, so the verdict still carries a receipt a
reader can grep for rather than a bare assertion.

A live-overlay param whose path only runs while the engine builds.

``engine.config_variant`` applies a non-capture-affecting flag delta
by installing an overlay on the ACTIVE member: no weight reload, no
recompile, no re-record. So a flag whose contract declares a
``boot`` or ``graph_record`` counter has already had its one chance
to fire by the time an A/B flips it, and both arms serve the value
the member was built with.

The override path itself rejects this delta on this config.

``config_overrides.apply_overrides`` refuses a delta touching
``gdn_prefill_capture`` while the resolved ``prefill_capture`` is
eager — the flag's two consumer sweeps both need a non-eager mode, so
the delta cannot mean what it says. A refusal is the strongest
possible INERT: the server will not even take the arm.

The flag is engaged and its own counter recorded a refusal.

The counter snapshot outranks every prose condition here: a declared
MustFire counter sitting at zero fires with a recorded refusal, while
the flag reads ON, is a direct measurement that the path declined.
Only consulted when the flag is currently engaged — an OFF flag reads
zero for the obvious reason and proves nothing.

Read a checkpoint's ``config.json`` into :class:`CheckpointFacts`.

Best-effort: an unreadable path yields an empty record, whose every
field is ``None``/empty, which drives the predicates that need it to
UNKNOWN rather than to a guess. The text tower's sub-config is
merged over the top level so a multimodal wrapper does not hide the
geometry.

Read cudagraph state off ``GET /v1/admin/capture_hist``, tri-state.

A NON-EMPTY per-batch histogram is proof: it is a replay the server
actually served, which no un-captured boot can produce.

An empty one proves NOTHING, and reading it as ``False`` is how two
perfectly reachable flags get called inert. That histogram is a debug
counter gated behind its own environment switch, so on an ordinary
boot it is empty whether or not the engine captured and is replaying
graphs; a missing ``by_batch`` key says only that the field was not
reported. Both answer UNKNOWN, and a param that needed the fact says
so instead of being refused on it.

Assemble a :class:`LiveConfig` from admin-API responses.

``config_override`` is ``GET /v1/admin/config_override``,
``server_info`` is ``GET /v1/admin/server_info``, ``flag_truth`` is
``GET /v1/admin/flag_truth``. Every one but the first is optional,
and each omission shows up as UNKNOWN on the params that needed it
rather than as a guess.

``cuda_graphs`` has no admin field of its own; the caller supplies
it, or :func:`live_config_from_server` proves it from a non-empty
``GET /v1/admin/capture_hist``.

Fetch the admin surfaces off a running server and read them.

``with_counters`` also pulls ``GET /v1/admin/flag_truth``, which
upgrades the config-derived verdicts with measured refusals.
Cudagraph state is read from ``GET /v1/admin/capture_hist`` by
:func:`cuda_graphs_from_capture_hist`, which answers UNKNOWN unless
the histogram positively proves capture — an empty histogram is a
debug counter that was switched off, not a boot that captured
nothing.

Flag-truth telemetry: prove that a perf path actually executes.

The invariant this module carries: **a flag's declared state is not
evidence that its code path runs on the config we serve.** A default-ON
flag can be structurally ineligible under production sampling, excluded by
a blanket gate elsewhere, overridden OFF by a deployment file, or chained
behind another flag that is OFF — and in every one of those states the
registry still reports it ``True``. The only honest evidence that a path
runs is a counter incremented at the path.

The cure has three parts; this module is part one:

1. **Fire-count telemetry** (here): every perf-relevant path registers a
   :class:`PathCounter` and bumps it when the fast path actually executes
   (``fired``) or when its gate routes away (``refused`` + a reason). The
   hot increment is a plain attribute ``+= 1`` on a ``__slots__`` object —
   no sync, no allocation, no dict lookup at fire time. This generalizes
   the counter pattern of
   :func:`arbi_serve.spec_decode.mtp_verify_sharded.sharded_steps_total`.
2. **A declared contract per flag** (:mod:`arbi_serve.flag_contracts`):
   for the canonical served config every :class:`RuntimeFlags` field is
   declared ``MUST_FIRE`` (counter must be ``> 0`` after a canonical
   smoke) or ``INAPPLICABLE(reason)``. No third, undeclared state exists.
3. **Automated gates**: ``tests/test_flag_contracts.py`` (CPU lane —
   contract exhaustiveness, counter wiring, compose-pin divergence) and
   ``tests/test_flag_truth_live_gpu.py`` (GPU lane, ``live_gpu_tp2`` —
   boots the canonical config and asserts every ``MUST_FIRE`` counter
   fired).

Import discipline: this module imports nothing heavy (no torch, no engine
modules) so any module — kernels, scheduler, boot — can register a counter
at import time without cycles.

When a counter can move. Reading it outside this window reads zero.

A fire count is evidence only if it was read in the window the counter
can fire in. This enum is the second half of the truth condition
:class:`~arbi_serve.flag_contracts.MustFire` already carries in ``on=``
(which config) — ``window`` is when, and a claim missing either half is
not a claim.

A GRAPH_RECORD counter fires only at capture time: its increment is
baked into the captured chain, so the recorded kernels run at replay
while the Python ``fire()`` call does not. A serving-window zero on a
GRAPH_RECORD counter is therefore not evidence of dormancy — it is
evidence of the wrong window — and it is guaranteed by construction, so
it can never be disconfirmed by re-running the smoke.

``GRAPH_RECORD``
    The increment sits inside a region traced under
    ``torch.cuda.graph(...)``. It fires once per capture and cannot fire
    at replay. Scrape at BOOT, pre-reset. Any post-warmup reset zeroes it
    forever.
``BOOT``
    A boot/configuration path that runs once while the engine builds.
    Scrape at BOOT, pre-reset.
``SERVE``
    A host-side path re-entered per step / per request under traffic —
    it runs eagerly outside any capture region, so it moves after a
    post-warmup reset. Scrape after the smoke.
``FAULT``
    Reached only inside an ``except`` / error branch; needs the fault
    injected, not just the flag set.
``SCRIPT``
    A CLI script invocation, never a server boot.
``UNDECLARED``
    No window declared at the registration site. Not a synonym for
    SERVE: it means nobody has classified this counter, so a zero from
    it is uninterpretable. Reported as-is rather than silently assumed
    — an assumed window is exactly the substitution this enum exists to
    forbid.

Fire/refuse counter for one perf path.

Hot-path contract: hold the object (module global / attribute bound at
boot) and call :meth:`fire` — a plain Python int increment on a slotted
attribute. Never look the counter up by name per step.

``fired``
    The fast path executed (not "is configured", not "was eligible").
    Bump it at the point of no return — after every gate has passed
    and the path's work is actually being done.
``refused`` / ``refusals``
    The path was reachable (its gate code ran) but the gate routed
    away, with a machine-readable reason. Refusals are how a
    default-ON flag that never fires becomes diagnosable instead of
    silent: the report shows *why* a path never ran, not just that it
    didn't.
``window``
    When this counter can fire (:class:`FiringWindow`). Carried on the
    counter — not on the contract — because the window is a property of
    where the increment sits, which only this site knows. Every consumer
    (the gate, a bench harness, a human reading the admin endpoint) then
    reads the same declaration instead of guessing from the name.

Return (creating if needed) the process-local counter ``name``.

Call at import time or boot and bind the result; do not call per step.
Idempotent: the same name always returns the same object, so a second
registration site shares the counter rather than silently forking it.

``window`` declares when the increment can fire (:class:`FiringWindow`),
at the only place that knows: the registration site next to the code.
Declaring it is how a reader learns that a zero is uninterpretable
rather than damning.

Two sites declaring different windows for one name raises. That is not
pedantry: one counter shared by a captured path and an eager path has no
single window, so no scrape can judge it — the counters must be split.
A site that declares a window upgrades an earlier ``UNDECLARED``
registration (import order is not a design decision, so it must not
decide the window).

Declared window of ``name``; ``UNDECLARED`` if never registered.

An unregistered name is UNDECLARED, not SERVE: a counter whose module
never imported has no window, and assuming one would manufacture the
exact false verdict this module exists to prevent.

Reset every registered counter (measurement-window bracket).

**This is a one-way door for GRAPH_RECORD counters.** They fire only while
a graph is recorded, which has already happened by the time anything
brackets a serving window. Resetting after warmup zeroes them for the life
of the process: every later scrape reads 0, and that 0 is unliftable — no
traffic can restore it. Scrape the boot window before calling this.

``sterilizes_graph_record`` is an acknowledgement, not a switch: pass
``False`` to reset only the counters that can refill (SERVE / UNDECLARED),
which is what a mid-run measurement bracket actually wants. The default
stays all-counters so existing callers keep their semantics.

Snapshot every registered counter: name -> fired/refused/refusals.

This is the metric surface the admin endpoint serializes. Counters
register lazily at import, so a path whose module never imported is
absent here — the contract gate treats absent as ``fired=0`` (a path
that never even loaded certainly never ran).

The fast path executed ``n`` times (e.g. n tokens committed).

Same cost as :meth:`fire` — one slotted int add — but O(1) instead of
an n-iteration Python loop at the call site. Hot-path safe.

Counts plus the window they are only meaningful in.

``window`` travels with every snapshot so a reader can never hold a
fire count without the condition that makes it evidence. A scrape
that shows ``fired=0, window=graph_record`` is self-labelling: it
says "this window cannot fire this counter", not "this path is
dead".

Model-shipped sampling defaults from ``generation_config.json``.

vLLM's ``--generation-config auto`` (its default) backfills a request's
unspecified sampling params from the model's ``generation_config.json``.
A model that ships ``{"do_sample": true, "temperature": 1.0,
"top_k": 20, "top_p": 0.95}`` therefore runs ``top_k=20`` even when the
client omits ``top_k`` — the tail-clip the model was tuned for.

This module is the torch-free CPU-importable loader the engine reads at
boot. It does not decide policy (whether to apply the defaults) — that
is the ``--generation-config auto|none`` switch; it only parses the
file. The request path consults the loaded :class:`GenerationDefaults`
to fill only the sampling params the client omitted (an explicit request
value always wins).

Sampling defaults parsed from a model's ``generation_config.json``.

Each field is ``None`` when the file did not specify it (or no file
shipped). A ``None`` field is never applied — the request keeps the
server's hardcoded default. ``do_sample`` is captured for logging /
completeness; the request path does not branch on it (temperature
and the top-k/p clips already encode the sampling behavior).

The model's sampling defaults, with its non-thinking profile attached.

The shipped ``generation_config.json`` (or the arch fallback) supplies
the base profile; :data:`ARCH_NONTHINKING_DEFAULTS` supplies the
separate profile a hybrid-reasoning checkpoint uses when a request
resolves to non-thinking. The request path selects between them with
:meth:`GenerationDefaults.for_thinking`.

Parse sampling defaults from ``<model_dir>/generation_config.json``.

Torch-free. When the file is absent or unreadable, falls back to
the arch-level table (:data:`ARCH_GENERATION_DEFAULTS`) keyed by the
checkpoint's ``model_type`` — :data:`EMPTY` (``loaded=False``) when
neither applies; the caller logs that and keeps the server's
hardcoded defaults. Unknown / non-sampling keys in the file are
ignored. A present-but-malformed sampling value (wrong type) is
dropped (stays ``None``) rather than crashing the boot.

This profile's set fields laid over ``base``'s.

Layering (rather than replacing) keeps any field ``base`` defines
that this profile is silent about — an arch profile transcribes the
model card, which does not necessarily enumerate every key the
checkpoint ships.

The profile that backfills a request whose regime is ``thinking``.

``None`` (no regime information — the raw-completions surface has no
chat template and therefore no reasoning block) keeps the file's
profile, which is what a checkpoint ships as its single default.

Post-ready JIT/compile detector: no kernel compilation during serving.

A kernel compile (Triton JIT specialization, Dynamo/Inductor (re)compile,
``cpp_extension`` build or first-in-process dlopen) that runs after the
engine declared ready is a warmup gap: the latency lands on a live
request as an unexplained TTFT/TPOT spike, invisible to any flag audit.
This module makes every such event self-announcing instead of a mystery.

Two-phase contract, keyed on :func:`mark_serving` (called at the end of
``Engine.build`` — the same instant ``eng._ready`` flips):

* **Before serving-ready** (boot / warmup / capture): compiles are
  expected. Counted in ``jit_compile_boot``
  (:attr:`~arbi_serve.flag_truth.FiringWindow.BOOT`);
  :func:`mark_serving` logs one summary line — "compiled N kernels at
  warmup" with a per-source breakdown.
* **After serving-ready**: a compile that does not match the explicit
  accepted-residual allowlist (:data:`_ACCEPTED_RESIDUALS`) fires
  ``jit_compile_serving`` (:attr:`~arbi_serve.flag_truth.FiringWindow.
  SERVE`) — an anti-invariant witness like
  ``decode_pad_borrow_alias_consumed``: it must read 0 on a warmed
  serving run — and logs a WARNING naming the source, the kernel, and
  the specialization so the missing warmup coverage is identified from
  the log alone. Read it at ``GET /v1/admin/flag_truth`` →
  ``counters.jit_compile_serving``; a bench brackets its timed window
  with ``flag_truth/reset`` and asserts the counter is still 0 after.
* **Accepted residual** (allowlist match): a post-ready compile whose
  ``(source, name, detail)`` matches a named :class:`AcceptedResidual`
  entry — a genuinely-unbounded specialization that cannot be
  pre-warmed (a Dynamo frame whose guard family the ``recompile_limit``
  bounds back to eager) — is counted in
  ``jit_compile_serving_accepted`` (SERVE window) and logged at INFO,
  not the hard ``jit_compile_serving`` gate. This keeps the hard
  counter a clean binary "warmup coverage complete" assert: a residual
  we have chosen to accept lands in its own named counter (reported,
  never a gate failure) instead of muddying the MUST-NOT-FIRE. Every
  allowlist entry names its reason and the task tracking its real fix;
  reserve it for compiles that cannot be pre-warmed, never for a
  fixable warmup gap.

  MEASURE the closure before calling one unbounded. "Prefill shape
  classes are unbounded (arbitrary prompt geometries)" was this
  module's canonical example of an un-warmable specialization, and it
  was WRONG: the CuTeDSL prefill dispatch's kernel-cache key carries
  neither ``S_q`` nor ``S_kv``, and its batch element is the literal
  ``0 if is_varlen else B`` — sequence length reaches the key only
  through the asym-long tile boolean. The reachable set is a handful of
  keys per deployment, now enumerated and pre-compiled at boot by
  ``arbi_serve.engine.capture_admin.compile_warmup.
  _prewarm_tkv_prefill_kernels``. So a post-ready ``cute`` compile is a
  warmup-coverage gap and fires the hard gate like anything else. A
  matcher that accepts EVERYTHING from its source is not an allowlist
  entry, it is a mute button; ``tests/test_jit_prefill_prewarm_closure.
  py`` fails the build on one.
  A hard fire is also RECORDED. Triton hands the post-compile hook its
  own ``specialization_data`` blob, which ``JITFunction.preload``
  consumes to reproduce the compile with no launch to reconstruct, so
  :mod:`arbi_serve.jit_replay` persists it to the deployment's durable
  cache and the next boot's serve-kernel warmup preloads it. The
  detector is therefore not only the alarm for a warmup gap but the
  mechanism that closes it: coverage the hand-written warm set misses by
  enumeration is picked up empirically, once, and never paid by a live
  request again. Only the hard branch records — an accepted residual is
  declared un-warmable and a mutation-window compile is requested work.
* **Sanctioned mutation windows** (``CriticalSection`` — backend swap,
  calibration/model reload — and a re-entered ``Engine.build``: boot
  pool-member prep): compiles are part of the requested mutation, not a
  warmup gap. Counted in ``jit_compile_mutation_window`` (SERVE window)
  and logged at INFO. The drained/pre-traffic window means no live
  request pays the latency.

Hook sites — all on the compile path only. The no-compile hot path is
untouched: Triton's launch fast path never calls the JIT hooks (they sit
behind the in-process kernel-cache miss), the Dynamo callback runs only
inside ``convert_frame``, and the ``cpp_extension`` wrap is reached only
when an extension module materializes. A compile is already
milliseconds-to-minutes slow; the hook adds a counter increment and a
log line to it.

* ``triton.knobs.runtime.jit_post_compile_hook`` — fires once per
  in-process Triton specialization compile (both a fresh build and an
  on-disk-cache load; either way the specialization was cold in this
  process and the caller stalled on it). Chains any pre-existing hook.
* ``torch._dynamo.callback.callback_handler`` start callback —
  Dynamo/Inductor compiles and guard-fail RECOMPILES, lazy backward,
  runtime Triton autotuning, and compile-driven cudagraph recording
  (the ``CallbackTrigger`` enum), each with its ``compile_id``.
* ``torch.utils.cpp_extension._jit_compile`` wrap — every C++/CUDA
  extension materialization: tkv codec kernels (``load_inline``), EXL3
  trellis / AWQ-Marlin / xgrammar (``load``). A nested
  ``_write_ninja_file_and_build_library`` wrap distinguishes a fresh
  ninja/nvcc BUILD from a cached-``.so`` load in the event detail — and,
  because that inner wrap is the one seam EVERY extension build passes
  through, it is also where a build running during a boot announces
  itself as an open boot phase
  (:func:`~arbi_serve.engine.boot_heartbeat.announced_boot_step`). No
  call site in the boot brackets these builds: they are reached lazily
  from a model import, so without this they are minutes of boot that
  report no phase, no heartbeat and no line.
* ``cutlass.cute.compile`` wrap — the CuTeDSL compile entry behind the
  the tkv Turbo prefill dispatch's per-shape-class kernel cache (a post-ready
  miss is a seconds-scale prefill stall).

Known residual (documented, bounded): the tkv split-K decode autotune
sweep (bench-and-pick, no compile) is tkv-owned and carries no hook on
this side; its per-variant ``.so`` compiles do flow through the
``cpp_extension`` wrap, so a cold-variant stall still announces itself.

Per-process state: each TP rank / engine process installs its own hooks
and owns its own counters. Rank 0's counts are the serve counts (the
flag-truth convention); other ranks still warn in the shared log.

One explicitly-accepted post-ready compile class.

A post-ready compile whose ``source`` equals :attr:`source` and whose
``(name, detail)`` satisfy :attr:`match` is counted in
``jit_compile_serving_accepted`` and logged at INFO, instead of firing the
hard MUST-NOT-FIRE ``jit_compile_serving`` gate. Reserve this for compiles
whose specialization is genuinely unbounded and cannot be pre-warmed; each
entry states why it is accepted (:attr:`reason`) and the task tracking its
real fix (:attr:`tracking`). A fixable warmup gap gets fixed, not listed.

Per-source compile counts recorded so far in the boot phase.

A snapshot copy, so a caller can diff two of them to attribute the
compiles a given boot step paid for.

One line naming the frame and why it re-specialized.

Built from the END-callback metrics context. ``recompile_reason`` is the
field that carries both halves — the guard that failed and the user stack
that reached it — so it is preferred; ``co_*`` are used when dynamo filled
them and as the fallback for a FIRST compile, which has no guard to blame.

Attach the frame identity to a post-ready dynamo compile.

Called from the dynamo END callback. Silent unless this ``compile_id``
actually fired the must-not-fire gate — a warmup-window compile needs no
second line, and every compile paying for one would be noise.

Record one compile event; phase decides expected vs violation.

Callable directly by in-repo JIT paths that do not flow through the
hooked seams. Never raises — a telemetry hook must not turn a slow
compile into a failed one.

``replay_record`` carries the kwargs of
:func:`arbi_serve.jit_replay.record_triton_specialization` — the
Triton hook supplies it — and is persisted ONLY on the hard
``jit_compile_serving`` branch: a boot compile is already covered, a
mutation-window compile is requested work, and an accepted residual
is declared un-warmable, so none of the three is a replay target.

Flip to the serving phase; log the boot-compile summary once.

Latched: the first call (end of the first successful ``Engine.build``)
wins; later builds (pool members, admin model reload) run inside
admin windows and are classified there.

Replay payload for a Triton compile, or ``None`` when not applicable.

Triton hands every post-compile hook a ``specialization_data`` JSON
blob — signature, constexprs, specialization attrs, compiler options,
cache key and GPU target — which ``JITFunction.preload`` consumes to
reproduce the compile without reconstructing a launch. That blob plus
the kernel's module/qualname and its Triton source hash is the whole
record. Built only in the serving phase, so a boot compile (the bulk
of them) pays nothing, and defensively: an upstream hook-API change
degrades to "no record", never to a raise inside a compile.

The Triton source hash of a kernel about to be autotuned, or ``""``.

Read WITHOUT asking Triton to hash. ``JITFunction.cache_key`` is a
property that takes ``fn._hash_lock`` and walks the kernel's source and
dependency tree; both of those are exactly what Inductor strips from a
``CachingAutotuner`` before pickling it into the FX-graph cache (its
``prepare_for_pickle`` sets ``fn.fn``, ``fn.__globals__`` and
``fn._hash_lock`` to ``None``), and the static-autotuner copy that a
later boot reloads from that cache is never given them back — Inductor
itself launches the cached cubin and never needs them. Reading the
property on such a kernel raises, so this reads ``fn.hash`` first: the
plain string the property memoises on first compile, which is pickled
with the kernel and therefore present precisely on the reloaded object
that cannot be re-hashed. The property is consulted only for a kernel
that was never hashed (a compile in this process that has not gone
through the worker), and is guarded because a diagnostic must not fail
the autotune it describes.

Arm the Triton AUTOTUNE seam.

The ``jit_post_compile_hook`` above sees a Triton kernel being COMPILED.
It does not see ``CachingAutotuner.autotune_to_one_config``, which
benchmarks configs that are already compiled — so a kernel whose code came
from the Inductor cache but which has never been EXECUTED reaches serving
fully compiled and still benchmarks its config ladder on the first call.
That is a compile-class stall the gate was blind to, and on a multi-rank
boot it is worse than a stall: the ladder runs on whichever ranks own that
specialization, so ranks that do not share it sail past while their peer
benchmarks, and the group desynchronises inside the next collective.

``max_autotune=False`` does not disable this. That knob governs Inductor's
exhaustive template search; the per-kernel heuristic autotune runs
regardless, which is why "autotune is off" is not a reason to leave this
seam unhooked.

Recorded through :func:`record_compile` like every other seam, so a
post-ready autotune lands on the hard ``jit_compile_serving`` gate and
names the kernel and its cache key rather than presenting as a hang.

The wrapper runs INSIDE Inductor's autotuner on the first call of every
kernel the served forward uses, so a fault in the observation is a
fault in the forward: the whole record is guarded, a failure is logged
once with its reason, and the wrapped autotune always runs with its
arguments and result untouched. The served numerics never depend on the
detector.

Arm the Dynamo compile callback when Dynamo loads.

``torch._dynamo`` is the whole compiler frontend, sympy included — seconds
of import on every boot of every architecture. Nothing can trigger a Dynamo
compile without importing it first, so arming on the import keeps the
"installed before the first compile" contract without putting the frontend
on the path of a boot that never compiles.

Wrap ``cutlass.cute.compile`` — the CuTeDSL compile entry.

The tkv Turbo prefill dispatch caches compiled cute kernels per shape
class and calls ``cute.compile`` on a miss — seconds-scale latency.
tkv's call sites resolve ``cute.compile`` as a module attribute at call
time, so rebinding it when the module loads observes every miss, and a boot
that never reaches the CuTeDSL path never imports it.

Path of the ``.so`` this build would (re)produce, or ``None``.

Resolved by BINDING the real signature rather than indexing a
position, so an upstream argument reorder loses the probe instead of
reading the wrong argument as a directory.

Register every compile hook. Idempotent; per-process.

Call before the first compile can happen (``Engine.__init__``) so the
boot population is fully counted and the serving phase inherits armed
hooks rather than racing their registration.

Record-and-replay for the Triton specializations warmup does not reach.

The boot warmup pre-compiles kernels by ENUMERATION: a hand-written list
of launches over the sampler/kernel wrappers at the shapes a served
config is believed to reach (:mod:`arbi_serve.sampler.kernel_warmup` and
the serve-kernel warmup phase of ``Engine.build``). Enumeration is
structurally lossy — a new kernel, a new specialization axis, or a
sampler mode nobody listed drifts out of coverage silently, and the only
signal is :mod:`arbi_serve.jit_detector` naming the escapee AFTER
serving-ready, once a live request has already paid the dispatch.

This module closes the loop empirically instead of by enumeration:

**Record.** Every post-ready Triton compile the detector classifies as a
hard ``jit_compile_serving`` fire (a real coverage gap, not an accepted
residual) is persisted, with enough detail to reconstruct the compile.
The detail is Triton's OWN ``specialization_data`` — the JSON blob
``JITFunction._call_hook`` hands every post-compile hook, carrying the
signature, the constexpr values, the specialization attrs, the compiler
options, the cache key and the GPU target.

**Replay.** At the next boot's serve-kernel warmup phase,
:func:`replay_recorded_compiles` hands each blob back to
``JITFunction.preload``, Triton's supported inverse of that
serialization. ``preload`` re-enters the same ``_do_compile`` the live
dispatch would, keyed on the same cache key, and populates the
in-process kernel cache — WITHOUT reconstructing a launch. That matters:
the recorded blob describes the compile, not the call, so no tensor
shapes, pointers, strides or grid have to be guessed. Replaying a
LAUNCH would require exactly the argument reconstruction the record
cannot supply; replaying a COMPILE requires only what it already holds.

The alternative — recording specialization KEYS and driving the existing
wrapper functions with shapes chosen to reproduce them — is rejected:
inverting a cache key back to a set of argument shapes is a per-kernel,
hand-written mapping, i.e. the same enumeration this replaces, with the
added failure mode of silently warming a NEIGHBOURING specialization.

Scope: Triton only. The detector's other sources cannot be replayed from
their events — a Dynamo compile is identified by a frame id whose guards
are not serialized, a ``cpp_extension`` build is already served from its
own on-disk ``.so`` cache, and a CuTeDSL compile is keyed by the prefill
dispatch's own enumerated shape closure.

Cost, honestly: the recorded compiles are in-process cache misses served
from the persistent Triton on-disk cache, so replay MOVES their latency
from the first matching request to boot rather than deleting it. What it
buys is that no live request pays it and that
``jit_compile_serving`` — the must-not-fire warmup-coverage gate — reads
0 on the second and every later boot at a given build/arch/config.

Bounding. A record that grows forever, or that replays kernels the
current config stopped using, is a boot tax paid for nothing:

* **Keyed, not shared.** The record lives under a directory keyed by the
  torch/Python/CUDA build (the same ``cpp_ext_cache_key`` the Inductor
  and cpp-extension cache roots use) and uid-namespaced, in a file named
  by a hash of the full build identity (torch, Triton, CUDA, GPU arch,
  Python). A record from a different build or a different arch is a
  different file, so it is never read — and if one is reached anyway,
  the embedded identity is compared on load and refused. Triton's
  ``preload`` independently refuses a blob whose GPU target differs.
* **Source-pinned.** Each entry carries the recording JITFunction's
  ``cache_key`` (Triton's source+dependency hash). An entry is replayed
  only while that hash still matches the live kernel; a source edit drops
  it instead of preloading a blob against a signature it does not
  describe.
* **Capped and aged.** At most :data:`_MAX_ENTRIES` entries survive a
  write, newest first, and an entry older than :data:`_ENTRY_TTL_DAYS`
  is dropped. Ageing is the config-drift eviction: an entry a config
  stopped reaching is never renewed (nothing re-records a kernel that
  never compiles), so it expires on its own, while an entry still in use
  is re-observed post-ready once after it expires and re-recorded. The
  self-healing costs one request one dispatch per still-live entry per
  TTL, which is the price of not keeping a permanent list of kernels
  nobody launches. The record is deliberately NOT keyed on the model or
  the serving config: most of what escapes the enumerated warm set is
  sampler-level and shared across models, and a config-keyed file would
  send a deployment back to cold every time it moved ``max_batch`` or a
  drafter width. A co-resident engine's extra entries cost a preload
  each and age out on their own.
* **Budgeted.** Replay stops at :data:`_REPLAY_BUDGET_S` and logs what
  it did not reach, so a pathological record cannot stall a boot.

Degradation. An absent, unreadable, mismatched or empty record is a
logged no-op: the boot proceeds exactly as it does without this module,
unwarmed but correct. Nothing here raises into a compile, a boot, or a
request.

One recorded Triton compile, replayable via ``JITFunction.preload``.

``spec_data`` is Triton's ``specialization_data`` JSON verbatim;
``module``/``qualname`` locate the ``JITFunction`` to hand it back to,
and ``src_hash`` is that function's Triton source+dependency hash at
record time, so a source change invalidates the entry instead of
preloading a blob against a signature that has moved under it.

Build axes a compiled Triton specialization is bound to.

A record produced under a different torch, Triton, CUDA, GPU arch or
Python is not replayable: the cache key it carries encodes compiler
options this build may not produce, and the cubin behind it was built
for another target. Missing axes resolve to ``"unknown"`` rather than
raising, which still keeps the two environments in DIFFERENT files.

Compute capability of the visible GPU, or ``"nocuda"``.

Anything that stops the capability from being read resolves to the SAME
string, so a record and a replay in one process always agree on the axis
— a half-resolved arch would silently split one deployment's record in
two.

Directory the replay record lives in.

Resolution order:
  1. Explicit ``cache_dir`` argument (tests).
  2. ``ARBI_SERVE_JIT_REPLAY_DIR``.
  3. ``<persistent cache root>/arbi-serve/jit-replay-<build key>/u<uid>``
     — the same versioned, uid-namespaced shape the Inductor/Triton and
     cpp-extension cache roots use, so a torch bump or another uid on a
     SHARED volume lands in its own subtree instead of being served a
     mismatched record.

Record one post-ready Triton specialization for the next boot.

Called from the detector's Triton hook on a hard ``jit_compile_serving``
fire. Returns whether the entry was newly accepted. Never raises: this
sits on the compile path of a live request.

Merge this process's recordings into the on-disk record; return its size.

Read-merge-write so co-resident engines and TP ranks sharing one record
union their findings instead of clobbering each other. The write is
atomic (tempfile → fsync → rename), so a crash never leaves a
half-written record a later boot would replay from. A racing sibling can
still lose an entry between the read and the rename; that entry is simply
re-observed and re-recorded on a later boot.

Entries in the record for this build identity; ``[]`` when there is none.

Refuses — loudly, with an empty list — a record whose format version or
embedded build identity does not match this process. A record from
another torch/Triton/CUDA/arch build describes compiles this build
cannot reproduce, and replaying one would at best waste a boot and at
worst populate the kernel cache with a mismatched key.

The live ``JITFunction`` a recorded entry names, or ``None``.

Triton registers every ``JITFunction`` it constructs under
``"<module>:<qualname>"``, which is exactly what the post-compile hook
reports, so importing the module and reading that registry resolves even
a kernel that is not a module-level attribute. Attribute lookup is the
fallback, unwrapping an autotuner/heuristics decorator to the underlying
``JITFunction``.

``None`` when the entry may be preloaded, else why it may not.

Both checks are cheap pre-flights for failures ``preload`` would raise
on: the recorded blob names its own kernel, and its signature keys are
the kernel's parameter names in order. A kernel whose source or
parameter list moved is dropped rather than preloaded.

Whether this specialization is already in the in-process kernel cache.

``preload`` compiles unconditionally, so skipping what an earlier warmup
step already compiled keeps replay from paying for the same kernel twice.

Pre-compile every recorded specialization; return ``(preloaded, skipped)``.

Called from the serve-kernel warmup phase, before serving-ready. An
empty, absent or refused record is a clean no-op. Individual failures
are logged and skipped — a kernel that will not preload costs its first
request a dispatch, which the detector then names, exactly as it does
without a record.

Triton / CUDA kernels self-contained inside arbi-serve.

Each subpackage exports one logical kernel surface (e.g. ``fused_moe``)
and is import-safe on CPU-only hosts: heavy GPU imports happen inside
the call sites, never at module-load time.

AOT-compiled Triton kernel artifact loader.

What this is
============

This package is the **runtime side** of the AOT pipeline described in
``docs/aot-triton-pipeline.md``.

CI runs ``scripts/aot_compile_triton.py`` once per (kernel, target,
shape) tuple we want to ship as a binary artifact. The script writes
the resulting ``.cubin`` / ``.json`` / IR-stage files under
``arbi_serve/kernels/_aot/<kernel>__<arch>__<shape_hash>.<ext>``. At
engine boot, kernel modules call :func:`try_load` to pick up the
prebuilt artifact; if no artifact is present (or CUDA is unavailable,
or the shape key has no entry) the kernel falls back to the JIT path.

What this is *not*
==================

* Not a JIT replacement. Every kernel module keeps its
  ``@triton.jit`` definition and the canonical ``kernel[grid](...)``
  call as a fallback. AOT is an opt-in fast path.
* Not a serialiser of shape keys. The kernel module owns its shape
  key derivation (``derive_shape_key`` in each module). This package
  only provides the lookup/load primitives.
* Not a target-architecture detector. The kernel module passes the
  current arch ("sm_89", "sm_90", "sm_120") in. The loader compares
  exact strings — there is no "sm_90 binaries are forward-compatible
  with sm_120" hand-waving.

Constraints
===========

1. **CPU-safe import.** Importing this module on a CPU-only host MUST
   NOT touch ``triton`` until ``try_load`` is actually called from
   the kernel hot-path. The CPU gate imports the kernel modules
   transitively via ``arbi_serve.engine.*``, so any ``import triton``
   at module top would explode the gate.
2. **Idempotent.** ``try_load`` is called from the kernel hot-path
   on every invocation; the result must be cached in a process-local
   dict so we pay the file-IO cost once per (kernel, arch, shape)
   triple.
3. **Fail-soft.** Any error during load (missing files, a source
   mismatch, Triton API skew) returns ``None`` so the kernel falls
   back to JIT. A loud failure would break the engine for users on
   architectures we haven't shipped binaries for yet.

The source check
================

``(kernel_name, target, shape_hash)`` names WHICH kernel and WHICH
shape, and nothing about the kernel's SOURCE — so an edit to a
``@triton.jit`` body moves the machine code the JIT would produce
without moving any of the three. A shipped artifact then keeps
loading and the old cubin runs under a cheerful INFO line.

So the writer records ``ASTSource.hash()`` — Triton's own digest over
the JIT function's source (``fn.cache_key``, which walks the
function's dependencies too) plus the signature and constexprs — into
``<stem>.provenance.json``, and :func:`try_load` recomputes it from
the live ``ASTSource`` the kernel module hands in and compares.
A disagreement is a MISS: the kernel JITs from the source actually in
the tree, which is self-healing, and the WARNING names both digests.
An artifact with no provenance file cannot prove which source built
it and is refused on the same terms.

This is the boot-artifact keying rule of ``docs/memory-accounting.md``:
the filename triple answers the coarse question, the artifact carries
its own content signature for the one the triple cannot.

The artifact layout
===================

Each (kernel, target, shape_hash) triple is a group of files:

    _aot/<kernel>__<target>__<shape_hash>.cubin
    _aot/<kernel>__<target>__<shape_hash>.json
    _aot/<kernel>__<target>__<shape_hash>.ttir
    _aot/<kernel>__<target>__<shape_hash>.ttgir
    _aot/<kernel>__<target>__<shape_hash>.ptx
    _aot/<kernel>__<target>__<shape_hash>.llir
    _aot/<kernel>__<target>__<shape_hash>.source
    _aot/<kernel>__<target>__<shape_hash>.__grp__.json
    _aot/<kernel>__<target>__<shape_hash>.provenance.json

The ``.json`` file is Triton's per-kernel ``CompiledKernel.metadata``
serialisation; the ``.__grp__.json`` is the metadata-group manifest
(maps suffix → relative filename) that ``triton.compile`` uses
internally. The ``.provenance.json`` is ours: the source digest the
loader checks (see "The source check" above). The ``.source`` file is
Triton's own dump of the kernel's initial MLIR module — a debugging
aid, not an identity: reproducing it needs the compiler front-end, so
it cannot be the thing a load compares against. We mirror Triton's
per-hash directory layout but flatten
into a single directory keyed on a stable triple-hash so artifacts
across (kernel, target, shape) tuples don't collide and so
``pyproject.toml`` sdist excludes are easy.

See ``scripts/aot_compile_triton.py`` for the writer side.

Return True if the AOT loader should be bypassed completely.

Useful for development (force-recompile after kernel edits) and
for the JIT-vs-AOT parity tests, which need to exercise both
paths in the same process.

Return True if a missing AOT artifact should raise instead of
falling back to JIT.

Off by default. Operators flip this on to verify their deploy
contains every artifact the live workload exercises (and to fail
loud if a kernel hot-path isn't covered yet).

Return the current GPU target as ``sm_<arch>`` string, or
``None`` on CPU-only hosts.

Used as the per-target key in the AOT artifact filename. We use
``sm_89`` (RTX 4090) / ``sm_90`` (H100) / ``sm_120`` (RTX 5090)
as the canonical strings — matches what the AOT compile script
writes.

The identity of the SOURCE ``src`` compiles from.

``ASTSource.hash()`` — sha256 over ``fn.cache_key`` (Triton's digest of
the ``@triton.jit`` function's source AND the JIT functions/constants it
references), the arg attrs, the signature and the constexpr values. It is
target-independent by construction, which is what makes it the right axis
to add: the target already keys the filename.

Both sides of the pipeline call THIS function — the writer to record it,
the loader to check it — so the two cannot drift into comparing different
digests.

The source digest recorded beside this artifact, or ``None``.

``None`` means the artifact cannot say which source produced it — a file
that is missing, unreadable, or carries no digest. The caller refuses it;
"cannot verify" is never "verified".

Record ``src_hash`` beside the artifact triple; return the path written.

The writer side of :func:`read_provenance`. Lives here rather than in
``scripts/aot_compile_triton.py`` so the file's name, shape and key are
one expression with two readers.

Why this artifact triple may not be loaded, or ``None`` to load it.

Two ways it cannot be: the files are not all there, or the source it was
built from is not the source this process would JIT. The second is the one
the filename triple cannot see (see "The source check" in the module
docstring), and it is the one that runs stale machine code rather than
merely missing a fast path — so it is logged at WARNING, naming both
digests, while a plain absence stays silent at INFO.

Look up an AOT-compiled kernel for ``(kernel_name, current_target,
shape_hash)``. Returns the ``CompiledKernel`` if found, or ``None``
if we should fall back to the JIT path.

Args:
    kernel_name: stable name used by the AOT compile script (the
        same string the script writes into the artifact filename).
        By convention, this matches the ``@triton.jit`` function's
        ``__name__``.
    shape_hash: per-kernel shape-key string. The kernel module
        owns the shape→hash derivation (so the writer + reader
        agree). Hexdigest of a stable JSON dump of the
        constexpr/spec dict is fine.
    src: ``triton.compiler.ASTSource`` for the kernel — built by
        the kernel module's loader. ``CompiledKernel.__init__``
        stores this for runtime arg-name introspection
        (``self.src.fn.arg_names`` is read in ``__getitem__``);
        the launcher needs it to map positional args to slots.

Behaviour:
    * Returns ``None`` if CUDA is unavailable, the loader is
      disabled (``ARBI_AOT_TRITON_DISABLED=1``), or no artifact
      for the (kernel, current_target, shape_hash) triple is
      shipped.
    * Returns ``None`` (and logs at warning) when the shipped
      artifact's recorded source digest is absent or disagrees
      with ``src``'s — the triple names the kernel and the shape,
      never the source it was compiled from.
    * Returns ``None`` (and logs at debug) if loading the
      artifact raises. This is the fail-soft path: deploys that
      ship a partial set of binaries should still boot.
    * Caches the result (positive or negative) in
      ``_LOAD_CACHE`` so repeated calls are O(1).

Strict mode:
    * If ``ARBI_AOT_TRITON_STRICT=1`` and the lookup misses,
      raises ``FileNotFoundError``. Used by deploy verification —
      operators flip strict on, exercise the live workload, and
      confirm every kernel hot-path has a binary shipped.

Why pass ``src`` in instead of building it here:
    Each kernel has a different signature + different constexpr
    names. The kernel module is the owner of the signature dict
    — pushing that here would either duplicate it or require a
    registry. Letting the kernel module build the ASTSource and
    hand it to us keeps coupling minimal: this loader only knows
    about the file-naming scheme and the on-disk layout.

Memoized per-launch resolution of an AOT kernel.

Same contract as :func:`try_load`, but takes a **lazy** ``src_builder``
(zero-arg callable returning the ``ASTSource``) instead of a prebuilt
``ASTSource``. The point is the hot path: kernel modules call this on
EVERY launch, and building the ``ASTSource`` (``triton.compiler.ASTSource(...)``)
is not free. On the second+ launch of a given ``(kernel_name, shape_hash)``
this is a single dict lookup — ``src_builder`` is never invoked and
:func:`target_str` is not re-queried.

The resolved entry is cached under a ``("__resolved__", kernel_name,
shape_hash)`` key in the shared ``_LOAD_CACHE`` dict, so the existing
``_LOAD_CACHE.clear()`` (used by the JIT-vs-AOT parity tests to force a
re-resolve) invalidates it in lockstep with the per-target entries.

The disable knob is honored on every call BEFORE the resolved-cache
lookup, so flipping ``ARBI_AOT_TRITON_DISABLED`` back off (after a
cache clear) re-enables the AOT path exactly as ``try_load`` does.

Pad a 1- or 2-tuple grid to 3D, like ``JITFunction.run`` does.

``CompiledKernel.__getitem__`` skips the JIT runner's
grid-canonicalization step — it indexes ``grid[0], grid[1],
grid[2]`` directly. JIT callers can pass ``(N,)`` because
JITFunction.run pads internally, but the CompiledKernel runner
cannot. We pad here so the AOT-fast-path call sites stay
symmetric with their JIT fall-through siblings.

Single-launch conv-window commit for the MTP masked-replay rollback.

The masked replay advances each slab row's rolling conv window to the
accepted prefix by a pure gather: the window after committing tokens
``0..n`` is window ``n`` of ``[conv_slab[..., 1:] ‖ x_0 .. x_{T-1}]``.
Expressed in torch that is a concatenation, an ``arange``, an index add,
a gather, a compare and a masked write — ten launches per layer, most of
them at ``grid=[1,1,1]``, and the concatenation materialises a
``(rows, channels, Kw-1+T)`` staging tensor that is read once.

This kernel does the same addressing in one launch and reads only the
``Kw`` columns it emits. It is the CUDA arm of
:func:`~arbi_serve.models.recurrent_common.advance_conv_window_masked`;
that function's torch body remains the reference (and the CPU path).

Values are element copies of the same source bytes in both arms, so the
two are byte-identical by construction rather than by tolerance.

Whether this launch is the shape the fused kernel is defined for.

The base window may be the destination slab itself (the GDN
masked-replay convention) or a separate snapshot frame (Mamba-2's);
the kernel loads before it stores, so both are safe. Every tensor
must share one dtype: a mixed-dtype call carries a cast the torch
body performs and this kernel does not.

Keep FLA's WY triangular inverse in true fp32 on consumer Ada.

``chunk_gated_delta_rule_fwd_kkt_solve_kernel`` inverts the block-triangular
WY factor with 14 ``tl.dot(..., input_precision=SOLVE_TRIL_DOT_PRECISION)``
calls, on operands that are genuinely fp32 (``tl.zeros([BC, BC],
dtype=tl.float32)``). Upstream picks that precision from ``IS_TF32_SUPPORTED``,
which is ``capability[0] >= 8`` — so sm_89 takes the TF32 path.

TF32 is a DATACENTER trade. On TFLOPS)
in exchange for 13 mantissa bits. Consumer Ada was given no such headroom: a
4090 runs fp32 and TF32 at the same 82.6 TFLOPS. So on this hardware the
capability gate spends 23-bit mantissas down to 10 and buys nothing — on the
one operation in the chunk that amplifies error most, a triangular inverse.

This is not a disagreement with upstream, it is upstream's OWN rule applied to
a path it missed: ``fla/ops/utils/matmul.py`` already declines TF32 for fp32
operands (``allow_tf32 = False if a.dtype == torch.float32 else True``). The
gated-delta-rule kernel is the outlier inside its own library.

WHY REACHING INTO THE MODULE IS SAFE. Triton folds ``constexpr`` module
globals into the JIT cache key (``triton/runtime/jit.py``: ``self.hash +=
str([(name, val) ... if isinstance(val, constexpr)])``), so changing this
changes the key — a cached TF32 binary cannot be served to an ``ieee`` build.
The kernel compiles at first forward, long after import time, so setting it
here is early enough. ``tests/test_fla_solve_tril_precision.py`` pins both
halves of that argument.

The override is a stopgap with a defined end: it goes when the pinned FLA
makes the gate capability-aware rather than ``>= 8``.

Point FLA's WY solve at true fp32; return the value it now holds.

``module`` is the FLA module carrying the constant, defaulting to
``fla.ops.gated_delta_rule.chunk_fwd``. It is a parameter so the test can
drive this against a stand-in whose starting value is TF32 — a CPU-only
box has ``IS_TF32_SUPPORTED`` False and already reads ``ieee``, which
would make a test against the real module pass without the override doing
anything.

Raises ``RuntimeError`` if the constant is gone. A silently-inert override
is worse than none: it reads as a precision guarantee while the kernel
keeps taking TF32.

Vendored / inference-only forward path for FLA's chunk_gated_delta_rule.

Provenance
----------
- Upstream: ``flash-linear-attention`` (``fla-core``) v0.5.0
- Upstream commit: matches the wheel installed under
  ``site-packages/fla`` at the time this directory was created
  (FLA 0.5.0 — ``flash_linear_attention-0.5.0.dist-info``).
- Upstream MIT license — see :file:`LICENSE.upstream`.

Why vendor
----------
The upstream ``fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule``
public entry point is decorated with ``@torch.compiler.disable``.  That
decorator forces Dynamo to skip the function frame at compile time,
which means the FLA kernel chain stays in pure-eager Python during a
``torch.compile`` / cudagraph capture window.  A captured graph's
contract requires every operation to be either
in-eager-without-side-effects or compiled-and-graph-safe, so the
whole-forward prefill cudagraph cannot span a
``@torch.compiler.disable``-d frame: a ``StateKind.GDN`` hybrid on
FLA's public entry point is refused by the capture pre-flight.

Beyond the decorator, the upstream public API also routes through
``ChunkGatedDeltaRuleFunction.apply`` (``torch.autograd.Function``),
which adds:

  * an autograd dispatch escape Dynamo cannot trace through cleanly,
  * ``@input_guard`` / ``@autocast_custom_fwd`` decorator wrappers
    that introduce Python frames inside the compile region.

This vendored copy provides a **forward-only inference** entry point
that:

  1. Drops ``@torch.compiler.disable``.
  2. Bypasses the autograd ``Function.apply`` indirection (arbi-serve
     runs under ``torch.inference_mode()``; no backward needed).
  3. Imports the underlying kernel-chain helpers directly from the
     installed FLA wheel — we do NOT re-vendor every ``@triton.jit``
     kernel source file, because each kernel-level helper
     (``chunk_gated_delta_rule_fwd_intra``, ``chunk_gated_delta_rule_fwd_h``,
     ``chunk_fwd_o``, ``chunk_local_cumsum``) carries no
     ``@torch.compiler.disable`` decorator upstream — only the public
     ``chunk_gated_delta_rule`` wrapper does.  Importing the helpers
     directly keeps the vendoring surface minimal and lets future FLA
     fixes flow through transparently.
  4. Provides an extension point for routing intermediate workspace
     allocations through a persistent ``NamedMemPool`` (Phase 2 of the
     vendoring sprint — see :mod:`arbi_serve._fla_persistent_cache`
     for the existing chunk-fwd workspace-pinning pattern this will
     subsume).

Status
------
- This is the PRODUCTION GDN path. ``ARBI_GDN_VENDORED_FLA`` defaults
  ON and is worth B=1 TTFT 275 -> 162.5 ms (-40.9%) on Qwen3.5-0.8B
  (see :mod:`arbi_serve.models._gdn_kernels`); ``=0`` falls back to
  FLA's public ``chunk_gated_delta_rule``.
- Numerics: byte-identical to FLA's ``chunk_gated_delta_rule`` public
  API on greedy ``temp=0`` ``seed=42`` prompts (see
  :mod:`tests.test_fla_vendored_parity`).
- Capture-safety: dropping ``@torch.compiler.disable`` is what makes a
  GDN hybrid admissible to the whole-forward prefill cudagraph at all.
  The pre-flight
  (:func:`arbi_serve.runtime.capture.preflight._can_capture_prefill`)
  admits GDN-ONLY hybrids under ``ARBI_GDN_PREFILL_CAPTURE``, which
  also defaults ON; Mamba and ShortConv state kinds still refuse.

Inference-only forward path for ``chunk_gated_delta_rule``.

Vendored equivalent of the upstream FLA 0.5.0 public entry point
(:func:`fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule`),
forward-only and **without** the upstream ``@torch.compiler.disable``
decorator or the ``torch.autograd.Function`` indirection.

Upstream reference (FLA 0.5.0, ``fla/ops/gated_delta_rule/chunk.py``):

    @torch.compiler.disable
    def chunk_gated_delta_rule(q, k, v, g, beta, scale=None, ...):
        ...
        o, final_state = ChunkGatedDeltaRuleFunction.apply(
            q, k, v, g, beta, scale, initial_state, output_final_state,
            cu_seqlens, cu_seqlens_cpu,
            use_qk_l2norm_in_kernel,
            cp_context, transpose_state_layout,
            use_gate_in_kernel, A_log, dt_bias,
        )
        return o, final_state

Where ``ChunkGatedDeltaRuleFunction.forward`` orchestrates::

    chunk_indices = prepare_chunk_indices(cu_seqlens, 64, ...)
    g, o, A, final_state, _, _ = chunk_gated_delta_rule_fwd(
        q,
        k,
        v,
        g,
        beta,
        scale,
        initial_state,
        output_final_state,
        cu_seqlens=cu_seqlens,
        cp_context=None,
        chunk_indices=chunk_indices,
        transpose_state_layout=False,
        use_gate_in_kernel=False,
    )
    return o.to(q.dtype), final_state

We replicate that orchestration here, but:

  * **No autograd ``Function.apply``** — inference-only path; arbi-serve
    runs under ``torch.inference_mode()`` so the autograd machinery
    is dead weight.  Saves ~3-5 Python frames per call AND removes
    the autograd dispatch escape that Dynamo cannot trace cleanly.
  * **No ``@torch.compiler.disable``** — the whole point of the
    vendoring; lets Inductor compile through this frame.
  * **No ``@input_guard`` / ``@autocast_custom_fwd``** — those
    decorators (FLA's :mod:`fla.utils`) wrap the function with
    contiguity-checks and autocast-state restoration that arbi-serve's
    inference path doesn't need (we control input contiguity at the
    GDN-block call site, and there's no autocast in the inference
    forward).
  * **No ``cp_context`` branch** — context-parallel is a training-time
    feature; arbi-serve is single-rank inference (TP-only when needed,
    no CP).
  * **No ``use_gate_in_kernel`` branch** — arbi-serve's GDN block
    pre-computes the log-decay gate before the call (matches FLA's
    default ``use_gate_in_kernel=False``).
  * **No backward pass** — saved tensors / ctx not needed.

This shaves ~7 Python frames per layer per prefill (per the
:mod:`project_arbi_serve_gdn_inline_prefill_parked` profile + the
upstream call-stack analysis in :mod:`arbi_serve.kernels.fla_vendored`)
and gives Inductor an unobstructed view of the kernel-chain.

Route allocations made inside the ``with`` body to ``pool``.

Mirrors the wrapping context manager in
:func:`arbi_serve._fla_persistent_cache._build_workspace_routing_wrapper`.
Falls through (no routing) when the pool is None, when CUDA is
unavailable, or when the underlying primitives aren't importable
(CPU/test paths).

**Cudagraph-capture awareness.**  When a ``torch.cuda.graph(pool=...)``
block is already active on the current stream, the allocator is
already routed to the captured graph's pool.  Calling
``_cuda_beginAllocateCurrentThreadToPool`` for a DIFFERENT pool
inside that block corrupts the capture state — every subsequent
allocation inside the capture lands in OUR pool instead of the
cudagraph's pool, the capture's ``endAllocateToPool`` never matches
a live ``begin``, and the capture aborts with
``cudaErrorStreamCaptureInvalidated``.  This is exactly the
failure observed on Phase-3 prefill capture (``perf/fla-vendor-
capture-safe`` tip ``dcff01f``): every bucket failed because our
workspace-pool routing redirected allocations away from the
``graph_pool`` mid-capture.

Fix: detect the in-progress capture via
:func:`torch.cuda.is_current_stream_capturing` and no-op the
begin/end pair.  Allocations during capture flow into the
cudagraph's pool, which is the intended behaviour — the captured
graph's mempool already keeps the workspace-tensor footprint
bounded across the multi-bucket sweep (each bucket's intermediates
live in the captured pool's segment for that bucket).  The
workspace-pool routing remains active OUTSIDE capture (live
forward path), where it serves its original purpose of bounding
the per-call allocator footprint.

Mirrors the same guard in
:meth:`arbi_serve.runtime.named_pool.NamedMemPool.use`.

Return the registered workspace pool, or ``None``.

Reads :data:`arbi_serve._fla_persistent_cache._GDN_WORKSPACE_POOL`
— the engine boot path registers it via
:func:`arbi_serve._fla_persistent_cache.set_gdn_workspace_pool`.
Lazily imported so the vendored package stays importable when the
persistent-cache module is missing (test isolation).

Forward kernel orchestrator — mirrors FLA's ``chunk_gated_delta_rule_fwd``
minus the cp_context / use_gate_in_kernel branches.

``emit_chunk`` / ``emit_out`` — the in-kernel state emit
(:mod:`arbi_serve.kernels.fla_vendored.chunk_delta_h_emit`): the fp32
running state at the start of the named chunk of each sequence, written
beside the fold with no second launch. Routes the state-evolution step
through the vendored emit kernel; every other step is upstream's.

All intermediate allocations route through the registered workspace
pool (when set via :func:`arbi_serve._fla_persistent_cache.set_gdn_workspace_pool`)
so the captured-graph mempool footprint stays bounded across
multi-bucket prefill capture sweeps.  When the pool is None, the
routing context is a no-op passthrough — same behaviour as the
upstream FLA call.

``o_out`` — a caller-owned ``(1, T, HV, V)`` buffer in ``q.dtype`` for
the output kernel to write into, in place of the buffer upstream's
``chunk_fwd_o`` allocates. A row-class-mixed batch hands in a SLICE of
the block's one core buffer here, so the prefill rows' output lands
beside the verify rows' with no concatenation. Same kernel, same grid,
same bytes; only the destination pointer differs.

Returns ``(o, final_state)`` cast to ``q.dtype``.

``chunk_fwd_o`` with a caller-owned destination.

The upstream wrapper is ``o = torch.empty_like(v)`` followed by ONE
launch of ``chunk_fwd_kernel_o``; this is that launch against ``o``.
``o`` must match ``v``'s shape, dtype and contiguity — a leading-dim
slice of a contiguous buffer qualifies.

Forward-only inference entry point — drop-in replacement for
:func:`fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule`.

``o_out`` — optional caller-owned output buffer; ``emit_chunk`` /
``emit_out`` — the in-kernel state emit; see
:func:`_chunk_gated_delta_rule_fwd_orchestrator`.

Differences from upstream:

  * No backward / autograd machinery (no ``Function.apply``).
  * No ``@torch.compiler.disable`` decorator.
  * No ``cp_context`` parameter (single-rank inference).
  * No ``use_gate_in_kernel`` parameter (arbi-serve's GDN block
    pre-computes log-decay externally).
  * ``transpose_state_layout`` fixed to ``False`` (matches the
    :class:`GdnLayerView` slab layout in
    :mod:`arbi_serve.cache.recurrent_pool`).
  * Optional ``chunk_indices`` parameter — when provided, skips
    the in-call ``prepare_chunk_indices`` (which contains a
    ``.tolist()`` host sync that breaks cudagraph capture).
    Callers can pre-compute ``chunk_indices`` outside the
    captured region (in the metadata-builder finalize step) and
    pass it through.  When ``None`` AND ``cu_seqlens`` is set,
    falls back to the in-call ``prepare_chunk_indices`` — same
    as upstream behaviour.

Args mirror upstream — see
:func:`fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule`.

Returns ``(o, final_state)`` where ``o`` has shape ``[B, T, HV, V]``
cast to ``q.dtype`` and ``final_state`` has shape ``[N, HV, K, V]``
if ``output_final_state`` else ``None``.

``chunk_gated_delta_rule_fwd_h`` with an in-kernel state EMIT.

The recurrent fold carries its state ``b_h`` in fp32 registers across the
chunks of a sequence, stores it to ``h`` in the activation dtype at every
chunk start (the output kernel reads that), and stores it in fp32 to ``ht``
once, at the end. Everything a caller can observe between those two is a
rounded copy. This variant adds ONE more fp32 store: at the start of chunk
``emit_chunk[seq]`` (``-1`` = none) the running state is written to
``emit_out`` — ONE ``(HV, V, K)`` buffer in the SLAB's layout (V outer, K
inner), not ``ht``'s ``(K, V)``. A restore writes these bytes straight back
into a slab row, so they leave the kernel already in the slab's order and no
transpose sits between the emit and the host copy. One buffer, not one per
sequence: a step arms at most one row (``Scheduler._arm_fold_split``), the
buffer is a boot allocation on the ledger, and sizing it per row would price
``max_batch`` copies of a state only one row ever writes. Two rows naming a
chunk in one launch would race on it, so the caller guarantees at most one.

That value is, by construction, the ``ht`` the fold would have produced had
the sequence been cut at that chunk boundary — the state the two-segment
fold (#2189) materialises with a second launch and a rebased ``cu_seqlens``.
The fold-grid constraint is the kernel's own addressing (a chunk index), so
an off-grid position has no representation here rather than a rounded one.
The emit table is data: ``emit_chunk`` is an ``(N,)`` int32 tensor and
``emit_out`` a preallocated buffer, so the same launch serves a step with or
without a savepoint, eager or captured — and a launch with no row armed is
issued WITHOUT the table, on upstream's kernel, so a step that takes no
savepoint pays nothing for the ones that do.

Kernel text is upstream's verbatim apart from the emit parameters, the
pointer setup and the store block; the store block is the ``ht`` store
retargeted. ``tools/gdn_row_class/check_emit_state.py`` measures the
bit-identity against the two-segment fold on a card.

Upstream ``chunk_gated_delta_rule_fwd_h`` (``state_v_first=False``,
``save_new_value=True``, no ``gk``) plus the emit table.

``emit_chunk`` — ``(N,)`` int32, the chunk index (relative to each
sequence's own start) at whose START the state is emitted, ``-1`` for
none; AT MOST ONE entry may be armed (the caller's contract — two armed
rows would race on the one buffer). ``emit_out`` — ONE ``(HV, V, K)``
buffer in the slab's layout, written only by the armed row; fp32 is the
bit-identical contract, any other float dtype receives the kernel's
rounding of it.

Triton device functions for the 4-bit float element codes.

``e2m1_decode`` maps an E2M1 nibble to its signed float value; it is the
element decode shared by NVFP4 (16-element blocks, fp8_e4m3 block scale
times an fp32 per-tensor scale) and MXFP4 (32-element blocks, E8M0 block
scale). ``e8m0_decode`` maps an E8M0 byte to ``2**(byte - 127)``.

Signed float value of a 4-bit E2M1 code ``[sign | exp(2) | mant(1)]``.

Magnitudes by the low 3 bits are ``{0, .5, 1, 1.5, 2, 3, 4, 6}``,
reconstructed as ``2**(e-1) * (1 + 0.5*m)`` for ``e >= 1`` and
``0.5*m`` for ``e == 0``.

``2**(code - 127)`` for a biased-exponent E8M0 byte.

The bias is written as an fp32 literal rather than read from
:data:`arbi_serve.kernels.fp4_format.E8M0_BIAS`: Triton binds a python
``int`` global as an integer constexpr, which changes the subtraction's
result type and mis-decodes the scale.

Fused residual-add + RMSNorm Triton kernel.

Computes ``residual += x`` in place and returns the RMSNorm of the new
residual (scaled by ``weight``, cast back to ``x``'s dtype), collapsing a
pre-norm block's residual add and the 7-8 unfused RMSNorm ops into a
single Triton launch.

One row per program; fused add + RMSNorm in a single launch.

Semantics (per row r):
    new_resid[r, :] = residual[r, :] + x[r, :]   # fp32-accumulated
    residual[r, :]  = new_resid[r, :]            # in-place store
    rms             = sqrt(mean(new_resid ** 2) + eps)
    out[r, :]       = (new_resid / rms) * weight # cast back to x's dtype

The fp32 accumulation pattern mirrors the eager Python reference
(``x32 = x.float(); var = x32.pow(2).mean(...)``) so the result is
bit-equivalent to the unfused path within bf16 ULP noise.

The kernel does TWO passes over the row:
  1. compute (residual + x), accumulate sum-of-squares, store the
     summed value into the residual buffer.
  2. reload residual (now == new_resid), normalize, multiply by
     weight, write to out.

Pass 1's store + pass 2's load communicate through the residual
buffer — this avoids a per-row scratch tensor (which would either
cost an extra global allocation or burn shared memory on large H).
Triton elides the round-trip on H <= ~16K rows in practice; for H
> BLOCK_SIZE we issue a tiled loop in both passes.

Fused (residual + RMSNorm) — single Triton launch.

Args:
    x: (..., H) input. bf16 / fp16 / fp32. Read-only at the Python
        level; the kernel uses x's pointer but does not write back
        into x.
    residual: (..., H) running residual stream. Same trailing dim
        as x. **Updated IN PLACE** to ``residual + x`` — callers
        who need the OLD residual must clone before calling. Same
        dtype as x (the kernel writes back in residual's dtype).
    weight: (H,) RMSNorm scale. fp32 (per the model's RMSNorm
        convention; the dense Python path also keeps this fp32).
    eps: numerical floor for the variance reciprocal-sqrt.

Returns:
    (..., H) tensor of the same shape and dtype as ``x``: the
    RMSNorm of the new residual (``residual + x``) multiplied by
    weight, cast back to x's dtype. The new residual is available
    via the in-place mutation of the ``residual`` argument.

Bit-equivalence: the kernel matches the unfused
``residual + x; (x32 / rms) * weight`` reference within bf16 ULP
noise on contiguous inputs (the two paths use the same fp32 add
+ fp32 sum-of-squares + fp32 rsqrt + fp32 weight-mul ordering).

Return the AOT shape-key string for ``(BLOCK_SIZE, x_dtype, weight_dtype)``.

Mirrors the writer side in
``scripts/aot_compile_triton.py::_spec_fused_add_rms_norm``. Keep
in lockstep — a writer/reader skew here means the loader can't
find the artifact and silently falls back to JIT.

Only the production combination
``(BLOCK_SIZE=1024, x=bf16, weight=fp32)`` is shipped; the rest
fall through to JIT.

Look up an AOT-compiled artifact for this kernel.

Returns the ``CompiledKernel`` if found, or ``None`` to fall
through to JIT. The resolution is memoized in
:mod:`arbi_serve.kernels._aot` keyed on the AOT shape key, so the
second+ launch of a given ``(block_size, x_dtype, weight_dtype)``
is a single dict lookup — neither :func:`_build_ast_source` nor the
per-target ``try_load`` filesystem recheck runs again.

Build the ASTSource the AOT loader needs at runtime to bind
arg names. The compiled-kernel runtime path consults
``self.src.fn.arg_names`` to translate positional args into the
Triton launcher slots — so the ASTSource we hand to ``try_load``
must wrap the same JIT function the writer compiled.

Fused MoE Triton kernel + thin invoker — vendored from vLLM.

The Triton kernels in :mod:`fused_moe_kernel` are a verbatim extract of
``vllm.model_executor.layers.fused_moe.fused_moe`` (Apache-2.0,
``Copyright contributors to the vLLM project``); the invoker, the
``moe_align_block_size`` binning kernels, the JSON config loader, and the
device-name helpers are re-implemented here so that nothing in this
subpackage imports vLLM at runtime. Calling code lives one layer up
in :mod:`arbi_serve.models.moe` (``FusedMoE``).

See ``THIRD_PARTY.md`` next to this file for the upstream commit hash
the kernel was synced from.

Fused MoE Triton kernels.

Three of the four ``@triton.jit`` kernels in this file —
:func:`write_zeros_to_output`, :func:`fused_moe_kernel`, and
:func:`fused_moe_kernel_gptq_awq` — are vendored from vLLM (Apache-2.0)
from ``vllm/model_executor/layers/fused_moe/fused_moe.py`` (commit
``25006e567f0e11d91c0cb5e23945666a8ac6a9e6``). The first two are
byte-for-byte; :func:`fused_moe_kernel_gptq_awq` carries one local
divergence — its K loop is chunked at the quantization group so the
scale / zero-point gathers are ``[1, BLOCK_SIZE_N]`` rather than
``[BLOCK_SIZE_K, BLOCK_SIZE_N]``. See ``THIRD_PARTY.md`` in this
directory for the divergence record and the re-sync procedure. The
vLLM-internal imports (``vllm.envs``, ``vllm._custom_ops``,
``vllm.triton_utils``, ``vllm.platforms``) are stripped and replaced
with direct ``triton`` imports — these kernels themselves only call
``triton.jit`` / ``triton.language`` primitives, no vLLM runtime hooks.

:func:`fused_moe_kernel_mxfp4` and :func:`fused_moe_kernel_nvfp4` are
first-party. The former is the OCP MXFP4 expert GEMM,
structured like the vendored ``use_int4_w4a16`` path (nibble unpack along
K via ``b_shifter``, grouped scale load) with the E2M1 decode and an E8M0
block scale in place of the zero-point subtraction, plus the per-expert
bias gpt-oss checkpoints carry. It is kept out of the vendored kernels so
those stay byte-identical to upstream. The latter is the same GEMM over
the modelopt NVFP4 pair (an FP8-E4M3 factor per 16 columns times the
linear's per-tensor ``weight_scale_2``).

Implements the fused computation for a Mixture of Experts (MOE) using
token and expert matrices, with on-the-fly INT4/INT8 weight dequant.
See the bf16/fp8 kernel ``fused_moe_kernel`` below for the parameter
semantics; the GPTQ/AWQ variant adds ``b_scale_ptr`` / ``b_zp_ptr``
for per-group scales / zero-points.

``b_scale`` / ``b_zp`` are per-``group_size`` along K, so the K-tile is
walked in ``min(group_size, BLOCK_SIZE_K)``-row chunks and each chunk
gathers a single ``[1, BLOCK_SIZE_N]`` scale / zero-point row, which
broadcasts over the chunk's dot. ``BLOCK_SIZE_K`` and ``group_size``
must therefore divide one another.

Implements the fused computation for a Mixture of Experts (MOE) using
token and expert matrices.

Key Parameters:
- A: The input tensor representing tokens with shape (*, K), where '*' can
    be any shape representing batches and K is the feature dimension of
    each token.
- B: The stacked MOE weight tensor with shape (E, N, K), where E is
    the number of experts, K is the input feature dimension, and N is
    the output feature dimension.
- C: The output cache tensor with shape (M, topk, N), where M is the
    total number of tokens post padding, topk is the number of times
    each token is repeated, and N is the output feature dimension.
- sorted_token_ids: A tensor containing the sorted indices of tokens,
    repeated topk times and arranged by the expert index they are
    assigned to.
- expert_ids: A tensor containing the indices of the expert for each
    block. It determines which expert matrix from B should be used for
    each block in A.
- naive_block_assignment: A boolean flag indicating whether to use naive
    token wise block assignment. If True, each block corresponds to a
    single token.
This kernel performs the multiplication of a token by its corresponding
expert matrix as determined by `expert_ids`. The sorting of
`sorted_token_ids` by expert index and padding ensures divisibility by
BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
multiplication across different blocks processed by the same expert.

Fused MoE GEMM against MXFP4-packed expert weights.

``K`` is the REAL (unpacked) reduction extent; ``b`` holds ``K // 2``
bytes per row, real column ``k`` living in byte ``k // 2`` (LOW nibble
when ``k`` is even). ``b_scale`` holds one E8M0 exponent per
``group_size`` real columns. ``b_bias`` is added after dequantization
and before the router-weight multiply.

Parameter semantics otherwise match :func:`fused_moe_kernel`.

Fused MoE GEMM against NVFP4-packed expert weights.

:func:`fused_moe_kernel_mxfp4`'s structure with the modelopt NVFP4
scale pair in place of the single E8M0 exponent: ``b_scale`` holds one
FP8-E4M3 factor per ``group_size`` real columns, and ``b_gscale`` holds
the per-tensor ``weight_scale_2`` of the checkpoint linear each output
row came from — ``GSCALE_ROWS`` consecutive rows of ``N`` share one
entry, which is how the gate/up seam inside ``w13`` keeps its two
source linears' factors apart. The row-constant factor is applied to
the f32 accumulator after the K loop rather than inside it.

Launch wrappers + token-to-expert binning for the fused-MoE kernels.

Pure-arbi-serve code: nothing here imports vLLM. The kernels live in
:mod:`fused_moe_kernel`; the bf16/fp8 and GPTQ/AWQ ones are vendored
verbatim (Apache-2.0, attribution preserved on that file) and the MXFP4
one is first-party.

Three pieces:

1. :func:`moe_align_block_size` — bins ``topk_ids`` by expert and pads
   each expert's token list up to ``BLOCK_SIZE_M``. Three first-party
   Triton kernels (histogram → scan → scatter) implement a stable
   counting sort; off-device operands run
   :func:`~arbi_serve.kernels.fused_moe.reference.moe_align_block_size_reference`,
   which is the executable spec both agree on. SYNC-FREE: every
   intermediate stays on device and the output buffers are sized to a
   static worst case, so no shape depends on device data and the decode
   graph stays capturable.

2. :func:`try_get_optimal_moe_config` — reads the JSON tuned configs
   shipped under ``configs/`` (tuned configs cover RTX 4090 fp8_w8a8;
   bf16 on Ada falls through to a heuristic default). Same shape-key
   format as upstream so tuned configs from vLLM's tree work as-is.

3. :func:`invoke_fused_moe_kernel` / :func:`invoke_fused_moe_kernel_awq`
   / :func:`invoke_fused_moe_kernel_mxfp4` /
   :func:`invoke_fused_moe_kernel_nvfp4` — kernel launchers. They take
   a pre-computed ``config`` dict and forward to the right Triton entry.

Pass 1: fill the pad sentinel and histogram each chunk of slots.

One program per ``BLOCK_T``-slot chunk of the OUTPUT, which is never
shorter than the input, so the same grid covers both jobs. The
histogram is per-chunk rather than global because pass 3 needs each
chunk's exclusive prefix to place slots in stable order.

Pass 2: turn per-chunk histograms into per-(chunk, bin) row bases.

Single program: the whole scan is over ``n_bins`` (a few hundred),
and the cross-bin cumsum is inherently serial in the bin axis.

Pass 3: place every slot, and label every block with its expert.

The two jobs are independent and both indexed by ``BLOCK_T``-sized
chunks (of slots / of blocks), so one grid covers them and the whole
binning costs three launches.

Sort tokens by their assigned expert and pad each expert's run.

Equivalent of vLLM's ``vllm._custom_ops.moe_align_block_size`` CUDA
kernel, as three Triton passes — histogram, scan, scatter. Called
twice per MoE layer per step, so the launch count IS the cost at
decode: the routed GEMMs it feeds are tiny.

The scatter places each slot at ``row_start[bin] + (number of earlier
slots in the same bin)``, i.e. a stable counting sort. An atomic
bump-pointer would be one pass shorter but would order each bin by
race, so the same routing would not reproduce the same
``sorted_token_ids`` twice.

Parameters
----------
topk_ids
    ``(num_tokens, top_k)`` int tensor; each entry is the expert id
    chosen for that ``(token, slot)`` pair. A **negative** id marks a
    slot this rank does not own (expert parallelism: the router runs
    replicated, so every rank sees ids for experts other ranks hold).
    Such slots are binned separately and their blocks carry expert id
    ``-1``, which makes the kernel write zeros for them — the owning
    rank supplies the real contribution through the trailing EP
    all-reduce.
block_size
    ``BLOCK_SIZE_M`` of the kernel — each expert's slice is padded
    to a multiple of this.
num_experts
    LOCAL expert count — the size of the ``(E, ...)`` weight tensor
    the kernel indexes with ``expert_ids``.

Returns
-------
sorted_token_ids, expert_ids, num_tokens_post_pad
    Same triple semantics as upstream:

    - ``sorted_token_ids[i]`` is the *flat* index into the
      ``(num_tokens * top_k)``-flattened ``topk_ids`` for the
      ``i``-th slot of the padded sequence. Slots whose value
      is ``>= num_valid_tokens`` are pad slots (the kernel's
      ``token_mask`` filters them out).
    - ``expert_ids[b]`` is the expert id this BLOCK_SIZE_M-sized
      block is processing (constant across the whole block by
      construction), or ``-1`` for a non-local / unused block.
    - ``num_tokens_post_pad`` is a 1-element int tensor giving
      the actual padded token count; the kernel uses this to bail
      out of trailing blocks.

The two buffers are sized to the STATIC worst case
``num_valid + min(num_experts + 1, num_valid) * (block_size - 1)``
(each occupied bin wastes at most ``block_size - 1`` pad slots, and
at most one bin per routed slot can be occupied), rounded up to a
whole number of blocks. Sizing them from the actual per-expert counts
would need those counts on the host — one device sync per MoE layer
per step, which is exactly the stall this path exists to avoid, and
a data-dependent shape no CUDA graph can capture.

Launch math for the three binning passes.

Split out from :func:`moe_align_block_size` so the parity test can
drive the real kernels under Triton's interpreter, where the operands
are CPU tensors and the device dispatch above would route around
them.

Same shape as vLLM's ``current_platform.get_device_name()`` —
underscores in place of spaces, no trailing decoration.

Memoized — the tuned-config lookup that calls this is reached from
inside the compiled decoder block on its first miss.

Heuristic config when no tuned JSON is available.

Lifted from the bf16/fp8 branch of vLLM's ``get_default_config``;
drops the ROCm and BATCH_INVARIANT branches since arbi-serve
targets CUDA only and never sets those toggles. The ``int4_w4a16``
block-shape branch is also dropped (the AWQ wrapper always loads
a tuned config or sets BLOCK_SIZE_K explicitly). See vLLM's
``vllm/model_executor/layers/fused_moe/fused_moe.py`` for the
architectural rationale behind each tile-size choice.

Pick the tuned config closest to this ``M`` (token count), or
fall back to a heuristic default. Mirrors vLLM's same-named helper.

``w1_shape`` / ``w2_shape`` are the stacked-expert tensor shapes;
we read ``E`` and ``N`` from them so the JSON key matches.

``configs`` accepts an already-loaded tuned table (or ``None`` for
"no table shipped"). Callers reached from inside a ``fullgraph``
torch.compile region MUST pass one: loading it here reads the device
name and a JSON file, neither of which Dynamo can trace, and a
fullgraph region cannot graph-break around them.

Launch :func:`fused_moe_kernel`.

Same launch math as vLLM's ``invoke_fused_moe_triton_kernel`` —
we just take a torch dtype for ``compute_type`` instead of the
Triton dtype so callers don't import ``triton.language`` at all.

``A`` is ``(num_tokens, K)`` (bf16 / fp16 / fp8_e4m3); ``B`` is
``(num_experts, N, K)`` packed; ``C`` is the expanded output
``(num_tokens, top_k, N)``. The caller is responsible for
finalize (sum across the ``top_k`` axis, weighted by router
weights).

Upper bound on the frame :func:`fused_moe_kernel_gptq_awq` claims.

Per K-tile the loop body loads ``BLOCK_SIZE_K // GROUP_ROWS`` chunks of
``[BLOCK_SIZE_M, GROUP_ROWS]`` activations (``elem_bytes`` each),
``[GROUP_ROWS, BLOCK_SIZE_N]`` packed weights (one byte per K row) and
one ``[1, BLOCK_SIZE_N]`` scale (``elem_bytes``) plus zero-point (one
byte) row, where ``GROUP_ROWS = min(group_size, BLOCK_SIZE_K)``. The
bound assumes every one of those is multi-buffered ``num_stages`` deep,
taken against the epilogue's ``[BLOCK_SIZE_M, BLOCK_SIZE_N]`` store
staging.

Triton packs by liveness and multi-buffers only what its pipeliner
hoists, so the real frame is smaller — 1.0-3.4x smaller over the tile
grid, never larger, which
``test_awq_shared_memory_bound_holds_against_the_compiler`` pins
against the compiler. An upper bound is what this needs to be: it is
consulted from inside the fullgraph-compiled decoder block, so it has
to be plain integer arithmetic rather than a probe of the compiler,
and erring high can only cost a smaller tile, never a refused launch.

Largest K-tile :func:`fused_moe_kernel_gptq_awq` accepts here.

Three hard constraints, applied by halving:

* ``BLOCK_SIZE_K`` and ``group_size`` must divide one another — the
  kernel walks its K-tile in ``min(group_size, BLOCK_SIZE_K)``-row
  chunks and gathers one scale / zero-point row per chunk. Both are
  powers of two, so this holds for every candidate and is asserted
  rather than searched.
* ``K % BLOCK_SIZE_K == 0`` — the B tile is loaded unmasked, so a tile
  that runs past K reads into the next expert.
* :func:`_awq_shared_memory_bytes` must fit ``budget``.

``_MIN_AWQ_DOT_K`` is the floor: below it ``tl.dot`` rejects a 16-bit
operand, and it does so as a compilation error raised from inside the
kernel where the checkpoint's group size is no longer visible.

Launch :func:`fused_moe_kernel_gptq_awq`.

Same as :func:`invoke_fused_moe_kernel` but for the AWQ / GPTQ
weight-only quant path: the kernel dequantizes int4/int8 in-tile
using ``B_scale`` (per-group) and optionally ``B_zp`` (zero-point).

``B`` is ``(E, N, K // 2)`` uint8 (two int4 per byte along K, LOW
nibble = even K), ``B_scale`` is ``(E, N, K // group_size)`` and
``B_zp`` is ``(E, N // 2, K // group_size)`` uint8 (two zero-points
per byte along N, LOW nibble = even N). ``C`` is
``(num_tokens, top_k, N)``.

``config["BLOCK_SIZE_K"]`` is a request, not a pin:
:func:`_resolve_awq_block_k` lowers it until it divides ``K`` and
:func:`_awq_shared_memory_bytes` fits the device's budget. It is NOT
bounded by ``group_size`` — the kernel walks a K-tile wider than the
group as several group-sized chunks.

NOTE: per the research doc's risk #2, this Triton AWQ path is
correct-but-not-optimal on Ada — Marlin would be ~15-30% faster
but lives in a separate vLLM CUDA op we have not extracted. Use
this path only as a fallback or for correctness validation.

Launch :func:`fused_moe_kernel_mxfp4`.

``A`` is ``(num_tokens, K)`` bf16/fp16; ``B`` is ``(num_experts, N,
K // 2)`` uint8 packed E2M1; ``B_scale`` is ``(num_experts, N,
K // group_size)`` uint8 E8M0; ``B_bias`` is an optional
``(num_experts, N)`` per-expert bias added after dequantization and
before the router-weight multiply. ``C`` is ``(num_tokens, top_k, N)``.

Launch :func:`fused_moe_kernel_nvfp4`.

``A`` is ``(num_tokens, K)`` bf16/fp16; ``B`` is ``(num_experts, N,
K // 2)`` uint8 packed E2M1; ``B_scale`` is ``(num_experts, N,
K // group_size)`` fp8_e4m3; ``B_gscale`` is ``(num_experts, G)``
fp32 holding the checkpoint's per-tensor ``weight_scale_2``, with
``N // G`` consecutive output rows per entry (``G == 2`` for a
``w13`` stack whose gate and up halves came from two linears, 1
otherwise). ``B_bias`` is an optional ``(num_experts, N)`` per-expert
bias added after dequantization and before the router-weight
multiply. ``C`` is ``(num_tokens, top_k, N)``.

Plain-torch reference for the fused-MoE kernel contract.

Triton has no CPU backend, so without this the whole MoE arch — router
semantics, expert binning, expert-parallel id remap, block-fp8 scale
convention — would only ever be executable on a GPU host, and the CPU
parity suites that gate this repo's model numerics could not run.

:func:`moe_align_block_size_reference` is the executable spec for the
token-to-expert binning: the ordering it produces (stable within a bin),
the static worst-case buffer size, and the ``-1`` non-local sentinel are
what the Triton binning kernels must reproduce BIT-IDENTICALLY.

:func:`invoke_fused_moe_reference` implements EXACTLY the contract of
:func:`arbi_serve.kernels.fused_moe.fused_moe_kernel` — same
``sorted_token_ids`` / ``expert_ids`` binning, same ``-1`` write-zeros
sentinel, same per-K-block dequantization, same in-kernel router-weight
multiply — in obvious torch ops. It is the executable spec, NOT a
performance path: :func:`invoke_fused_moe_kernel` dispatches here only
when the operands are off-device, and a GPU parity test asserts the
Triton kernel agrees with it.

Static worst-case length of ``sorted_token_ids``.

Each occupied bin wastes at most ``block_size - 1`` pad slots, and at
most one bin per routed slot can be occupied — hence
``num_valid + min(n_bins, num_valid) * (block_size - 1)``, rounded up
to a whole number of blocks. Sizing from the ACTUAL per-bin counts
would need those counts on the host (one device sync per MoE layer
per step) and would give the decode graph a data-dependent shape.

Off-device twin of
:func:`arbi_serve.kernels.fused_moe.invoker.moe_align_block_size`.

Bin 0 collects the non-local (``id < 0``) slots; bin ``e + 1``
collects expert ``e``. Folding "not mine" into an ordinary bin keeps
the binning branch-free — the ``- 1`` when reading expert ids back
turns bin 0 into the kernel's ``-1`` write-zeros sentinel.

Every intermediate stays on device (sort → searchsorted → cumsum →
scatter) so this path is sync-free and capture-safe in its own right;
it is the executable spec the Triton binning is pinned to, and the
only implementation on a CPU host.

Fill ``C`` exactly as the Triton kernel would. See module docstring.

``B_bias`` is the optional ``(num_experts, N)`` per-expert bias, added
after dequantization and BEFORE the router-weight multiply — the same
order the kernel's ``HAS_BIAS`` branch uses.

``use_mxfp4`` reads ``B`` as ``(E, N, K // 2)`` packed E2M1 with
``B_scale`` ``(E, N, K // 32)`` E8M0. ``use_nvfp4`` reads the same
packed ``B`` with ``B_scale`` ``(E, N, K // 16)`` fp8_e4m3 and
``B_gscale`` ``(E, G)`` fp32, one entry per ``N // G`` output rows.

Fused single-launch GDN gated-delta-rule packed decode kernel.

Computes one recurrent decode step over a packed ``[Q‖K‖V]`` input,
fusing the gate (``g = -exp(A_log) * softplus(a + dt_bias)``,
``beta = sigmoid(b)``), the delta-rule recurrence, and the in-place slab
gather/scatter (via ``ssm_state_indices``) into a single Triton launch.
Vendored from vLLM's FLA ops; see the header for the upstream attribution
and the arbi-serve slab-layout / call-surface adaptations.

The ``ARBI_GDN_DECODE_NUM_WARPS`` flag, read as a per-process constant.

Isolated into a 0-arg ``assume_constant_result`` helper so a Dynamo trace
that reaches :func:`_resolve_decode_num_warps` bakes the flag instead of
tracing ``RuntimeFlags.from_env`` (``dataclasses.fields`` is untraceable on
torch-2.12 when the runtime-flags cache is disabled).

Pick the Triton ``num_warps`` for the packed-decode kernel.

The kernel grid is ``(NV, B * HV)`` and each program processes one
``[BV, BK]`` tile (BV rows of the V axis, BK = the full K axis).
Upstream vLLM hard-codes ``num_warps=1`` — tuned for the B=1 latency
regime, where the grid is tiny (``NV * HV`` programs) and a single
warp per program keeps launch + occupancy overhead minimal.

At high batch the grid is large (``NV * B * HV`` programs) but every
program is still a single warp, so each SM is under-occupied relative
to the available register/lane budget. Giving the program more warps
lets Triton spread the ``[BV, BK]`` tile across more lanes.

Not bit-exact: raising ``num_warps`` changes the tile-to-lane layout,
which reassociates the two ``tl.sum(..., 1)`` K-axis reductions
across the warp split. The result is deterministic but not
byte-identical to ``num_warps=1`` at fp32. Because GDN state is
recurrent and compounds over decode steps, this drift accumulates —
so the heuristic is opt-in and the default stays at the upstream
``num_warps=1`` (bit-exact).

Resolution order:
  * ``ARBI_GDN_DECODE_NUM_WARPS`` env, if set:
      - an integer (1/2/4/8) forces that value;
      - ``"auto"`` selects the batch-adaptive heuristic below.
  * otherwise ``default`` (1 — the upstream value) is returned, so
    the heuristic is OPT-IN until the win is validated on the target
    deployment.

The heuristic scales warps with the per-grid work. ``BV`` is capped
at 32, so a single warp (32 lanes) already covers exactly one BV row;
more warps help when there are multiple BV blocks AND a wide batch to
fill the SMs. We cap at 8 (256 lanes) — beyond that the [BV<=32, BK]
tile has no more parallelism to exploit and extra warps only add
scheduling overhead.

Live-flip note: this function itself re-reads ``runtime_flags()`` on
every call (no caching in this function) — the resolved value is
always current AT THE TIME THIS FUNCTION RUNS. What is NOT guaranteed
is that this function runs again after a live admin-API flip: on the
canonical cudagraph-captured decode, the caller
(:func:`fused_recurrent_gated_delta_rule_packed_decode`) is invoked
once per (member, shape-bucket) at capture time, and every subsequent
decode step of that shape replays the captured ``torch.cuda.CUDAGraph``
directly — no Python, so no call into this function at all. This is
why ``gdn_decode_num_warps`` is ``scope="backend"`` (capture-affecting)
in :mod:`arbi_serve.config_overrides`: a live flip must rebuild+
recapture the member to actually reach a fresh call here, exactly like
its sibling ``gdn_decode_kernel``.

vLLM's GDN packed-decode kernel, adapted to arbi-serve's slab.

Args:
    mixed_qkv: (B, q_dim + k_dim + v_dim) post-conv packed
        ``[Q || K || V]`` tensor, contiguous in the last dim.
    a, b: (B, HV) raw projections (BEFORE softplus/sigmoid) — the
        kernel will compute ``g = -exp(A_log) * softplus(a + dt_bias)``
        and ``beta = sigmoid(b)`` itself.
    A_log, dt_bias: (HV,) per-head decay + bias parameters.
    scale: ``1/sqrt(K)`` (callers are expected to compute this once).
    initial_state: 4-D recurrent slab. Updated IN-PLACE at row
        ``ssm_state_indices[i]`` for each i in [0, B). Shape is
        ``(N, HV, V, K)`` when ``transpose_state=True`` (the
        arbi-serve production layout, set by ``RecurrentStatePool``)
        and ``(N, HV, K, V)`` when ``transpose_state=False``.
    out: (B, 1, HV, V) preallocated output buffer (contiguous).
    ssm_state_indices: (B,) int64 row indices into ``initial_state``.
        A negative entry is treated as NULL — the corresponding
        output row is zero-written and the slab is left alone.
    use_qk_l2norm_in_kernel: matches FLA semantics — L2-normalize
        Q and K inside the kernel before scaling by ``scale``.
    transpose_state: if True, treat ``initial_state`` as
        ``(N, HV, V, K)`` (V outer, K inner — production arbi-serve
        layout, also vLLM upstream). If False, treat it as
        ``(N, HV, K, V)`` (the legacy FLA default layout). Default
        False; production callers in :mod:`arbi_serve.models._gdn_kernels`
        pass ``True``.
    snapshot_out: optional MTP rollback snapshot slice
        ``snap_rec_full[0]``, SAME shape/layout as ``initial_state``.
        When given, the kernel writes the just-computed final state
        into it IN THE SAME LAUNCH (at rows ``ssm_state_indices``) —
        bit-identical to a post-kernel ``slab.index_select`` +
        ``snap.index_copy_`` but without the separate slab-sized fp32
        copy that dominates the DFlash seed-decode step. ``None``
        (default) skips the snapshot write (non-MTP decode).
    num_warps: Triton launch ``num_warps``. ``None`` (default) defers
        to :func:`_resolve_decode_num_warps` (env-driven; falls back
        to the upstream value 1 unless ``ARBI_GDN_DECODE_NUM_WARPS``
        is set). An explicit int forces that value — used by the
        micro-benchmark to sweep {1,2,4,8}. Changing ``num_warps``
        perturbs numerics (the K-axis reductions reassociate
        across the warp split) — see the resolver docstring; this is
        why the default stays at the bit-exact upstream value 1.
    conv_input_v: Optional pre-convolution V stream ``(B, HV*V)``.
        Supplying it together with ``conv_state`` and ``conv_weight``
        computes the V short convolution inside the recurrence launch.
        ``mixed_qkv`` then contains only post-convolution Q and K.
    conv_state: Complete indexed convolution slab ``(N, C, W)``.
    conv_weight: Complete depthwise convolution weights ``(C, W)``.
    conv_snapshot_out: Optional rollback snapshot with the same shape
        and layout as ``conv_state``. The V-state owner writes its
        updated channels in the recurrence launch.

Returns:
    ``(out, initial_state)`` — both are the input tensors (in-place
    updates), returned for API symmetry with the FLA call.

Run the packed recurrence with its V short convolution in-launch.

Q and K share convolution state across multiple recurrence programs, so
they are advanced by the indexed short-conv kernel before the recurrence
launch. Each V channel has exactly one recurrence-program owner and can be
advanced safely inside that launch. The split removes the global V
convolution output while retaining deterministic state ownership.

Fused sigmoid-gating gated-delta-rule update Triton kernel.

Fuses the GDN decode step into one launch: gathers the recurrent state
from a slab (via ``ssm_state_indices``), computes the gates in-kernel
(``g = -exp(A_log) * softplus(a + dt_bias)``, ``beta = sigmoid(b)``), runs
the per-token gated-delta-rule recurrence, and scatters the final state
back into the slab in place. Vendored from flash-linear-attention via
vLLM; see the header for upstream attribution.

Resolve a ``(B, T, H, D)`` stream to ``(tensor, token pitch)``.

The kernel reads each token's ``(H, D)`` block contiguously and steps
tokens by a pitch argument, so a view whose only non-contiguity is that
pitch — a head-group slice of a wider ``(T, C)`` activation, which is
what the conv split hands the verify launch — is consumed in place.
Any other layout is copied once here. This is the single owner of the
kernel's layout contract: callers pass their views and never add a
defensive ``.contiguous()`` (each one is an eager launch inside the
opaque custom op, invisible to the compiled region around it).

Resolve a ``(T, HV)`` gate stream to ``(tensor, token pitch)``.

A column half of the fused ``[B|A]`` projection has a unit inner stride
and a row pitch of ``2 * HV``; the kernel reads it in place. See
:func:`_token_strided_heads` for why the copy is owned here.

Fused triton implementation of sigmoid gating delta rule update.
This function uses a single fused kernel that combines both sigmoid gating
computation and the recurrent delta rule update for better performance.

GDN MTP masked-replay extensions (arbi-serve, snapshot-ladder cut):

* ``save_inputs_kvba=(saved_k, saved_v, saved_b, saved_a)`` — persist the
  RAW per-token activation-dtype inputs at slot ``row * T_save + t`` of
  the ``(N_rows, T_save, …)`` replay buffers (``row`` = the per-seq slab
  index from the 1-D ``ssm_state_indices``). Free relative to the ladder
  it replaces: ~(H·K + HV·V + 2·HV) activation-dtype elements/token vs
  the ladder's HV·V·K fp32 state/token.
* ``save_base_state_out`` — persist the EFFECTIVE fp32 h0 (post
  ``init_state_mask`` / post ``round_h0_to_act_dtype``) into the 1×
  committed BASE frame ``base[row]``.
* ``final_state_last_only=True`` — store only the post-token-(T-1) state
  at ``ht[row]`` (no per-token ladder stores).
* ``n_steps`` — per-SEQUENCE trip count, replacing the launch-wide depth
  ``cu_seqlens`` implies. The masked tail steps are exact identities on
  the loop-carried state, so stopping at the accepted length reaches the
  same state without executing them, and the caller can drop the masking
  writes that made them identities. ``final_state_last_only`` follows it:
  the stored state is the one after step ``n_steps - 1``.
* ``commit_final_state=False`` — run the recurrence for its outputs and
  store NO final state anywhere. The slab passed as ``initial_state`` is
  left holding the pre-recurrence state, which is what lets it serve as
  the rollback base with no separate base frame.

The masked replay is not a kernel mode: it re-invokes this same
call signature (hence the same compiled binary — the property that makes
the replayed prefix byte-identical) over the saved buffers with steps
``t > n_accepted`` masked to the ×1.0/+0.0 identity at the inputs
(``k → 0`` kills the rank-1 update and the v-correction; ``a → -inf``
makes ``softplus → 0`` so the decay is ``exp(-0.0) = 1.0`` exactly),
and rows with nothing to commit routed to the zero-sentinel slot via
``ssm_state_indices`` (the kernel's ``state_idx <= 0`` skip). The
save-* stores during a replay launch write back the very bytes just
loaded — benign self-writes.

All three require the continuous-batching h0 path (1-D
``ssm_state_indices`` + ``initial_state``) and ``is_kda=False``.

Single-token Triton causal-conv update for GDN decode.

The kernel gathers one persistent conv-state row per request, shifts the
rolling window, appends the current input, evaluates the depthwise convolution
in fp32, applies SiLU, and scatters the updated window in one launch.

The launcher deliberately accepts a strided ``(batch, channels)`` input.  The
production GDN projection exposes ``[Q | K | V]`` as a narrow view of a wider
``[Q | K | V | Z]`` allocation, so requiring row-major contiguity would add a
copy to the hot path.

Run one indexed GDN causal-conv decode update.

Args:
    x: ``(batch, channels)`` activation tensor. Only the channel stride
        must be one; a wider parent allocation is supported.
    conv_state: Persistent ``(num_slots, channels, kernel_width)`` slab,
        mutated in place.
    weight: Depthwise weights ``(channels, kernel_width)``.
    state_indices: ``(batch,)`` int32/int64 slab-row indices.
    snapshot_out: Optional rollback snapshot slab with the same shape as
        ``conv_state``. Selected updated rows are written in the same
        kernel launch as ``conv_state``.
    coalesced_state: Select complete-state tiled transactions. ``None``
        selects the batch-adaptive production specialization.

Returns:
    A contiguous ``(batch, channels)`` tensor in ``x.dtype``.

Return the kernel that will ACTUALLY run for this GDN geometry.

``runtime_flags().gdn_decode_kernel`` is a REQUEST, not the outcome. The
head-parallel ablation kernels are specialized to
``PARALLEL_DECODE_LOCAL_SHAPE`` (the 0.8B shape); every other geometry —
the served 27B included — downgrades to the shape-general ``baseline``.

Exported so a caller that must branch on the kernel choice BEFORE it has
operands to validate reads the SAME rule the dispatcher applies, instead
of comparing the raw flag string and being silently overridden a step
later. Idempotent: resolving an already-resolved name is a no-op.

Launch the selected implementation for one captured decode shape.

``baseline_decode`` lets the caller hand in the shape-general
``fused_conv_recurrent_gated_delta_rule_packed_decode`` entry point to
use for the ``baseline`` kernel (including the ablation-shape
fall-through). The GDN block passes its
``_gdn_kernels.fused_conv_recurrent_gated_delta_rule_packed_decode``
module attribute so the single-namespace monkeypatch contract (see the
``_gdn_kernels`` module docstring) is honoured — a fresh ``import`` here
would shadow the tests' spy. Left ``None`` on the
compiled custom-op path (no monkeypatch there); the direct import below
is the fallback.

Return whether the head-parallel ablation kernels accept this geometry.

The ONE downgrade rule. Both the dispatcher (which validates full
operands) and any caller that must know the kernel choice BEFORE it has
operands to validate (``GDNBlock._forward_decode_packed`` sizes its conv
path off it) read it here, so the two can never disagree.

Two-warp, two-half GDN decode research kernel for high batch.

The first launch advances Q/K convolution state without materializing Q/K
activations.  Two CTAs per ``(batch, key head)`` then reconstruct Q/K from the
updated rolling state, and each CTA advances one 64-wide V/recurrent half with
two one-warp K-reduction tiles.  Q/K state therefore has one writer while the
recurrence has twice the independently schedulable CTA granularity of the
four-warp head kernel.

The callable is isolated from production dispatch and specialized to the
0.8B local shape.

Dispatch-free head-parallel GDN decode research kernel.

One four-warp CTA owns a ``(batch row, key head)`` pair.  Q/K convolution
channels are partitioned across its warps and advanced exactly once.  The
normalized vectors are staged through shared memory, then each warp advances
one disjoint 32-wide value/recurrent tile.  Each tile retains the production
one-warp reduction order across K.

The callable is intentionally isolated from production dispatch.  Gluon is
lazy-imported so importing this module remains safe on CPU-only hosts.

``causal_conv1d_fn``-backed depthwise causal-conv1d + SiLU for the GDN
prefill path — an A/B alternative to the in-tree varlen Triton kernel
(:mod:`arbi_serve.kernels.gdn_prefill_conv_varlen`).

WHY. arbi already depends on Dao-AILab's tuned CUDA ``causal_conv1d``
extension for the short-conv block and for the GDN *decode* path
(``causal_conv1d_update``). This module routes the GDN *prefill* conv
through the same tuned ``causal_conv1d_fn`` kernel so we can A/B the two
on GPU.

GATING. Behind ``ARBI_GDN_CONV_CAUSAL_CONV1D=1`` (default OFF → the
Triton path stays canonical). When ON we still fall back to Triton for
the rare *non-uniform* per-row-T batch (chunked-prefill admission
boundary) — ``causal_conv1d_fn`` only takes ``initial_states`` /
``final_states`` in its *batched* (uniform-seqlen) mode; its varlen
(``seq_idx``) mode in the installed Dao-AILab wheel forbids
``initial_states``, so it cannot fold the prior conv state into a
single packed launch.

I/O CONTRACT (identical to the Triton kernel — see
``gdn_prefill_conv_varlen``):

  Inputs:
    x_conv     : (T_total, C)  flat token-major conv input.
    prior_conv : (B, C, K)     K-wide prior conv state, already
                               ``has_initial_state``-masked to zero by
                               the caller for fresh-prefill rows.
    w          : (C, K)        depthwise weights.
    cu_seqlens : (B+1,) int32  per-row token boundaries.
  Outputs (written into caller-allocated buffers):
    out        : (T_total, C)  SiLU-activated conv output.
    final_buf  : (B, C, K)     K-wide final rolling buffer the decode
                               path's ``causal_conv1d_update`` consumes.

ARG MAPPING onto ``causal_conv1d_fn`` (batched / uniform branch):

  * ``x``               := x_conv reshaped (T_total,C) → (B,T,C) →
                           (B,C,T).
  * ``weight``          := w  (C,K) — Dao's batched ``causal_conv1d_ref``
                           does ``F.conv1d(x_aug, weight.unsqueeze(1),
                           padding=0, groups=dim)`` (cross-correlation, no
                           kernel flip) over ``x_aug = cat(initial_states,
                           x)`` — bit-for-bit the same op + weight ordering
                           as arbi's ``F.conv1d`` reference, so no flip is
                           needed.
  * ``initial_states``  := prior_conv[..., 1:K]  (B,C,K-1) — arbi's
                           reference seeds history from the K-1 most
                           recent prior columns (drops the oldest column
                           ``prior[..., 0]``); Dao's ``initial_states`` is
                           exactly ``(B, dim, width-1)``.
  * ``activation``      := "silu".
  * ``return_final_states`` := False — we do not use Dao's K-1-wide final
                           state. We reconstruct arbi's full K-wide
                           ``final_buf`` ourselves from the same
                           ``x_aug = cat(prior[...,1:K], x_ct)`` slice the
                           reference uses (``x_aug[..., -K:]``). This is
                           pure data movement (gather/slice, no
                           arithmetic) → exactly equal to the reference,
                           including column 0 (``x_aug[..., -K]``) which
                           Dao's ``(B,C,K-1)`` final state omits. (Column
                           0 is dropped on the first decode rolling-buffer
                           shift and never enters a conv output, but we
                           keep the slab bit-identical regardless.)

CUDA-only (``causal_conv1d_fn`` is a CUDA extension). The caller gates
this on ``x_conv.is_cuda`` and ``_HAVE_CAUSAL_CONV1D``.

Return the common per-row token count if the packed batch is
uniform AND non-empty, else ``None``.

Reads ``cu_seqlens`` on the HOST. This is acceptable here: the
``causal_conv1d_fn`` path is an opt-in A/B branch (default OFF) and is
not used under the B==1 ``ARBI_GDN_PREFILL_CAPTURE`` stream-capture
region — the canonical Triton kernel (capture-safe, no host sync)
stays default. The uniform check needs the per-row lengths regardless;
a host read is the cheapest correct way to make the uniform/non-uniform
dispatch decision the Triton op makes internally.

Run the GDN prefill depthwise causal-conv1d + SiLU via Dao-AILab's
tuned ``causal_conv1d_fn`` (batched / uniform branch).

Writes into the caller's pre-allocated ``out`` (conv output) and
``final_buf`` (K-wide rolling state). Returns ``True`` when the
``causal_conv1d_fn`` path handled the batch, ``False`` when it
declined (non-uniform per-row T, K==1, or extension absent) so the
caller can fall back to the Triton / reference path.

Math is bit-identical (fp round-off only) to the per-row
``F.silu(F.conv1d(...))`` reference; ``final_buf`` is exactly equal
(pure data movement).

Varlen depthwise causal-conv1d + SiLU for the GDN prefill path.

This module provides one varlen Triton launch over the flat
``(T_total, C)`` packed input keyed by the device ``cu_seqlens``
(int32, ``(B+1,)``) — no host sync, no per-row Python iteration.

Two kernels live here:

  * ``_kern_tiled`` — a 2-D chunked grid. The grid is
    ``(n_chunks, cdiv(C, BLOCK_N))``; ``program_id(0)`` indexes a flat
    list of ``(seq, chunk_offset)`` pairs (the *chunk map*), each chunk
    covering ``BLOCK_M`` tokens. Each program loads its ``BLOCK_M``-token
    chunk plus the ``K-1`` left-context as coalesced ``(BLOCK_M, BLOCK_N)``
    tiles, then the depthwise conv is ``K`` elementwise FMAs over the tile
    (no serial per-token loop). This grid parallelizes across both chunks
    and channels, so it stays parallel at low batch / long context where
    the serial-time kernel below would only launch one program per
    sequence. Specialized for ``K == 4`` (Qwen3.5 GDN).

  * ``_gdn_prefill_conv_varlen_kernel`` — the serial-time kernel, kept as
    the fallback for ``K != 4`` (one program walks a whole sequence in
    time, holding the rolling window in registers).

The ``(seq, chunk_offset)`` chunk map is built on device from
``cu_seqlens`` by ``_kern_build_chunk_map`` — no host sync, so the whole
path is cudagraph-capture-safe (relevant under
``ARBI_GDN_PREFILL_CAPTURE=1``). The map grid is bounded by
``n_chunks_max = cdiv(T_total, BLOCK_M) + B`` (both ``T_total`` and ``B``
are static tensor shapes at call time, so the launch grid is fixed at
capture time). Map slots past a sequence's real chunk count are filled
with a sentinel ``coff`` large enough that ``tok_start >= T_i`` → the
conv program stores nothing.

Numerics are bit-identical (fp round-off only) to the per-row
``F.silu(F.conv1d(...))`` reference, on both the conv output and the
K-wide final rolling buffer the decode ``causal_conv1d_update`` consumes.
The history convention matches the reference exactly:

  * The GDN slab stores a K-wide conv state ``prior_conv[:, :, 0:K]``;
    the reference effective history is ``prior_conv[:, :, 1:K]`` (K-1
    samples). A within-sequence position ``p < 0`` maps to prior column
    ``p + K`` (so ``p = -1`` → col ``K-1`` = newest history). The caller
    has already ``has_initial_state``-masked ``prior`` to zero for fresh
    prefills.

  * ``F.conv1d`` is cross-correlation (no kernel flip):
    ``out[c, t] = sum_k w[c, k] * x_aug[c, t - (K-1) + k]`` where
    ``x_aug = concat(history, x)``. The tiled kernel computes the same
    dot product as ``K`` FMAs over shifted tiles.

  * The K-wide final state is the last K columns of ``x_aug``
    (``x_aug[..., -K:]``) — written only by the program covering each
    sequence's last chunk, reconstructed by direct loads (within-seq
    positions ``T_i - K + k`` for ``k in [0, K)``; ``p < 0`` pulls prior
    column ``p + (K-1)``).

CUDA-only. The caller falls back to the ``F.conv1d`` reference path on
CPU / when Triton is unavailable.

One program per flat chunk slot ``p``.

Recomputes the per-seq chunk-count prefix sum (cheap: ``B`` small,
<= 64) and locates ``p`` within it, deriving ``(seq, chunk_offset)``.
Beyond the last real chunk → sentinel.

``B`` and ``sentinel_coff`` are ``do_not_specialize`` runtime scalars:
``sentinel_coff`` tracks ``n_chunks_max = cdiv(T_total, BLOCK_M) + B``,
which varies with every distinct ``(T_total, B)`` — as a constexpr (or
a value-specialized int) it is an unbounded specialization axis that
recompiles this kernel repeatedly under live traffic (caught by the
post-ready JIT detector). As plain scalars the compiled-variant set
collapses to the bounded ``B_PAD`` power-of-two ladder, which boot
warmup pre-compiles (:func:`warm_chunk_map_kernel`).

Build the (seq_of_chunk, coff_of_chunk) flat maps on device.

Returns ``(seq_map, coff_map, n_chunks_max)``. ``n_chunks_max`` is a
static upper bound — ``cdiv(T_total, BLOCK_M) + B`` — so the conv
launch grid is fixed at cudagraph-capture time. Slots past a
sequence's real chunks carry a sentinel ``coff`` (== n_chunks_max,
guaranteeing ``tok_start >= T_i`` → the program stores nothing).

Pre-compile ``_kern_build_chunk_map`` for every served ``B_PAD``.

``B_PAD`` (power-of-two ``>= B``) is the kernel's only remaining
specialization axis (``B`` / ``sentinel_coff`` are do_not_specialize
scalars), bounded by the served ``max_batch`` — so the whole variant
set is pre-compilable at boot. Launches one tiny map build per rung;
returns the number of rungs warmed. A post-ready compile of this
kernel is a warmup bug the JIT detector flags.

Load a (BLOCK_M, BLOCK_N) tile at within-seq positions ``pos``.

``pos >= 0`` → raw ``x``; ``pos < 0`` → prior conv state column
``pos + K`` (so ``pos = -1`` → col ``K-1`` = newest history).

Launch the varlen depthwise causal-conv1d + SiLU kernel.

One conv launch covers all ``B`` sequences in the packed batch — no
per-row Python loop, no ``cu_seqlens`` host sync. For ``K == 4`` (the
Qwen3.5 GDN width) it uses the 2D chunked-grid ``_kern_tiled``
(saturates the SMs at low concurrency / long context); the chunk map
is built on device first (one tiny launch, no host sync). For other
``K`` it falls back to the serial-time kernel. Writes into the
caller's pre-allocated ``out`` and ``final_buf``. Math is bit-identical
(fp round-off) to the per-row ``F.silu(F.conv1d(...))`` reference on
both the conv output and the K-wide final state.

Vendored / inference-only forward path for mamba-ssm's SSD chunk scan.

Provenance
----------
- Upstream: ``mamba-ssm`` 2.3.2.post1 (the wheel built into the arbi-serve
  server container), ``mamba_ssm/ops/triton/ssd_combined.py`` and
  ``mamba_ssm/ops/triton/ssd_state_passing.py``.
- Upstream Apache-2.0 license — see :file:`LICENSE.upstream`.

Why vendor
----------
The savepoint fold emit needs the SSD scan to store its running fp32
recurrent state at a chosen chunk boundary
(:mod:`arbi_serve.cache._fold_emit_staging` states the seam). Upstream
materialises a per-chunk state — ``_state_passing_fwd(..., out_dtype=C.dtype)``
— but in the ACTIVATION dtype, because the chunk-scan kernel consumes it as
a matmul operand; every boundary reachable from outside the kernel is
therefore a rounded copy of a register that was fp32 the whole way. One
extra store inside that kernel is the whole change
(:mod:`~arbi_serve.kernels.mamba_vendored.ssd_state_passing_emit`).

Reaching that kernel means reaching past ``mamba_chunk_scan_combined``,
which is a ``torch.autograd.Function`` that returns only ``out`` and
``final_states``: the per-chunk states are its private intermediate. So the
four-kernel orchestration is vendored too
(:mod:`~arbi_serve.kernels.mamba_vendored.ssd_combined`) — forward-only, with
three of the four kernels still imported from the installed wheel so
upstream fixes to them flow through.

Inference-only forward path for ``mamba_chunk_scan_combined``, with the emit.

Vendored equivalent of upstream's public entry point
(:func:`mamba_ssm.ops.triton.ssd_combined.mamba_chunk_scan_combined`),
forward-only and without the ``torch.autograd.Function`` indirection, so
that ONE of the four kernels in the chain can be ours
(:mod:`~arbi_serve.kernels.mamba_vendored.ssd_state_passing_emit`) while the
other three stay the installed wheel's and follow its fixes.

Upstream's ``MambaChunkScanCombinedFn.forward`` orchestrates::

    dA_cumsum, dt = _chunk_cumsum_fwd(dt, A, chunk_size, ...)
    states = _chunk_state_fwd(B, x, dt, dA_cumsum, states_in_fp32=True)
    states, final_states = _state_passing_fwd(..., out_dtype=C.dtype)
    CB = _bmm_chunk_fwd(C, B, chunk_size, output_dtype=torch.float32)
    out, out_x = _chunk_scan_fwd(CB, x, dt, dA_cumsum, C, states, D=D, z=z)

Only the third line changes: the state passing gains the emit table, and
nothing else about the chain — order, dtypes, blocking — moves, which is
what makes an emitting launch's ``out`` and ``final_states`` the plain
launch's own bytes. The autograd context this drops is dead weight under
:func:`torch.inference_mode`, and the frames it drops are the same ones the
GDN vendoring dropped for the same reason.

WHAT THIS DOES NOT VENDOR
=========================
The backward pass, the ``cu_seqlens`` varlen-states return, and ``z``
(the block applies its gate through its own gated RMSNorm). arbi's Mamba-2
prefill calls the kernel once per row with ``cu_seqlens=None`` and
``z=None`` (:meth:`~arbi_serve.models.mamba2_block.Mamba2Block._mamba2_kernels`),
so a vendored path that accepted them would be untested surface. They are
refused by name rather than silently ignored.

``(out, final_states)`` — upstream's forward with the savepoint emit.

``emit_chunk`` / ``emit_out`` are the state-passing emit's table and
buffer; see :func:`~arbi_serve.kernels.mamba_vendored.ssd_state_passing_emit.state_passing_fwd_emit`
for the contract. Passing them costs an unarmed launch one extra
predicated compare per chunk and no store, so a boot that CAN emit
hands them to every launch and the traced graph is the same with and
without a savepoint.

``_state_passing_fwd`` with an in-kernel state EMIT.

The SSD scan's state passing carries a sequence's SSM state in fp32
registers across chunks. It stores that state to ``out`` at every chunk
start and to ``final_states`` once, at the end. ``out`` is materialised in
``out_dtype`` — the caller passes the ACTIVATION dtype (``C.dtype``) because
the chunk-scan kernel reads it as a matmul operand — so every boundary a
caller could read out of ``out`` is a rounded copy of the register. This
variant adds ONE more store: at the start of chunk ``emit_chunk[seq]``
(``-1`` = none) the running fp32 state is written to ``emit_out``, ONE
``(H, P, N)`` buffer in the SLAB's own layout, so a restore writes those
bytes straight back into a slab row with nothing between the emit and the
host copy.

Layout is why the store is a single line here and a four-way unrolled block
in the GDN twin: ``states`` is already flat over ``dim = P * N`` and the
Mamba slab row is ``(nheads, head_dim, state_dim)`` contiguous, which is the
same address arithmetic. There is no transpose to fold into the pointer.

One buffer, not one per sequence: a step arms at most one row
(``Scheduler._arm_fold_split``), the buffer is a boot allocation on the
ledger, and sizing it per row would price ``max_batch`` copies of a state
only one row ever writes. Two rows naming a chunk in one launch would race
on it, so the caller guarantees at most one.

That value is, by construction, the ``final_states`` this scan would have
produced had the sequence been cut at that chunk boundary: the recurrence
``states = exp(dA_cs[c]) * states + states_in[c]`` reads only chunk ``c``'s
own inputs, and chunks are cut at multiples of ``chunk_size``, so chunks
``0..c-1`` of a cut sweep are the same chunks with the same values in the
same order. The fold-grid constraint is the kernel's own addressing (a chunk
index), so an off-grid position has no representation here rather than a
rounded one.

The emit table is data: ``emit_chunk`` is an ``(N,)`` int32 tensor and
``emit_out`` a preallocated buffer, so the same launch serves a step with or
without a savepoint, eager or captured, and nothing the forward reads to
decide "emit here" is a Python value Dynamo could guard on.

Kernel text is upstream's verbatim apart from the emit parameters, the
pointer setup and the store; ``tools/mamba_fold_emit/check_emit_state.py``
measures the bit-identity against the two-segment scan on a card.

Upstream ``_state_passing_fwd`` plus the emit table.

``emit_chunk`` — ``(N,)`` int32, the chunk index (relative to each
sequence's own start) at whose START the state is emitted, ``-1`` for
none; AT MOST ONE entry may be armed (the caller's contract — two armed
rows would race on the one buffer). ``emit_out`` — ONE ``(H, P, N)``
buffer in the slab's layout, written only by the armed row; fp32 is the
lossless contract, any other float dtype receives the kernel's rounding
of the fp32 register, which is the same rounding the slab scatter
applies to ``final_states``.

Fused gated RMSNorm Triton kernel.

Computes ``rmsnorm(x) * weight * silu(z)`` (``norm_before_gate=True``) in a
single forward-only Triton launch, replacing the ~11 elementwise launches
of the pure-PyTorch ``_RMSNormGated``. Vendored from flash-linear-
attention via vLLM; see the header for upstream attribution.

One 2-D tile per program: ROWS_PER_BLOCK rows × BLOCK_N cols.

``norm_before_gate=True``: out = (x * rstd * w) * act(z), where
``act`` is ``silu`` (GDN — Qwen3.5 / 3.6) or ``sigmoid`` when
``GATE_SIGMOID`` (KDA — Ling-3.0's ``FusedRMSNormGated`` with
``activation='sigmoid'``).

Math is fp32 internally regardless of input dtype — matches our
pure-PyTorch reference (which casts to .float() for the variance
reduction) within the asserted atol/rtol.

Gated RMSNorm on the gate stream, ``norm_before_gate=True``.

out = norm(x) * weight * act(z), ``act`` ∈ {``silu``, ``sigmoid``}.

Args:
    x:       ``(..., N)`` activation; promoted to fp32 for the
             variance reduction internally.
    weight:  ``(N,)`` learnable scale (no bias — matches HF
             :class:`Qwen3_5RMSNormGated` semantics).
    z:       ``(..., N)`` gate stream — same shape as ``x``.
    eps:     numerical stability term.
    gate_activation: ``silu`` (GDN) or ``sigmoid`` (KDA — Ling-3.0
             builds its ``FusedRMSNormGated`` with
             ``activation='sigmoid'``).

Returns:
    ``(..., N)`` tensor, same dtype as ``x``.

Raises on shapes the kernel can't currently handle (``N >= 64KB``
or ``z.shape != x.shape``) — no silent fallback.

Early-H2D decoupling: start the warm weight transfer before the graph exists.

The problem
-----------
On the warm flat-dump path the boot spends, MEASURED on a single RTX 4090
booting ``Qwen3.8-27B-exl3-4.0bpw``::


and the ``weights DMA + rebind join`` phase is ~= the DMA duration, i.e. the
transfer is essentially 100% UNHIDDEN. It is also already near the floor: the
read mode, so there is no meaningful bandwidth left to win. The only lever is
OVERLAP.

Why the transfer does not need the model graph
----------------------------------------------
``load_flat_weights`` starts only after ``load_model`` returns because it needs
per-tensor GPU destination views. But that dependency is real for the SCATTER,
not for the TRANSFER: the dump's safetensors header carries every tensor's
name, dtype, shape and byte range, and the data section is a gap-free tiling of
``[0, total)`` (:func:`_assert_flat_dump_valid` proves it at dump time). So the
destination LAYOUT is fully derivable from a ~400 KiB header read, with no
model in hand. This module does exactly that: plan, allocate, and stream —
concurrently with the 1.2-1.9s of pure-CPU model build that already happens.

Why the slabs are the FINAL storage, not a staging buffer
---------------------------------------------------------
The obvious shape — "H2D into a big staging slab, D2D-scatter into the graph
later" — is VRAM-infeasible for the case that motivates it: the staging slab
early-allocated slabs ARE the weights: the post-build step rebinds each
placeholder buffer onto its slab view (exactly what
:func:`~arbi_serve.loader.flat_dump._presize_placeholder_buffers` already does,
minus the separate fill) and copies only the handful of already-correctly-
allocation exists.

Why not a verbatim GPU image of the file
----------------------------------------
A single uint8 slab holding the file's data section byte-for-byte would need no
scatter at all — but the dump packs tensors with zero padding, and MEASURED on
the dump above 1836 of 2456 tensor offsets are not even 16-byte aligned (2418
are not 256-byte aligned), because 802 four-byte EXL3 scalar seeds sit between
the multi-MiB trellis packs. Views into such a slab would be unusable by the
kernels (and ``Tensor.view(dtype)`` rejects most of them outright). So the plan
groups by dtype and pads each slot to :data:`_SLAB_TENSOR_ALIGN_BYTES`, the same
rule the cold path's ``Exl3DirectSlabBinder`` uses. The padding cost on that

Rebinding, and why it is safe HERE
----------------------------------
Every non-tied tensor is REBOUND onto its slab view rather than copied into the
storage the meta-materialize walk built. Copying instead leaves two resident
copies of those bytes — the graph's own allocation and the slab slot that fed
it — and MEASURED on the 27B EXL3 boot that duplication cost 20 KV pages
is the same operation the COLD path already performs on every persistent tensor
at the same boot phase (``stream_compact_per_layer`` re-points ``param.data`` /
``mod._buffers[attr]`` onto its slab views), through the same
:func:`~arbi_serve.loader.flat_loader._common._set_attr` helper, which preserves
Parameter object identity.

The invariant it must not break — captured cudagraphs bake tensor addresses, so
a live arena storage may never be re-pointed — is about POST-capture compaction.
This runs strictly before any capture exists: ``load_model_into_engine`` is
called from ``_load_model_phase``, ``prepare_derived_weights`` (the only
main-thread reader of these buffers) is deferred to after the weights future
joins, and ``_phase_start("cudagraph capture")`` is several boot phases later.
The one thing rebinding CAN break is a tie: two ``state_dict`` names over one
storage (a shared embedding aliased into ``lm_head``) would become two storages.
Tied entries are therefore detected by storage identity and filled in place.

Contract
--------
Otherwise correctness is unchanged, deliberately: the same per-tensor
dtype/shape validation runs (against the same header), the same post-fill
bit-signature gate runs over the same manifest, the same
``rebind_after_compaction`` hooks fire at the same seam, and a real parameter
whose shape disagrees with the dump is still refused loud.

Measured
--------
GMU 0.99. Interleaved A/B, n=3 per arm, medians of ``build + join``::

    OFF   build 5.07s · join 4.2s = 9.27s · 1104 KV pages · late_growth 0
    ON    build 5.59s · join 1.2s = 6.79s · 1104 KV pages · late_growth 0

-2.48s end-to-end at full KV parity. The join wait itself is 0.00s — the
transfer is 100% hidden — but the win is smaller than that delta because the
build absorbs part of the transfer and the transfer runs ~20% slower sharing

Opt-in via ``ARBI_EARLY_H2D`` (default OFF).

Rebind ``mod.<attr>`` to ``view``, preserving Parameter-vs-buffer semantics.

Thin re-export of the compaction helper so both the cold path's slab
compactor and this one re-point a tensor the same way: a Parameter keeps
its object identity (``param.data = view``), a buffer routes through
``setattr`` so the persistent flag survives.

Plan per-dtype slab placement for every tensor in a flat-dump header.

Pure function of the HEADER — no model, no device, no allocation. Returns
``(slots, slab_elems)`` where ``slab_elems[dtype]`` is the element count the
dtype's slab must hold. Tensors are placed in header order (which is
``state_dict()`` order, so the layout is deterministic for a given dump) and
every slot base is padded up to :data:`_SLAB_TENSOR_ALIGN_BYTES`, so every
view handed to a kernel is at least 256-byte aligned.

Fails loud if a tensor's declared byte range disagrees with
``numel * itemsize`` — that would mean the header and the data section
describe different things, and no byte should move until it is resolved.

Lifecycle: :meth:`start` (main thread, before ``load_model``) →
:meth:`finish` (after the graph exists) or :meth:`abort`. Exactly one of
``finish``/``abort`` must run; both join the reader thread FIRST, because
dropping the slabs while a D2D copy is still in flight would free storage
the device is writing into.

Arm early-H2D for ``blob_path``, or return ``None`` if it cannot arm.

Never raises into the boot: a header that will not parse, a device that
cannot hold the plan, or any other setup failure is logged and returns
``None``, and the caller takes the unchanged in-place load path.

Release everything without binding (incompatible dump, cold fallback).

Joins first — the reader thread holds byte views into the slabs and may
have D2D copies in flight; freeing under it is a use-after-free.

Bind the pre-filled slabs into ``model`` and gate the content.

Returns a stats dict shaped like :func:`load_flat_weights`' so the
existing ``boot.weights`` breakdown logger keeps working, plus
``early_wait_s`` — the part of the DMA this thread actually WAITED for,
which is the honest "still unhidden" number.

Make this path's orphans COLLECTABLE, on the boot thread, in phase.

The bind re-pointed every tensor it touched, so the storage the
meta-materialize walk had allocated for them is now garbage — as are
the graph-build transients the in-place path returns in its own
pre-fill reclaim (which this path skips, having nothing to reclaim FOR:
the weights pool maps first, into clean space).

What that requires is a ``gc.collect``, and ONLY a ``gc.collect``. The
rebinds themselves drop the old storages by refcount the moment
``_set_attr`` replaces the slot, so the bytes are already back in the
caching allocator; what needs collecting is whatever the build left in
reference CYCLES, which refcounting cannot reach and which no other
step on this path runs a collector for.

Returning those bytes to the DRIVER is deliberately NOT done here.
``_load_model_phase`` runs an unconditional ``torch.cuda.empty_cache()``
a few milliseconds later, on the main thread, in this same phase and
four phases before ``settle_vram_for_kv_sizing`` — and ``settle`` is
itself a ``gc.collect`` + ``empty_cache`` backstop INSIDE the barrier.
Doing a second full ``empty_cache`` here only duplicates that
``cudaFree`` work on the join's critical path, which is measured boot
time. The invariant the ``kv_sizing_late_growth`` counter enforces is
"nothing releases AFTER the barrier", and two unconditional releases
already sit before it.

Returns the MiB the driver took back across the collect (usually 0 —
the return happens at the ``empty_cache`` above).

Re-check each bound tensor against the dump's sampled bit signature.

Byte-for-byte the same gate ``load_flat_weights`` runs: it is the only
thing standing between a corrupt cache and silently-wrong weights, so
the early path must not skip it.

Placeholder-vs-real discrimination for the flat-dump bind seams.

Two seams rebind a live module tensor onto a flat dump's storage: the warm
reload's presize (:func:`~arbi_serve.loader.flat_dump._presize_placeholder_buffers`)
and early-H2D's ``EarlyWeightDMA.finish``. Both must GROW a deferred-bind
PLACEHOLDER and both must REFUSE a real tensor whose dtype/shape disagrees
with the dump — rebinding there discards a tensor the live graph built and
serves the dump's bytes under its name, which is a wrong answer rather than a
slow boot.

The predicate and the refusal message live here so the two seams cannot drift
apart: a guard landing on one copy only is the failure mode this module
exists to prevent.

True when ``t`` carries no bound payload, so rebinding it loses nothing.

A graph built under ``skip_weight_load`` registers its quantized weights as
zero-element buffers (``torch.empty(0, dtype=...)``) and the
meta-materialize walk leaves meta tensors; in both cases the backend bind
was deferred to the dump, so the dump's shape/dtype is the authority.
Anything else holds bytes the live graph produced.

Refusal text naming the tensor and BOTH dtypes and BOTH shapes.

A mismatch here is a graph/dump divergence an operator has to diagnose from
the log alone, so the message carries every value the comparison used.

Flat-dump cache-root resolution + content fingerprinting.

Decides where a flat weight dump lives on disk and computes the content
fingerprint that keys a cache hit. The public names
(``default_flat_cache_dir``, ``_FLAT_CACHE_VERSION``, ...) are re-exported
from :mod:`arbi_serve.loader.flat_dump`.

Return True if ``root`` exists and a temp file can be written there.

Creates the directory (parents) if missing, then writes and unlinks a
probe file. Returns False on any OSError instead of raising.

Resolve the root directory holding flat-dump cache subdirs.

Resolution order:

1. :envvar:`ARBI_SERVE_FLAT_CACHE_DIR` non-empty → use it verbatim
   (explicit operator override). Not writable → raise ``RuntimeError``.
2. :envvar:`ARBI_SERVE_FLAT_CACHE_USE_NFS` opted in AND a shared mount
   is present → ``/cache/flat-weights`` (NFS symlink), else
   ``/mnt/k8scache/flat-weights``.
3. Otherwise ``~/.cache/arbi-serve/flat-weights`` (node-local).

NFS is OPT-IN, not the default. The post-ready background warm dump
minutes, stalls live decode through the shared CUDA context, and trips
the soft-mount timeout (EIO) mid-write — so a node-local target is the
safe default and the shared share is reserved for hosts whose
``/cache`` is a HARD-mounted, fast share (opt in with
``ARBI_SERVE_FLAT_CACHE_USE_NFS=1``). Production TP deployments whose
``/cache`` is a durable LOCAL mount set ``ARBI_SERVE_FLAT_CACHE_DIR``
explicitly (case 1), so they are unaffected by this default.

Auto-detected NFS roots (case 2) that turn out not to be writable log a
warning and fall back to the node-local default. The explicit env
override fails loud — silently ignoring an operator-named dir is a
footgun.

The ``awq_a8`` digest field, or ``None`` when it does not apply.

A W4A8 repack REWRITES the weight bytes — it folds the zero-points into the
4-bit nibbles and pre-multiplies the group scales by 512 — so an a8 dump and
an a16 dump of the same checkpoint hold genuinely different bytes and must
not share a cache dir. Without this field they do, and the warm-reload
provenance check catches it as a REFUSAL: correct, but it makes an operator
delete a valid dump by hand every time the a8 policy moves.

Emitted ONLY when a8 actually resolves ON, which keeps the change
backwards-compatible in the direction that matters: every dump written
before a8 became a default carries no field, so an a16 boot (including
``ARBI_SERVE_AWQ_NO_A8=1``) still hits it. An a8 boot resolves to a
different digest, builds its own dump, and both coexist — flipping the
policy back and forth costs one extra build each way instead of a manual
``rm -rf``.

Resolved here rather than threaded through the eight call sites: this is
the one seam that already declares itself "the single source of both
_content_fingerprint and the on-disk manifest", and a second copy of the
rule is precisely how a boot comes to load bytes baked under a different
decision.

Best-effort by construction — a checkpoint that cannot be inspected, a
non-AWQ checkpoint, or a machine without CUDA all return ``None`` and get
the pre-existing key.

Ordered ``(tag, value)`` fields the content digest is folded from.

The single source of both :func:`_content_fingerprint` and the on-disk
manifest, so a dir's recorded provenance cannot describe a different key
than the one that named it.

Field ORDER and the exact bytes are part of the on-disk contract: adding a
field unconditionally re-keys every dump on every deployment. Fields that
are not the shipped default are therefore emitted only when they deviate.

Return a short hex digest fingerprinting ``model_path``'s content.

Hashes cheap metadata only — never weight bytes. Folds in the fields
:func:`_fingerprint_fields` builds: the loader version constant, the
parallel topology ``(tp_size, ep_size, rank)`` plus ``moe_ep_size`` /
``attn_dp_size`` / ``embed_quant`` / ``dtype`` when they are non-default,
the safetensors index json, each shard's ``(name, size, mtime_ns)``, and
``config.json``. Re-quantizing or re-exporting the same model (same
shapes, new bytes) changes shard mtime, so the digest changes and the
stale dump misses.

``enable_vision`` is accepted but NOT folded in — a vision dump is a
superset a text-only boot can use, arbitrated by
:func:`~arbi_serve.loader.flat_dump.flat_dump_compatible` rather than by
the key. The field is emitted as the frozen byte ``b"0"`` so pre-existing
dumps keep their digest.

Falls back to hashing the absolute path plus a marker when none of
index/config/shards exist (e.g. a unit-test fake dir), so distinct
paths still produce distinct keys.

Manifest rendering of one field value.

Short values are recorded verbatim so the manifest is readable; long ones
(``config.json``, a sharded index) are recorded as ``sha256:<hex>`` so the
file stays small while still pinning the exact bytes the digest folded.

The manifest object recorded beside a published dump.

``fields`` is a LIST of ``[tag, value]`` pairs, not a mapping: ``index:``
and ``shard`` repeat, and the digest is order-sensitive, so a mapping
would neither round-trip nor detect a reordering.

A resolved flat-dump cache entry: where it lives and what keys it.

Carries the field list the digest was folded from so the publish seam can
record provenance and the load seam can re-check it, without either having
to re-derive the key from the caller's arguments a second time.

LRU stamp for one cache dir: the later of its use marker and its blob.

The blob's mtime is the fallback for a dump published before use markers
existed, and for one nothing has warm-loaded yet.

Record that this dump was USED, so LRU ranks it by reads, not writes.

Best-effort: a read-only or full cache root must never fail a boot that is
otherwise serving-ready — the cost of a missed touch is that this entry
looks older than it is and may be evicted early, which is one cold boot.

Delete least-recently-used dumps until the root fits under its cap.

Returns the entries removed. Called after a successful publish — the one
moment the root is known to have just grown — so a long-lived node's cache
is bounded by the cap rather than by how many checkpoints it has ever
served.

``keep`` is never evicted: it is the entry this boot just published or is
serving from, and evicting it would make the sweep undo the publish that
triggered it. Neither is an entry used within :data:`_EVICT_MIN_AGE_S`
(this process cannot see another's references), nor one the stable-VA park
owns (see :data:`_PARK_OWNED_SUFFIXES` — those are not caches).

Safe by construction for everything that remains: a flat dump is a pure
latency cache, so losing one costs a cold boot and can never make an answer
wrong. Removing a blob a concurrent boot has open is also safe — POSIX
unlink drops the name, not the mapping, and NFS silly-renames it.

Resolve ``model_path``'s cache entry — dir, digest, and key fields.

The seam every caller should use: it hands back the field list the digest
was folded from, which is what makes the on-disk manifest (and therefore
any later audit of why a dir exists) possible.

Return the cache dir for ``model_path``'s flat dump.

The root is resolved by :func:`_resolve_flat_cache_root` — the
explicit :envvar:`ARBI_SERVE_FLAT_CACHE_DIR`, else (only when
:envvar:`ARBI_SERVE_FLAT_CACHE_USE_NFS` is opted in) the NFS-shared
``/cache/flat-weights`` (or ``/mnt/k8scache/flat-weights``), else the
node-local ``~/.cache/arbi-serve/flat-weights`` default. The dump carries
no SM code, so one NFS dump warms every node of the same arch — it is not
arch-agnostic (see :data:`_NFS_CACHE_CANDIDATES`) — and NFS is opt-in
because the background warm dump must not stream ~GBs over a soft-mounted
share (see :func:`_resolve_flat_cache_root`).

A cache hit is decided by a content fingerprint
(:func:`_content_fingerprint`) over the model's index json,
``config.json``, per-shard (name, size, mtime), the parallel topology,
the loader version, and — when non-default — ``embed_quant`` and the
engine ``dtype``; never the model path. Two copies of the same model at
different paths hit the same dump; re-quantizing or re-exporting the same
model (same shapes, new bytes) misses, because the shard mtime and index
json change.

``enable_vision`` is NOT part of the key: a vision dump is a superset a
text-only boot can warm-load, so both resolve to one dir and
:func:`~arbi_serve.loader.flat_dump.flat_dump_compatible` arbitrates.
There is correspondingly no ``.vis`` filename suffix.

The model name, the readable ``.v3`` version, and the ``.eq`` /
``.tp.ep.r`` suffixes stay in the filename for grepability; the same
values are folded into the digest, so the digest alone disambiguates.

Prefer :func:`flat_cache_key` in new code — it also hands back the fields
the digest was folded from, which the on-disk manifest records.

Record provenance beside a freshly published dump.

Best-effort and atomic: a cache root that has gone read-only must not
fail a boot that is otherwise serving-ready, and a process killed
mid-write must not leave a truncated manifest that the next boot reads
as a disagreement.

Whether an on-disk manifest matches the freshly recomputed key.

``(True, "no_manifest")`` for a dump published before manifests
existed — refusing those would send every shipped deployment cold for
no correctness gain. ``(False, reason)`` means this dir's dump was
produced under a DIFFERENT key than the one that just resolved to it
(a hand-copied dir, a renamed cache entry, a truncated digest
collision), so its bytes must not be bound to this graph.

Post-warmup tensor dump + fast reload for instant restart.

Boot through the safetensors+``weight_loader`` path is CPU-bound: tensor
parse, dtype cast, per-rank shard slice, dispatch through every parallel
linear's loader.

For instant restart we side-step that: dump every persistent tensor
once (post-warmup) into a single safetensors file shaped exactly like
the model's ``state_dict()``. On reload,
``safetensors.safe_open(..., device='cuda:0')`` does the mmap +
multi-threaded H2D itself; we just ``copy_`` into the existing GPU
storage. No Python iteration over storage; no bespoke header format;
no dtype gaps.

Capturing ``state_dict()`` instead of ``named_parameters()`` matters
for quant backends (AWQ Marlin, EXL3 trellis, FP8 dynamic): they bind
extra persistent buffers post-load — Marlin-reordered weight,
codebook LUTs, FP8 inverse scales — that a ``named_parameters()``
capture would miss, forcing the quant-backend bind path to recompute
them on every warm boot.
``state_dict()`` includes them; non-persistent buffers (RoPE freq
caches, etc.) stay excluded because their semantics are
deterministically rebuildable.

Tensor names match ``model.state_dict()`` using the fused-linear
convention (``qkv_proj.weight``, ``gate_up_proj.weight``, etc. — never
``.fused_weight``). Shards are baked in at dump time, so a TP=1 dump
cannot be loaded by a TP=2 deployment without reslicing.

Both functions are CPU/IO-bound and are wrapped in ``asyncio.to_thread``
by the engine wrappers so the FastAPI event loop never blocks.

Return a contiguous 1-D ``uint8`` view of ``t`` for raw byte DMA.

``t.view(torch.uint8)`` raises on a 0-dim (scalar) tensor — quant
backends carry scalar buffers (EXL3 mcg/mul1 codebook seeds, NVFP4
weight_scale_2). Reshape to 1-D FIRST so the byte view is always
valid; a scalar becomes a 1-element row, ``nbytes`` is unchanged.

Clean up a gated-abort partial dump and return an ``aborted`` result.

Removes the half-written temp file (a partial dump is never published —
the next boot must not read it) and reports the elapsed time so the
caller can log the abort. Shaped like a normal result minus the
byte/throughput keys, plus ``aborted=True`` for the caller to branch on.

Dump every persistent tensor to ``out_dir/weights.safetensors``.

The capture is ``model.state_dict()`` — parameters AND persistent
buffers — so a quant backend's post-load buffers (Marlin-reordered
weight, EXL3 trellis LUTs, FP8 inverse scales) round-trip instead of
being recomputed on every warm boot. Names match ``state_dict()`` so
the reload is a by-name copy. Returns ``{out_dir, num_tensors, total_bytes,
write_s}`` on success, or ``{aborted: True, ...}`` if ``gate`` asked
to stop (the partial temp file is removed — a partial dump is never
published).

``only_names`` restricts the dump to that subset of ``state_dict()``
(fail-loud if a requested name is absent from the model) — a tiny
SUPPLEMENT dump carrying just the live tensors a shared bulk dump lacks
(e.g. an in-memory-quantized drafter's marlin buffers), so a cross-model
park never re-serializes the whole model. ``None`` dumps the full
``state_dict()``.

``gate`` is an optional cooperative checkpoint consulted before the
GPU is first touched and again before each tensor's D2H copy. The
callable MAY BLOCK to throttle the dump (the engine passes one that
pauses the stream while live requests are in flight, so the ~GB D2H
never contends with decode through the shared CUDA context) and MUST
return ``True`` to proceed or ``False`` to request a clean abort
(e.g. server shutdown). ``None`` (the default) streams straight
through — the legacy behavior used by tests and the calibrate path.

Synchronous CPU/IO. Callers in async contexts must wrap in
``asyncio.to_thread``.

Dump-time structural gate: the data section must exactly tile the file.

The streaming writer assigns contiguous per-tensor offsets; this re-reads
the just-written header and proves the tensors form a gap-free, non-
overlapping cover of ``[0, expected_data_bytes)`` and that the file size
equals ``header + data``. A writer offset bug fails HERE (before the
atomic rename publishes the dump), never silently at a future warm boot.

Fraction of ``blob_path`` currently in the page cache, or ``-1.0``.

Samples evenly spaced windows with ``mincore(2)`` rather than walking the
good enough to choose a read mode. Returns ``-1.0`` when residency cannot
be determined (no ``mincore``, an unreadable file, a platform without
it), which callers treat as "unknown" and resolve to ``O_DIRECT``.

Resolve ``ARBI_FLAT_DIRECT_IO`` for ``blob_path``, logging the choice.

``auto`` prefers whichever source is actually faster for THIS dump on
THIS boot: the page cache when the dump is already in it, the device
otherwise.

Open ``blob_path`` for ``O_DIRECT`` reads, or ``None`` if unavailable.

Returns a descriptor only when the platform exposes ``O_DIRECT``, the
destination buffers are alignment-satisfying, the filesystem accepts the
flag and an aligned read through ``probe`` actually succeeds. Every other
outcome yields ``None`` so the caller falls back to a buffered read.

Read ``[span_lo, span_hi)`` of ``blob_path`` into a ring of buffers.

Reserves a buffer index from ``free_q``, fills it, and publishes
``(index, file_offset, nbytes)`` on ``full_q``; the consumer recycles the
index back onto ``free_q``. Slabs are emitted in ascending file order and
together cover the whole span.

``O_DIRECT`` is used when :func:`_open_direct` grants it: the read then
bypasses the page cache, which both removes a full-size kernel memcpy the
pinned ring makes redundant and lets one slab-sized request reach the
device instead of being decomposed into read-ahead windows. Because
``O_DIRECT`` constrains the offset, the length and the buffer, the direct
path starts at the block boundary at or below ``span_lo``, rounds the
final read up to a block, and fills only the block-aligned prefix of each
buffer. The extra head and tail bytes belong to no range, so the
consumer's overlap scatter ignores them. Otherwise the span is read
buffered, which has no alignment constraint.

Returns the mode used, ``"direct"`` or ``"buffered"``.

Sequential-read + overlapped H2D-then-D2D-scatter of file byte ranges.

Shared DMA core for both the flat-dump loader and the direct-DMA dense
fill. ``ranges`` is ``(file_offset, nbytes, dest_byte_view)`` where each
``dest_byte_view`` is a contiguous uint8 CUDA view exactly ``nbytes``
long and ``file_offset`` is absolute in ``blob_path``. The ranges are
sorted by offset here, then the file's covered span is read once,
sequentially, in ``_SLAB_BYTES`` slabs by a producer thread (see
:func:`_read_span` for the O_DIRECT-or-buffered read itself); the consumer
copies each pinned slab to a paired GPU scratch slab in one H2D and
scatters it from there to the overlapping destination views with D2D
copies. A slab may carry alignment padding outside ``[span_lo, span_hi)``;
it overlaps no range, so the scatter ignores it.

Correctness: every source byte lands at exactly the destination its
range names — a slab is split at range boundaries and each piece copied
with ``narrow`` into the right offset of the right view (the H2D→D2D
indirection moves the same bytes through GPU scratch, byte-for-byte). Ranges
may OVERLAP — two destination views can name the same source bytes (a tied
embedding aliased to ``lm_head`` reads one ``src_key`` into both the embed
and the head storage), so each slab scatters into every range it overlaps,
not just the next one in offset order. The final ``copied == total_dst``
gate proves full coverage. ``data_base`` is informational (the caller has
already folded it into ``file_offset``).

Returns total bytes copied; if ``progress`` is given, bumps
``progress[0]`` as slabs are consumed so the boot poller can report live
%. If ``log_total`` is given, this function also emits the periodic
"Boot: loading weights" line itself (used when the caller has no poller).

``staging_pool`` is the :class:`~arbi_serve.runtime.named_pool.NamedMemPool`
the GPU landing ring is allocated from. It MUST NOT be the destination
transit scratch: it is dropped when this function returns, but a named
cuMem-backed ``MemPool`` never hands freed segments back
(``torch.cuda.empty_cache()`` does not visit a private MemPool's block
pools — pytorch#145168, re-measured on torch 2.12.1). So whichever pool is
life of the process. On the warm flat-dump path the caller's active pool was
Routing the ring to the loader's own scratch pool lets the post-load
``NamedPoolRegistry.release_empty_pool`` return the bytes to the driver
before the KV budget is sized, so they become KV pages instead of a
permanent free list. ``None`` keeps the caller's active pool (the cold path, which runs
outside any pool scope and so lands in the default allocator where
``empty_cache`` genuinely reclaims).

``ring_bufs`` / ``dma_streams`` override :data:`_RING_BUFS` /
:data:`_DMA_STREAMS` for this call. The defaults are tuned to saturate the
read on a boot that is WAITING for these bytes. A caller that is NOT
waiting — early-H2D streams under a graph build with seconds of slack —
can trade transfer rate for less interference with whatever it is hiding
behind: each ring buffer is a pinned ``cudaHostAlloc`` (which
device-synchronizes) plus a GPU slab, and each stream adds per-slab launch
work on the consumer thread, which holds the GIL against the concurrent
build.

Rebind any persistent buffer whose live shape/dtype differs from
the flat dump's to a fresh empty tensor at the dump's shape/dtype.

Used by the quant warm-reload path: a quantized model built with
``skip_weight_load`` has placeholder buffers (``torch.empty(0)``)
because the swap deferred the backend bind. The chunked DMA workers
assert an exact shape match against the dump, so the destination must
be the right shape BEFORE the fill.

Only a genuine PLACEHOLDER is grown (:func:`is_growable_placeholder` —
meta or zero-element). A buffer that already carries BYTES and still
disagrees with the dump is a graph/dump divergence: rebinding it would
discard the live tensor and serve the dump's bytes under its name, so it
raises instead. An already-matching param/buffer is skipped, so dense
weight storage (and any cudagraph capture pinned to it) is never
disturbed.

A mismatch on a tensor that is NOT a registered buffer (i.e. a real
parameter whose shape disagrees with the dump) is a genuine error —
the chunked worker will raise on it; we do not silently rebind
parameters here.

Returns the number of buffers rebound (for the caller to refresh its
``state_dict`` snapshot + the boot log).

Cheap header-only check: can the dump at ``in_dir`` warm-load ``model``?

Reads NO weight bytes — parses only the safetensors header. Returns
``(ok, reason)``. ``ok=False`` means a warm :func:`load_flat_weights` would
fail-loud (missing tensor / parameter shape mismatch), so the caller should
DELETE the stale dump (:func:`remove_flat_dump`) and take the cold load path
instead of aborting the boot — the "fall to cold with a clear log" contract.

The check mirrors the loader's own fail-loud conditions:

  * every live ``state_dict()`` tensor name must appear in the dump (a dump
    MISSING a live tensor cannot fill the graph). A post-head-quant dump is
    missing the dense ``mtp.fc.weight`` and so is correctly rejected here;
  * every live PARAMETER's dtype+shape must equal the dump's. Parameters are
    filled in place and never resized;
  * every live persistent BUFFER that is not a PLACEHOLDER must likewise
    match. A placeholder (meta or ``torch.empty(0)``) is exempt because the
    quant warm reload grows it to the dump's shape via
    :func:`_presize_placeholder_buffers` before the fill, so its pre-fill
    shape is expected to differ. A buffer that already carries BYTES is not
    exempt: :func:`_presize_placeholder_buffers` refuses to rebind it, and
    rejecting it here is what turns that refusal into a cold fallback
    rather than a failed boot.

A byte-level corruption that matches shape+dtype is NOT caught here (it needs
the filled buffer) — that stays the loader's post-fill bit-signature gate,
which fail-loud aborts. This function only decides cold-fallback vs warm.

Delete a stale/incompatible flat dump so the next boot re-publishes it.

Removes the published blob, its provenance manifest, and any leftover
``*.partial.*`` temp files, but leaves the (fingerprinted) directory so the
dense-seam dump can rewrite it.
Best-effort: an OSError is swallowed (a read-only cache dir just means the
stale dump stays and every boot falls to cold — safe, if slow).

Inverse of ``dump_flat_weights``. ``model`` must already be on
GPU with parameters allocated at the same shapes / dtypes. Quant
backends are an exception: a model built with ``skip_weight_load``
carries empty placeholder buffers for its quantized weights, which
:func:`_presize_placeholder_buffers` grows to the dump's shape before
the fill (the backend's ``rebind_after_compaction()`` then rebuilds
the kernel descriptor over the filled buffers — fired by the caller).

``supplement_dir`` names an optional SECOND dump whose tensors take
priority over ``in_dir``'s: it fills the live tensors the bulk dump at
``in_dir`` lacks (an in-memory-quantized drafter's marlin buffers the
head-quant-independent warm dense cache never held), so a cross-model
park reloads bulk-from-warm + delta-from-supplement without ever
re-serializing the whole model. The union of the two dumps must cover
every live tensor — a live tensor in NEITHER is still refused loud.

Uses a single sequential reader overlapped with multi-stream H2D
(:func:`_sequential_dma`): one producer thread pulls the file's data
region in large contiguous slabs at full cold-disk read-ahead bandwidth
into a bounded pinned ring, while the consumer scatters each slab into
the GPU destination views across a couple of CUDA streams and recycles
buffers via a non-blocking event query. ``n_workers`` / ``chunk_mb``
intentionally single-threaded; extra readers thrash a cold disk).

allocated from (see :func:`_sequential_dma`). Pass the loader's own scratch
pool: this function is called under the WEIGHTS pool scope (the destination
storage must be there), and a cuMem-backed MemPool never returns a freed
free list on ``model.weights``.

Synchronous CPU/IO + GPU sync. Callers in async contexts must wrap
in ``asyncio.to_thread``.

CPU-only fallback. The multi-stream pinned-DMA path doesn't
apply when the destination isn't a CUDA device (mostly used by
tests). Walks ``safe_open(device='cpu')`` over each source.

``supplement_blob`` mirrors the CUDA path's ``supplement_dir``: a second
dump whose tensors take priority, filling the live tensors the bulk dump
lacks. Their union must cover the live state_dict — a tensor in neither is
refused loud.

DMA raw byte ranges from safetensors shards straight into
pre-allocated CUDA destination views, in place.

This is the cold-boot dense-fill counterpart of
:func:`load_flat_weights`: instead of reading a single flattened blob
it reads the *original* safetensors shards, and instead of a
name-keyed ``state_dict`` it writes each range into a caller-supplied
destination tensor (a flat-pool slab view bound by
:class:`arbi_serve.loader.flat_loader.AllocateThenFillBinder`). It shares
the same sequential-reader + multi-stream H2D core (:func:`_sequential_dma`)
as ``load_flat_weights``.

Args:
    items: ``(shard_path, file_offset, nbytes, dest_view)`` tuples.
        ``dest_view`` MUST be a contiguous CUDA tensor whose byte
        length equals ``nbytes`` and whose on-disk wire dtype matches
        (the caller guarantees both — see ``_direct_dma_fill``). The
        copy is raw bytes, so no cast / slice is performed here.
    device: target CUDA device string (e.g. ``"cuda:0"``).
    n_workers: accepted for back-compat; ignored (the read is a single
        sequential pass per shard, which is read-bound on a cold disk).
    chunk_mb: accepted for back-compat; ignored (a fixed large slab is
        used; see ``_SLAB_BYTES``).

Synchronous (sequential read + H2D + sync). Wrap in ``asyncio.to_thread``
from async contexts.

Flat-buffer compaction for cold-boot weight loading.

The generic :func:`arbi_serve.loader.weights.load_model_weights` plus
the per-quant-backend bind path (AWQ/FP8/EXL3) issues dozens of small
``cudaMalloc`` calls per layer (raw safetensors reads, intermediate
Marlin repack outputs, the final persistent Marlin tensors, dense
``nn.Parameter`` storage). PyTorch's caching allocator rounds each
allocation up to its segment-class minimum, so the many small
irregular-sized allocations across all layers leave a fragmented,
reserved-but-unallocated tail in ``weights_pool``.

``torch.cuda.empty_cache()`` does NOT reclaim segments inside a named
``MemPool``; per-pool release isn't exposed in public torch. Compaction
is the practical fix:

  1. Boot routes weight loading through a SCRATCH pool
     (``weights_scratch_pool``) instead of ``weights_pool`` directly.
     Every transient (raw int32 qweight, dropped intermediates inside
     Marlin repack) plus every persistent (Marlin tensors, dense bf16
     weights) lands here first.
  2. After loading completes, :func:`compact_into_flat_pool` walks
     every persistent tensor (parameters + non-persistent buffers
     created by quant backends), groups them by dtype, allocates ONE
     flat slab per dtype inside ``weights_pool``, and rebinds each
     tensor's ``.data`` to a ``narrow()`` view into the slab. The
     copies from scratch storage to slab views are a single
     ``copy_(non_blocking=True)`` per tensor.
  3. The caller then drops ``weights_scratch_pool`` (release +
     gc.collect + ``torch.cuda.empty_cache()``); every scratch
     segment goes back to the driver.

Result: ``weights_pool`` ends up with one segment per dtype holding
the full weight footprint with near-zero fragmentation, matching
``kv_pool`` and ``recurrent_pool``.

Bit-identity invariant
----------------------
Compaction is a pure data move: ``param.data.copy_(scratch_tensor)``
is the only mutation. Shapes / dtypes / strides are preserved (the
narrow() view has the same shape as the source, and slabs are
allocated with identical dtype). The numerical content of every
tensor is unchanged — the slab views point at the same byte sequences
the scratch tensors held.

Memory headroom invariant
-------------------------
Two compaction strategies live in this module:

  - :func:`compact_into_flat_pool` (batch): scratch and the
    destination slabs are held simultaneously; the per-rank peak
    during compaction is roughly ``scratch_reserved + slab_reserved``.

  - :func:`stream_compact_per_layer` (default cold-boot path):
    pre-allocates slabs in ``dst_pool``, then walks the model
    per-layer and copies each layer's tensors into its slab slots,
    dropping the scratch refs immediately. Periodic ``gc.collect()``
    + ``torch.cuda.empty_cache()`` returns released scratch segments
    to the driver as we go, so the peak is roughly ``slab_reserved +
    one-layer-of-scratch + transients``.

Operators select between the two via ``cfg.cold_load_strategy``
(``"batch_compact"`` | ``"stream_compact"``); default is
``"stream_compact"``.

Pre-bind a model's meta parameters to flat-pool slab views so the
weight loader fills them in place — cold-boot peak ≈ 1× model size.

The streaming compactor (:func:`stream_compact_per_layer`) packs an
already-materialized model into flat slabs: the source weights and
the destination slab are both resident while the copy runs, so for a
single-dtype model (one slab == the whole model) the peak spikes to
~2× and OOMs a card that the model otherwise fits on.

This binder removes the spike for dense checkpoints. Used before the
weight loader runs:

  1. :meth:`bind_all` walks every meta parameter / buffer, groups by
     dtype, allocates one flat slab per dtype inside ``dst_pool``,
     and rebinds each tensor to a ``narrow()`` view into its slab via
     :func:`_set_attr`. The slabs are the only persistent weight
     allocation; no per-tensor ``cudaMalloc``, no scratch.
  2. The normal :func:`arbi_serve.loader.weights.load_model_weights`
     then runs. Every weight loader writes in place
     (``param.copy_`` / ``param.weight[...].copy_``), so each copy
     lands directly in the slab. Sliced params (TP / fused / expert)
     transit one CPU/GPU temp before the in-place copy; peak stays at
     ``slab + one transient tensor``.

The binder is meta-only by construction: it asserts every tensor it
visits is on the meta device, so it can't silently clobber a quant
backend's already-materialized buffers. The quant cold path keeps the
load → :func:`stream_compact_per_layer` route instead.

:meth:`verify` re-hashes the filled slab views against the pre-fill
signatures captured at bind time and checks the per-dtype cursor
advance, mirroring the streaming compactor's bit-identity + cursor
gates so silent corruption still fails the boot.

One-line diagnosis of WHY the donor has no tensor named ``qname``.

A gap is filled from the checkpoint rather than refused, so the reader needs
the two cases separated without rebuilding anything: a PREFIX SKEW (the donor
holds the same leaf under a different root — the weights exist and the paths
disagree) versus a genuinely absent tensor (the member's config built
structure the donor never had). The suffix match names the skewed donor keys.

Whether the donor holds ANY tensor under the module ``path``.

Separates the two reasons a donor tensor can be absent: the donor built the
module but not this one optional buffer (nothing to share — skip), versus
the donor never built the module at all (a coverage gap the caller must
fill from the checkpoint). The root module (``""``) is always covered.

Bind a freshly-constructed meta module graph to a donor module's
live tensors — same-model weight sharing, zero new weight VRAM.

The stable-VA residency pool's same-``(model.path, dtype)`` members
(e.g. a live config-override variant that differs only in
capture-affecting knobs like ``max_batched_tokens``) would otherwise
re-load and re-allocate the full weight set. Weights are immutable at
serve time, so a second member can bind the same storage into its
own module graph:

  * the new member gets its own ``nn.Module`` graph (so per-member
    module state — the compile trampoline, per-call side-channels,
    derived-weight buffers, layer-dispatch registrations — stays
    distinct),
  * every meta param/buffer the donor also holds is replaced by a
    tensor sharing the donor's storage (:func:`_rebind_meta_to_view`),
    matched by qualified name; a name the donor holds at a DIFFERENT
    shape aborts the build rather than serving a mis-bound model,
    while a name the donor does not hold at all is a coverage GAP the
    caller fills from the checkpoint (:attr:`unshared_params` /
    :attr:`unshared_modules`).

Used via the same :func:`set_active_binder` seam as
:class:`AllocateThenFillBinder` (the ``_materialize_meta_to_empty``
funnel), with ``set_skip_weight_load(True)`` so the safetensors fill
is skipped entirely — the donor's bytes are the weights.

Dense, weight-quantized (EXL3 / AWQ / FP8 / NVFP4), and fused-MoE
checkpoints. A fused-MoE arch (Qwen3.5-MoE, DeepSeek-V3, LFM2-MoE) holds
its stacked expert weights (``mlp.experts.w13_weight`` / ``w2_weight`` +
the fp8 ``*_scale``) as persistent buffers on a plain :class:`FusedMoE`
module, registered ``torch.empty(0)`` under ``skip_weight_load`` and sized
only after this binder runs; they alias the donor's populated stack by the
same zero-element-placeholder rule as the quant weight buffers below.
A quant member builds its module graph under ``skip_weight_load`` — the
quant swap runs but ``backend.bind()`` is deferred, so every persistent
buffer of a quant linear is an ``__init__`` placeholder/default, exactly
like a warm flat-dump reload. The binder aliases the donor's immutable
buffer for each (zero new weight VRAM): the zero-element weight buffers
(EXL3 ``trellis``/``suh``/``svh``, AWQ/NVFP4 packed weight + scale buffers)
and the small non-empty derived-state marker buffers the rebind hook reads
to reconstruct the python kernel flags (AWQ's ``_format_marker``, NVFP4's
scalar ``weight_scale_2`` / ``_weight_only_flag`` / ``_marlin_ready_flag``).
Sharing only the zero-element weight buffers is not enough for AWQ/NVFP4:
their rebind hook keys on the marker (not on a weight buffer, the way EXL3
keys on ``trellis``), so a stale-default marker makes the hook no-op and the
first profiling forward aborts with "AWQ linear forward called before any
load()". The caller then fires the standard post-load rebind hook
(:func:`_run_compaction_rebind_hooks`) so each quant linear rebuilds its
derived kernel descriptor over the shared buffers (read-only — the
descriptor references the donor storage; it never mutates it). This makes a
live config-override flip of a capture-affecting knob on a large EXL3/AWQ
model (27B) share the immutable weight residency with the parent member
instead of paying a second full weight copy.

``binds_quant_placeholders`` tells the materialize seam this binder can run
even when a quant backend swapped modules (``AllocateThenFillBinder`` cannot
— it pre-sizes slabs the quant repack would invalidate).

Install (or clear) the process-wide meta-binding weight binder.

Set by the cold-boot dense path (:class:`AllocateThenFillBinder`) or
the same-model weight-share path (:class:`DonorWeightBinder`) in
:func:`arbi_serve.engine.build_helpers.load_model_into_engine`
around the model load; cleared in a ``finally``. The model-load seam
:func:`arbi_serve.models._layer_stack_model_mixin._materialize_meta_to_empty`
reads it via :func:`get_active_binder`.

Re-hash every bound slab view AFTER the loader filled it and
return the streaming-compactor stats shape (``bit_identical`` +
cursor totals) so the caller's existing gate works unchanged.

The signatures are read AFTER fill, so a corrupt write (wrong
narrow offset, dtype mismatch) shows up as a slab whose bytes
don't round-trip through a fresh hash of the same view; we hash
each view twice (pre/post a synchronize) to surface in-flight
copies that never landed.

The live donor module this binder shares storage from.

Exposed so the model-load seam can reproduce a donor's runtime
head-quant module swap (``apply_head_quant`` — a post-load build
phase, so the fresh member graph builds a dense head the donor no
longer carries) on the member before this binder walks it, sharing
the donor's immutable INT4 head modules instead of aborting the
exact-name bind on the missing dense ``mtp_head.fc.weight``.

Rebind every meta param/buffer of ``model`` to the donor tensor
of the same qualified name. Idempotent guard: refuses a second bind.

A donor covers the SHARED part of the graph, not necessarily all of it:
a member whose config ADDS parameters (``mtp.enabled`` off → on builds
the checkpoint's bundled MTP head the donor never constructed) presents
names the donor has no tensor for. Those are recorded on
:attr:`unshared_params` / :attr:`unshared_modules` and left for the
caller to fill from the checkpoint, so the share still saves every
weight the two graphs DO have in common. A dense gap param is
materialized here as empty ``device``/``dtype`` storage (the same shape
the dense meta-fill seam would have made); a quant module with zero
donor coverage keeps its deferred-bind placeholders, which only
``backend.bind()`` can size.

Divergences that are NOT a coverage gap stay fail-loud: a name the
donor holds at a different SHAPE, and a donor tensor still on meta.

Shared helpers for the flat-buffer cold-boot loader package.

Constants, the persistent-tensor walkers, the bit-identity signature, the
slab-slot alignment math, the attribute-rebind primitives, and the
per-dtype slab packing core shared by the binders and the compactors.
See the package ``__init__`` for the module-level overview.

Invoke the post-compaction rebind hook on every module that
declares one.

Called at the end of a compaction pass — after every parameter /
buffer has been rebound to its slab view — so a hook implementation
sees the final slab-backed buffers when it re-derives its secondary
references (see :data:`_COMPACTION_REBIND_HOOK`). Returns the number
of hooks fired (for the boot log / tests).

Invoke the pre-compaction secondary-reference release hook on every
module that declares one.

Called at the start of a compaction pass — before any slab is
allocated — so a module that pins its source tensors via a secondary
GPU reference (the EXL3 inner + its BC descriptor) drops that
reference, letting the host-staging guard free the original Returns the number of
hooks fired (for the boot log / tests).

Elements to reserve for one slab tensor, padded so the NEXT tensor's
base lands on a :data:`_SLAB_TENSOR_ALIGN_BYTES` boundary.

The cursor starts at 0 (aligned) and advances by this padded count, so
every tensor's narrow offset is a multiple of the alignment — a no-op
``round_up`` arithmetic only (no allocation). Returns ``numel`` unchanged
when it already fills whole alignment units.

Padded byte size one ``numel``-element tensor reserves in a ``dtype``
slab: the alignment-padded element slot times the dtype size.

Mirrors the per-slot cursor advance (:func:`_slab_slot_elems`) so a
dtype's pre-walk byte total equals the pack-pass cursor advance. Shared
by the per-dtype pre-walks of the slab binders and the compactors.

Return the streaming-compactor group key for a qualified
``module.attr`` path.

``model.layers.7.self_attn.q_proj.weight`` →
``model.layers.7``. Any path that doesn't match the per-layer
pattern (embeddings, norms, lm_head, MTP/drafter heads attached
outside the per-layer prefix) lands in the ``__top__`` group,
which is compacted last.

Order layer groups numerically (``layers.0`` < ``layers.10`` <
``layers.11``) and place ``__top__`` last so per-layer scratch
is released before the small top-level group is moved.

Returns a ``(bucket, idx)`` tuple where ``bucket=0`` is per-layer
and ``bucket=1`` is the top-level group.

Like :func:`_iter_persistent_tensors` but also yields the
fully-qualified name (``model.layers.7.self_attn.q_proj.weight``)
used by the streaming compactor to compute per-layer group keys.

De-dup is identical (by ``data_ptr``); the qualified name comes
from the first visit's owning module path.

Yield ``(parent_module, attr_name, tensor)`` for every parameter
AND non-zero-numel buffer in ``model``.

Walks ``named_modules()`` and inspects each module's ``_parameters``
and ``_buffers`` dicts directly so we get the (mod, attr) tuple for
in-place rebinding via ``setattr(mod, attr, ...)``. Skips zero-numel
placeholders (e.g. AWQ's raw ``self.qweight = torch.empty(0, ...)``
sentinels left after Marlin repack frees the raw int32 input).

Non-persistent buffers (``register_buffer(..., persistent=False)``)
are included too: compaction must pack every live weight tensor into
the slab regardless of its persistence flag, or the post-load pool
keeps the un-compacted originals and stays fragmented. (The quant
backends' core weight buffers are persistent so the warm flat
dump captures them, but a few derived caches stay non-persistent —
those still need compacting while they hold real storage.)

De-duplication: tracks ``data_ptr()`` so a Parameter registered in
two modules (no current arch does this, but the future
LoRA-as-Parameter path might) is moved exactly once. The first
visit performs the move; subsequent visits to the same storage
rebind to the already-moved view via :func:`_set_attr`. Tied
lm_head doesn't go through this path — :class:`TiedLMHead` holds
the embedding in a 1-element list to dodge nn.Module's auto
re-registration, so the embedding storage shows up exactly once
under ``model.embed_tokens.weight``.

Cheap per-tensor signature for the bit-identity smoke check.

Returns a 64-bit integer derived from up to 8 fixed-position bytes
of the tensor's storage. The signature is invariant under
``copy_`` of identical bytes — if compaction silently corrupts
data (wrong narrow offset, dtype mismatch, mid-row TP slice
error), the post-copy signature will differ from the pre-copy
signature and the boot log lights up.

We sample 8 strategic positions (first byte, last byte, and 6
evenly-spaced midpoints) so a transposition or partial overwrite
can't pass undetected. The check is O(1) per tensor regardless
of size.

Rebind ``mod.<attr>`` to ``new_tensor`` in place, preserving
Parameter-vs-buffer semantics.

For an :class:`nn.Parameter`, we update ``param.data`` so the
Parameter object identity is preserved (any module / cudagraph /
LoRA handle that captured the Parameter ref keeps working). For a
plain buffer, we re-issue ``setattr`` which routes through
:meth:`nn.Module.__setattr__` to update ``_buffers``.

Rebind a meta param/buffer ``attr`` on ``mod`` to a real ``view``.

Unlike :func:`_set_attr` (which preserves Parameter identity via
``.data =`` for the already-materialized compaction path), the binder
runs on the pre-load meta graph where no cudagraph / LoRA handle has
captured the parameter yet. A meta Parameter is fp32 by default while
the slab ``view`` is the engine dtype (bf16), so ``.data =`` raises
``set_data ... incompatible tensor type``. We therefore replace the
Parameter object outright (new dtype/device is fine — nothing captured
it). Buffers route through ``setattr`` as usual.

Copy ``t``'s bytes into ``slab[cursor : cursor + t.numel()]`` (reshaped
to ``t``'s shape), rebind ``mod.<attr>`` to that view, and return it.

The per-tensor move shared by the batch (:func:`compact_into_flat_pool`)
and streaming (:func:`stream_compact_per_layer`) compactors: a single
``copy_(non_blocking=True)`` into the slab view followed by
:func:`_set_attr` so the old scratch storage drops.

Allocate one flat slab per dtype and place every grouped tensor into it.

Shared packing core for the slab binders (:class:`AllocateThenFillBinder`,
:class:`Exl3DirectSlabBinder`). For each dtype — largest total first so the
layout is deterministic — it allocates one slab in ``dst_pool``, then walks
``by_dtype_group[dtype]`` in numeric layer-group order (``__top__`` last)
and, per tensor, calls ``place(slab, dtype, elem_size, cursor, item)`` for
the caller's site-specific view/bind/record step before advancing the
alignment-padded cursor (:func:`_slab_slot_elems`). After a slab is filled
it asserts the cursor matches ``slab.numel()`` (a pre-walk/pack-pass byte
disagreement is a dedup bug — refuse to start) and calls the optional
``on_slab(dtype, slab, total_elems)`` for per-slab bookkeeping.

Args:
    by_dtype_group: ``{dtype: {group_key: [item, ...]}}`` in arrival order.
    bytes_by_dtype: padded slab byte total per dtype (see
        :func:`_aligned_slab_bytes`).
    numel_of: ``item -> int`` element count, for the cursor advance.
    place: ``(slab, dtype, elem_size, cursor, item) -> None`` callback.
    on_slab: optional ``(dtype, slab, total_elems) -> None`` callback.
    label: class name for the cursor-mismatch error message.

Returns ``(n_groups, n_tensors, n_bytes)`` where ``n_bytes`` is the raw
(unpadded) byte sum the binders report in their stats.

Copy each ``(mod, attr, tensor)``'s data to a host tensor and
rebind the model attribute to the host copy, freeing the GPU source
storage.

Used by :func:`stream_compact_per_layer` when the destination slabs
would not fit alongside the resident GPU source set: staging the
source tensors to host first lets the caller drain the GPU source
pool to the driver, allocate the slab, and copy host→slab — removing
the transient 2× peak for the compaction path as well.

``pin_memory`` ASKS for pinned host copies (fast async H2D); whether it
gets them is the host's decision, not the caller's. Staging a whole dtype
group is model-scale, and page-locked pages can be neither swapped nor
reclaimed, so a request the host cannot hold does not degrade — the
kernel's global OOM killer picks a process. The group is therefore priced
through :mod:`arbi_serve.runtime.pinned_host_budget` and falls back to
PAGEABLE host memory when it does not fit: the copy is still correct (the
H2D just runs synchronously) and the bytes stay swappable, which is what
makes the fallback safe rather than merely smaller. That fallback is why
``pin_memory`` is a request and not a promise, and why no caller has to
guess at the host's size.

Returns the same ``(mod, attr, host_tensor)`` triples (in input
order) so the caller can copy them into slab views and rebind to the
final GPU views.

Move every persistent tensor in ``model`` into one flat slab per
dtype inside ``dst_pool``.

Strategy: group tensors by ``dtype``; for each group, allocate one
contiguous buffer in ``dst_pool`` of size ``sum(numel * elem_size)``;
for each tensor in the group, copy its bytes into a ``narrow()``
region of the slab and rebind the tensor's ``.data`` (Parameter) or
re-set the attribute (buffer) to the view.

Args:
    model: the fully-loaded model. ``state_dict()`` and
        ``named_buffers()`` already reflect the post-load
        tensor layout (raw AWQ slots are zeroed-out
        ``empty(0)`` sentinels and skipped automatically).
    dst_pool: the named MemPool that will own the slabs. Typically
        the engine's ``weights_pool``; the caller is expected to
        have routed the original load through a separate scratch
        pool so the source storage can be reclaimed by dropping
        the scratch pool after compaction.
    device: target CUDA device.

Returns:
    ``{"groups": int, "tensors": int, "bytes": int, "dtypes":
    {dtype: bytes}}`` — one entry per dtype slab.

Streaming variant of :func:`compact_into_flat_pool`.

Strategy
--------
The batch compactor (:func:`compact_into_flat_pool`) holds the
full scratch pool and the destination slabs simultaneously while
it copies tensor-by-tensor, so the per-rank peak is roughly
scratch plus all destination slabs.

The streaming variant pre-allocates the slabs in ``dst_pool``
sized exactly to the post-load total per dtype, then walks the
model per layer. For each layer's tensors it copies into the
pre-allocated slab views and rebinds ``param.data`` /
``mod._buffers[attr]`` so the old scratch tensor's storage drops
immediately. Every ``empty_cache_every`` layers the function calls
``gc.collect()`` + ``torch.cuda.empty_cache()`` so the released
scratch segments return to the driver as we go (and not just at
the very end via :func:`drop_scratch_pool`), keeping the peak
roughly at ``slab + one-layer-of-scratch + transients``.

``new_param`` is a ``narrow()`` view into the pre-allocated slab
so the destination is a single flat buffer per dtype rather than
a per-tensor ``cudaMalloc``.

Args:
    model: fully-loaded model graph. The persistent-tensor walk
        is post-quant-swap so AWQ Marlin shapes are final.
    dst_pool: target named MemPool (typically ``weights_pool``).
    device: target CUDA device.
    empty_cache_every: number of layer groups between
        ``empty_cache`` calls. Smaller = lower scratch peak,
        higher driver-call overhead. ``4`` is a reasonable default
        for 64-layer dense models; for huge MoE models we can
        tune down to ``1`` for tighter peak control.

Returns:
    ``{"groups": int, "tensors": int, "bytes": int, "dtypes":
    {dtype: bytes}, "pre_hash": int, "post_hash": int,
    "n_hashed": int, "bit_identical": bool, "layer_groups":
    int, "empty_cache_calls": int}`` — superset of the batch
    compactor's stats so callers can use a single log-line
    format for either path.

Release every segment held by ``scratch_pool`` and reclaim the
bytes back to the driver.

The dance: scratch tensors must already be dereferenced by the
caller (all model weights now point at slab views in
``weights_pool``); we then:

  1. Call ``_cuda_releasePool`` to drop the use-count increment
     that ``MemPool.__init__`` planted (use_count: 1 → 0).
  2. Replace the wrapper's ``_mempool`` attribute with C++ destructor run and return segments to the driver.
     Without this, the wrapper holds the last ref and the
     segments stay reserved (allocated=0 but reserved!=0) —
     which can collide with the new ``weights_pool`` slab and
     OOM the activation profile.
  3. ``gc.collect`` to break any Python-level cycles.
  4. ``torch.cuda.empty_cache`` to actually free the segments
     to the driver.

The ``NamedMemPool`` wrapper itself stays alive in the engine's
registry so metric callbacks keep reporting a stable label set,
but its underlying ``MemPool`` is gone — subsequent ``.use()``
on the wrapper would error out (calling code shouldn't reach for
it after this point). Per the boot path, the wrapper is referenced
only by the registry; metric reporting reads ``allocated_bytes``
via :meth:`NamedMemPool.snapshot` which gracefully returns 0 for
a dropped pool.

What a weight-materialization route owes the boot.

Five routes put a model's weights into ``weights_pool``: the dense
allocate-then-fill bind, the donor share, the EXL3 direct-slab bind, and
the two compactors. They differ in where the bytes come from and in when
the destination view is handed out, and those differences are real. What
they do NOT differ in is the verdict they owe the boot: every one of them
reports whether the bytes it moved round-tripped, and every one of them
is refused when they did not.

The two bind protocols below are separate because the routes really do
split on WHEN the slab view is handed over:

  * :class:`MetaBindingBinder` PUSHES — ``bind_all`` walks the model and
    rebinds every meta tensor before a byte is read, so the loader fills
    a view that is already in place.
  * :class:`SlabPlanningBinder` PULLS — ``plan`` sizes and allocates the
    slabs, and each filler asks for its own view by ``(module, attr)`` at
    the moment it has the bytes. The EXL3 buffers do not exist as meta
    tensors at bind time, so there is nothing to walk and rebind.

Collapsing those into one protocol would only describe the difference,
not remove it. The verdict is what they share, and
:func:`assert_weights_bit_identical` is the one place it is judged.

Refuse a boot whose weights did not survive the move into the pool.

Every weight-materialization route is gated here, so the check cannot
be present on one route and absent on the next. A route that reports
no verdict is refused for that reason: a missing key is not evidence
of an identical copy, and serving wrong weights is silent.

Stream every

The legacy cold quant route is *load → compact*: the EXL3 backend reads
the full ``trellis / suh / svh`` set off disk into standalone GPU
tensors (the "source"), and only then does
:func:`stream_compact_per_layer` pack them into per-dtype slabs and
free the source. Source and slab coexist during the copy, so the load
peak roughly doubles; the streaming compactor's host-staging guard and
progressive free claw some of that back, but the source set still
fully materializes before being freed, and the compaction copy is pure
overhead.

This binder removes both costs. It mirrors the warm flat-dump path
(which mmaps the pre-packed slab directly): load once, into the slab.

  * :meth:`plan` — runs after the quant swap (every EXL3 mirror linear
    exists with its final local ``in_features / out_features`` but empty
    placeholder buffers) and before any weight byte is read. Walks the
    model in the same order and grouping the streaming compactor uses
    (:func:`_iter_persistent_tensors_qualified` → ``(dtype, layer_group)``)
    and computes each persistent tensor's final numel:

      - dense meta params (embed / norms / dense lm_head): ``t.numel()``
        of the meta tensor; the slab dtype is the engine dtype for a
        float tensor, else the tensor's own dtype (mirrors
        :class:`AllocateThenFillBinder`).
      - EXL3 buffers (``trellis / suh / svh / mcg / mul1``, registered
        empty at swap time): the post-TP-slice numel, derived from the
        linear's local in/out features + the on-disk trellis ``K`` (see
        :meth:`_exl3_buffer_plan`). No index math is reproduced — only
        final sizes, which a test pins equal to the bound shapes.

    then allocates one slab per dtype in ``dst_pool`` and records each
    tensor's ``(dtype, offset_elems, shape)``.

  * :meth:`slab_view_for` — the EXL3 ``exl3_load`` and the dense
    meta-fill seam call this with a (module, attr) pair to obtain the
    pre-allocated slab view they must write into. ``exl3_load`` copies
    the (TP-sliced, contiguous) transient into the view and binds the
    buffer to it; the transient frees immediately. The dense seam binds
    the meta param to the view so the weight loader fills it in place.

  * :meth:`verify` — re-hashes every bound slab view after fill against
    the post-fill signature and checks the per-dtype cursor advance,
    reusing the streaming compactor's bit-identity + cursor gate so a
    wrong offset / dtype still fails the boot loudly.

Slab packing order does not affect warm-reload correctness — the flat
dump captures ``state_dict()`` per named tensor, so each tensor's bytes
must be identical (same data, dtype, shape) but their offset within the
slab is a private detail. We keep the compactor's ``(dtype,
layer_group)`` ordering anyway so a side-by-side slab dump matches.

Return ``{buffer_name: final_numel}`` for one EXL3 linear.

The final (post-TP-slice) numels are a pure function of the
linear's LOCAL feature dims and the trellis codebook width ``K``:

  - ``trellis``: ``(in_local//16) * (out_local//16) * (K*16)``
  - ``suh``:     ``in_local``
  - ``svh``:     ``out_local``
  - ``mcg`` / ``mul1``: the on-disk scalar-seed numel (shared, never
    TP-sliced), or ``0`` when the codebook is the default.

``in_local`` / ``out_local`` are the EXL3 linear's stored
``in_features`` / ``out_features``, which the linear classes already
set to the rank-local shard at construction — so this reproduces the
sliced shape without re-deriving the per-rank index ranges. A test
pins these numels equal to the shapes a real ``exl3_load`` binds.

Record one

Called from the quant swap loop (``apply_quant_if_present``) right
after the mirror linear is constructed — before any trellis byte is
read. ``trellis_k`` / ``mcg_numel`` / ``mul1_numel`` come from the
on-disk safetensors header (a microsecond read), so planning never
materializes weight data.

Return ``(embed_module, qualified_name)`` for the fp8 input embedding,
or ``(None, "")`` when ``cfg.embed_quant`` resolves ``off`` on this model.

The mode is settled by
:func:`~arbi_serve.weight_quant.embed_quant.resolve_embed_mode`, the SAME
decision the eager and warm paths take — so ``auto`` lands fp8 here
exactly where it lands fp8 there (never a silent ``off``), and a NAMED
mode on an ineligible model (absent / non-:class:`VocabParallelEmbedding`
/ TIED embedding) fails the boot loudly rather than falling back to a
bf16 slot.

Yield ``(mod, attr, qname, tensor)`` for every tensor the binder
will fill into a slab, in the streaming compactor's walk order.

That is exactly: (a) EXL3 buffers registered via
:meth:`register_exl3` (empty placeholders at plan time, sized from
the registered plan, filled by the deferred bind), and (b) dense META
parameters the weight loader will fill in place (embeddings / norms /
dense lm_head). Everything else is left untouched:

  - already-materialized tensors (arch-built persistent buffers like
    RoPE caches) keep their own storage — the dense meta-fill seam
    never rebinds them, so slab-binding them would orphan a
    never-filled view;
  - non-registered EMPTY placeholders (an absent mcg/mul1 seed) carry
    no data — skipped so they stay empty buffers, matching today.

De-dup: materialized EXL3 buffers by ``data_ptr``; meta params by
``(id(mod), attr)`` (a meta tensor's ``data_ptr`` is not unique).

Return the pre-allocated slab view for ``mod.<attr>`` reshaped to
``shape``, or ``None`` if the binder did not plan this tensor.

The caller (EXL3 ``exl3_load`` for quant buffers, the dense meta-fill
seam for dense params) writes its bytes into the returned view. A
shape mismatch against the plan fails loud — the planner's size and
the bind-time size must agree or a wrong offset would corrupt the
neighbouring tensor.

Slice a fused safetensors tensor into per-sub-tensor chunks.

A fused tensor packs multiple logical sub-projections along one axis
(typically the output dim, axis 0 for column-parallel weights stored
as ``(out, in)``). Each sub-projection has a distinct ``shard_id``
that maps to the corresponding chunk of a
:class:`MergedColumnParallelLinear`'s local weight.

The naive approach of "slice the whole fused tensor by tp_rank along
dim 0" is wrong when the fused output stacks ``[A | B]`` and TP is 2:
rank 0 would receive all of A, rank 1 all of B, rather than each
holding the per-shard A_local + B_local stacks the merged-linear
expects.

This module implements the correct path: slice the fused tensor by
the per-sub-tensor widths first, then dispatch each sub-slice through
:meth:`MergedColumnParallelLinear.weight_loader` with its own
``shard_id`` — the merged linear's existing per-rank slicing inside
``weight_loader`` (see ``linear.py``) handles the TP cut.

Split ``full`` into per-shard slices keyed by ``shard_id``.

Args:
    full: the fused source tensor read from safetensors.
    sub_tensors: declared per-sub-tensor widths + shard_ids in
        stack order along ``split_dim``.
    split_dim: axis along which sub-tensors are stacked.

Returns:
    ``{shard_id: slice_tensor}`` — each slice has the declared
    sub-tensor's width along ``split_dim`` and the source tensor's
    full extent on every other axis.

Raises:
    ValueError: if the sum of declared widths does not match
        ``full.shape[split_dim]``, or if ``split_dim`` is out of
        range for ``full``.

Slice + bind a fused tensor into a :class:`MergedColumnParallelLinear`.

Iterates the sub-tensor declarations, slices ``full`` accordingly,
and calls ``target.weight_loader(slice, shard_id=sub.shard_id)``
once per sub-tensor. The merged linear's own per-rank TP slicing
runs inside each ``weight_loader`` call.

Use this when the safetensors checkpoint stores the merged linear
as a single fused tensor (Qwen 3.5 q_proj's ``[Q | gate]``,
Mamba's ``in_proj``'s ``[x | z]`` — and any future arch that
stacks multiple logical projections into one tensor).

Args:
    target: the destination merged-column-parallel linear.
    full: the fused safetensors tensor.
    sub_tensors: per-sub-tensor widths + shard_ids.
    split_dim: axis along which sub-tensors are stacked. Must be
        ``0`` for column-parallel layouts (where the fused stack
        is along the output dim of ``(out, in)``).

Safetensors collection — directory index with on-demand tensor load.

A thin wrapper over the standard ``safetensors`` Python library: scan
a directory of shards, build a ``key -> shard`` index, and resolve
``get_tensor(key)`` lazily.

Raw on-disk location of one tensor inside its safetensors shard.

A chunked DMA worker mmaps ``shard`` once and reads
``[file_offset, file_offset + nbytes)`` straight into a pinned host
buffer, then copies host→GPU. No safetensors materialization, no
dtype cast, no Python tensor allocation on the read path.

Fields:
    shard: absolute path to the shard file holding the tensor.
    dtype_tag: the safetensors wire dtype tag (``"BF16"`` etc.).
    shape: tensor shape from the shard header.
    file_offset: absolute byte offset of the tensor's first byte in
        ``shard`` (already includes the 8-byte length prefix + JSON
        header — caller seeks here directly).
    nbytes: number of raw bytes the tensor occupies on disk.

Index + on-demand load for a directory of .safetensors shards.

Scans every ``*.safetensors`` in ``directory`` once at construction
and builds a flat ``{tensor_key: shard_path}`` map. ``get_tensor``
opens the relevant shard, reads the slice for the requested key,
and returns it on the requested device with optional dtype cast.

The HF model.safetensors.index.json is consulted opportunistically
(it's cheaper than reading every shard's header) but the live
index built from headers is authoritative.

Resolve a requested key to its stored form via the optional
alias (e.g. strip the ``model.`` prefix for ST checkpoints).

Returns ``key`` unchanged when it already exists or no alias
applies — callers handle a genuine miss.

Install a ``model.``-prefix alias for sentence-transformers checkpoints.

ST embedding / reranker checkpoints (Qwen3-Embedding,
Qwen3-Reranker exported via sentence-transformers) save the inner
transformer WITHOUT the ``model.`` prefix that the ForCausalLM
param names carry, and omit the tied ``lm_head``. When that
layout is detected, install an alias that resolves requested
``model.X`` keys to the bare ``X`` stored key; the tied
``lm_head.weight`` (whose source key is ``model.embed_tokens.weight``)
resolves through the same alias. Returns True if installed.

Return True if any indexed key starts with ``prefix``.

Quant-agnostic projection detection: a dense checkpoint stores
``<linear>.weight`` while a weight-quantized one stores packed
tensors under the same linear path (``.trellis`` / ``.qweight``
/ ``.weight_scale`` / ...). Callers that only need "is there a
projection at this path, in any quant format" pass the
``<linear>.`` prefix instead of a specific ``.weight`` key.

Sum the on-disk bytes of every tensor whose key starts with ``prefix``.

Returns ``(count, total_bytes)``. Reads ONLY the safetensors shard
headers (cached per shard) — each tensor's byte length comes straight
from its ``data_offsets`` span, so no tensor is materialized or moved
to a device. Used to report the VRAM a skipped sub-module (e.g. the
``visual.*`` vision tower under a text-only boot) would have cost,
without paying to load it.

The shard's open ``safe_open``, opening it on first use.

The mapping is read-only and the handles live as long as the
collection, so a caller that loads one tensor pays one open and a
caller that loads a hundred thousand still pays one.

Parse + cache one shard's safetensors header.

The on-disk layout is ``[u64 hdr_len | json hdr | raw data]``.
``data_offsets`` in the JSON are relative to the start of the
data section; the second return value is the absolute file
offset of that data section (``8 + hdr_len``).

Resolve ``key`` to its raw on-disk byte range.

Returns a :class:`TensorByteRange` a chunked DMA worker can use
to mmap the shard and copy the tensor's bytes straight into a
destination GPU view — bypassing ``safe_open`` /
``get_tensor`` (which allocate a Python tensor and serialize the
read through one Rust CPU thread).

Raises ``KeyError`` if the key isn't in the collection.

Return a zero-copy CPU view of a tensor's mmap-backed file bytes.

The mapping is retained by this collection so callers can keep tensor
views without concatenating large logical tables. No dtype conversion,
device transfer, or contiguous copy is performed.

Shared safetensors wire-format primitives.

One home for the on-disk safetensors details used across the loader:

  * the **bidirectional dtype map** (wire tag ↔ ``torch.dtype``),
    including the float8 quant dtypes (FP8 dynamic / NVFP4 scales);
  * the **header pack / parse** pair for the
    ``[u64 hdr_len | json hdr | raw data]`` container (the 8-byte
    alignment pad rule lives in one place, symmetric for write & read);
  * a small **bit-signature hex** helper for the dump's per-tensor
    content manifest.

Serialize a safetensors header dict to the on-disk header bytes.

Returns ``struct.pack("<Q", hdr_len) + padded_json`` — the u64
little-endian header length followed by the JSON header padded with
spaces so the data section that follows starts 8-byte aligned (the
header length itself is a multiple of 8). Inverse of
:func:`parse_safetensors_header`.

Read the safetensors header JSON + return (header, data_offset).

The on-disk layout is ``[u64 hdr_len | json hdr | raw tensor data]``.
``data_offsets`` in the JSON are relative to the start of the data
section, not the file — caller adds the returned ``data_offset``
(``8 + hdr_len``) to get an absolute file position. Inverse of
:func:`pack_safetensors_header`.

Per-parameter shard spec for the generic safetensors loader.

Each model declares ``weight_map() -> dict[str, WeightShardSpec]``
mapping its ``named_parameters()`` keys to a :class:`WeightShardSpec`
that tells the loader:

  - which safetensors key to read (``src_key``);
  - whether to slice along an axis by ``tp_rank`` (``shard_dim``);
  - which sub-chunk of a fused weight to bind (``shard_id``);
  - which expert to pick out of an MoE ``(n_experts, ...)`` tensor
    (``src_expert``);
  - or, for fused safetensors tensors that pack multiple
    column-parallel sub-projections along one axis (e.g. Qwen 3.5's
    ``q_proj`` which stacks ``[Q | gate]``, or HF Mamba's ``in_proj``
    which stacks ``[x | z]``), declare the per-sub-tensor split via
    ``fused_split`` so the loader can slice the fused tensor and
    dispatch one ``weight_loader(slice, shard_id=...)`` call per
    sub-tensor.

The loader walks the map; for each (param_name, spec) it slices
the source tensor and calls ``param.weight_loader(slice,
shard_id=spec.shard_id)`` for parallel linears or ``param.copy_(slice)``
for plain params. When ``fused_split`` is set, the loader iterates the
declared sub-tensors and invokes ``weight_loader`` once per shard_id.

One sub-tensor inside a fused safetensors weight.

Used as an entry in :attr:`WeightShardSpec.fused_split` to declare
how a single fused safetensors tensor decomposes into multiple
logical sub-projections that each map to one ``shard_id`` of a
:class:`MergedColumnParallelLinear`.

Args:
    shard_id: the ``shard_id`` to pass to the target's
        ``weight_loader`` for this slice. Must match one of the
        ``shard_ids`` declared at :class:`MergedColumnParallelLinear`
        construction.
    width: number of elements along ``split_dim`` this sub-tensor
        occupies in the fused source tensor.

Layout convention: sub-tensors are stacked along ``split_dim`` in
declaration order. Cumulative offsets are computed from the
preceding entries' widths.

How to bind one parameter from the safetensors collection.

Args:
    src_key: safetensors key to read. Use the model's HF naming
        convention (e.g. ``model.layers.0.self_attn.q_proj.weight``).
    shard_dim: axis along which to slice by ``tp_rank``.
        ``None`` = replicated (no slicing). Mutually exclusive with
        :attr:`fused_split` — when ``fused_split`` is set the loader
        handles per-sub-tensor TP slicing inside each shard rather
        than a single whole-tensor cut.
    shard_id: qualifier for fused parallel linears
        (``MergedColumnParallelLinear``: ``"gate"`` / ``"up"``;
        ``QKVParallelLinear``: ``"q"`` / ``"k"`` / ``"v"``). Used
        when the safetensors stores each sub-projection as a
        separate key — the model declares one ``WeightShardSpec``
        per ``shard_id`` and the loader binds them into the fused
        local weight one at a time.
    src_expert: for MoE — which expert's slice to READ out of a
        fused ``(n_experts, ...)`` SOURCE tensor (symmetric with
        ``dst_expert``, which names the slot to WRITE in a stacked
        destination).
    fused_split: tuple of :class:`FusedSubTensor` entries
        describing how a single fused safetensors tensor splits
        into multiple logical sub-projections (each sub-projection
        maps to one ``shard_id`` of the target
        :class:`MergedColumnParallelLinear`). When set, the loader
        slices the source tensor along :attr:`fused_split_dim`
        into the declared widths and dispatches one
        ``weight_loader(slice, shard_id=sub.shard_id)`` call per
        sub-tensor — each of which handles its own per-rank TP
        slicing internally. Mutually exclusive with ``shard_dim``
        and ``shard_id`` (the per-sub-tensor entries carry the
        shard_id).
    fused_split_dim: axis of the source tensor along which the
        sub-tensors are stacked. Defaults to ``0`` which is the
        convention for column-parallel projections (output dim
        stacked along dim 0). Ignored when ``fused_split`` is None.
    dst_param: the ``named_parameters()`` key this spec binds INTO,
        when it differs from the map key. Set it when several
        checkpoint tensors fill disjoint slices of ONE destination
        parameter — a stacked MoE expert tensor is filled by
        ``n_local_experts × 3`` separate checkpoint keys, and a dict
        cannot hold three thousand entries under one parameter name.
        The map key is then a unique binding id (which must NOT end
        in ``.weight``, so the quant filter keeps treating it as
        opaque) and this field names the real destination.
    dst_expert: index along dim 0 of the destination parameter to
        write into — the expert slot of a stacked ``(E, ...)`` MoE
        weight. Selects ``param[dst_expert]`` as the destination view.
    dst_offset: offset along dim 0 of the ``param[dst_expert]`` view
        at which the source is written; the width is the source's own
        dim-0 extent. This is how gate (offset 0) and up (offset N)
        land in the single ``w13`` tensor. Requires ``dst_expert``.
    shard_axis: which per-rank split ``shard_dim`` is indexed by.
        ``"tp"`` (default) = the whole TP group
        (``tp_rank`` of ``tp_size``) — every dense projection.
        ``"moe_tp"`` = only the ranks of the TP group that share this
        rank's expert shard (``moe_tp_rank`` of ``moe_tp_size``) —
        the MoE intermediate dim, which under
        ``--enable-expert-parallel`` is split by FEWER ranks than the
        attention weights because the rest of the TP group is holding
        different experts instead. The two coincide whenever
        ``moe_ep_size == 1``, so declaring the axis is free for every
        pre-existing topology.
        ``"attn_tp"`` = only the ranks of the TP group that share this
        rank's REQUESTS (``attn_tp_rank`` of ``attn_tp_size``) — the
        attention head dim, which under ``attn_dp_size > 1`` is split
        by FEWER ranks than the FFN weights because the rest of the TP
        group is holding different requests instead. It too coincides
        with ``"tp"`` at the default ``attn_dp_size == 1``.
    keep_src_dtype: read the source tensor at its ON-DISK dtype
        instead of casting to the engine dtype. Required for a
        quantized payload (fp8 bytes) and its scale grid, where a
        cast to bf16 would corrupt / inflate the tensor.
    transform: pure ``tensor -> tensor`` repack applied to the source
        right after the read, for a checkpoint whose on-disk
        container differs from the destination's — an AWQ int32
        nibble grid feeding the fused-MoE kernel's uint8 expert
        stack. It runs BEFORE the ``shard_dim`` cut, so ``shard_dim``
        indexes the TRANSFORMED tensor's axes and the shard extent is
        counted in the destination's storage units.

Return whether the skip-weight-load flag is currently set.

Live accessor — reads the module global at call time so callers in
other modules see the current value (a plain ``from .weights import
_SKIP_WEIGHT_LOAD`` would capture the value at import time and never
update). The warm flat-dump path uses this to skip the quant
backend's safetensors bind (the dump supplies the bytes instead).

Fill the straight-copy tensors via chunked multi-stream DMA.

Returns ``{weight_map name: bytes filled}`` for what was handled here;
the caller skips those in the per-tensor loop. Returns an empty dict
(a no-op) unless every safety precondition holds:

  * a CUDA device (the multi-stream pinned-DMA engine is CUDA-only);
  * the destination param is already real, contiguous CUDA storage
    (the cold-boot allocate-then-fill binder bound it to a slab
    view — on the legacy load→compact path params are still on meta
    here, so this returns empty and the per-tensor loop runs as
    before);
  * the spec needs NO SOURCE transform (``shard_dim`` / ``shard_id``
    / ``src_expert`` / ``fused_split`` all ``None``) — TP slice, fused
    QKV split, and source-side MoE expert slice stay on the per-tensor
    path. A ``dst_expert`` DESTINATION slice IS eligible: it selects a
    contiguous sub-view of a stacked MoE parameter and leaves the
    source untouched, so the copy is still raw bytes;
  * the source shard's on-disk dtype EXACTLY matches the
    destination dtype (no fp32→bf16 cast);
  * source byte length matches the destination view's byte length.

Anything that fails a check is simply left out of the DMA batch and
handled by the existing per-tensor loop — correctness first.

Return this rank's slice of ``full`` along ``shard_dim``.

``shard_axis`` selects WHICH per-rank split indexes the cut — the
whole TP group (``"tp"``, every dense projection), only the ranks
sharing this rank's expert shard (``"moe_tp"``, the MoE intermediate
dim under ``--enable-expert-parallel``), or only the ranks sharing
this rank's requests (``"attn_tp"``, the attention head dim under
``attn_dp_size > 1``). See
:class:`~arbi_serve.loader.shard_spec.WeightShardSpec`. All three
coincide whenever ``moe_ep_size == 1`` and ``attn_dp_size == 1``.

The divisibility check fires loudly rather than silently handing a
rank a ragged shard.

Bind every parameter in ``model.weight_map()`` from ``stc``.

For each ``(param_name, spec)``:

  1. Read ``spec.src_key`` from the safetensors collection. (For
     tied weights — e.g. lm_head when not present — fall back to
     the source-aliased key the spec carries.)
  2. If ``spec.src_expert is not None`` (MoE), slice the
     ``(n_experts, ...)`` tensor at the expert dim.
  3. If ``spec.dst_expert is not None``, copy into that expert
     slot's row range of the stacked MoE parameter named by
     ``spec.dst_param`` — many checkpoint tensors, one destination.
  4. If ``spec.fused_split`` is set, slice the source tensor by the
     declared per-sub-tensor widths and dispatch one
     :meth:`MergedColumnParallelLinear.weight_loader` call per
     shard_id (each handling its own per-rank TP cut internally).
  5. Otherwise, if the target parameter belongs to a
     :class:`LinearBase`, dispatch through its ``weight_loader``
     which handles per-rank sharding (and ``shard_id`` for fused
     linears) internally.
  6. Otherwise the parameter is plain — replicated by default, or
     per-rank-sliced along ``spec.shard_dim`` if set (used for
     non-Linear TP-aware buffers like depthwise conv1d kernels and
     per-head SSM scalars).

When :data:`_SKIP_WEIGHT_LOAD` is set, returns immediately — the
caller (instant-restart resume path) is responsible for filling
parameter storage via :func:`arbi_serve.loader.flat_dump.load_flat_weights`.
``ignore_skip_flag`` overrides that for a caller that fills a SUBSET
itself: the donor-share path skips the bulk load (the donor's bytes are
the weights) and then fills only the params the donor does not cover.

``only`` restricts the pass to the named DESTINATION params
(``spec.dst_param or name``); every other ``weight_map`` entry is left
untouched. ``None`` (the default) binds the whole map.

Architecture registry.

Each known ``architectures[]`` value in ``config.json`` maps to a
:class:`ModelBase` subclass via the single declarative
:data:`_ARCH_TABLE` below. **Adding a model is one table row** — no
separate edits to ``__all__``, the lazy-import map, or a builder
function (they are all derived from the table).

This package's ``__init__`` is intentionally torch-clean: the per-arch
model modules all import torch at module load, so eagerly importing them
here would force every caller of
``from arbi_serve.models.layer_spec import StateKind`` to pay the torch
tax (and break the torchless ``dump_openapi`` CI smoke). The table holds
only **strings** (module path + class name + a short "not yet wired"
reason for stubs); modules are imported lazily on first
:func:`load_model` / first attribute access via a PEP 562
``__getattr__``.

To add an architecture:
  1. Write ``arbi_serve/models/<arch>.py`` (copy the nearest sibling;
     dense decoders subclass ``DenseDecoderConfig`` + mix in
     ``LayerStackModelMixin`` — see ``docs/adding-a-model.md``).
  2. Add ONE row to :data:`_ARCH_TABLE`. Done.

One registry row: where the impl lives, or why it's a stub.

``module`` / ``cls`` name the concrete :class:`ModelBase` subclass
(imported lazily). When ``stub_reason`` is set the arch is recognised
but unimplemented — :func:`load_model` raises a precise message
instead of a generic "unsupported".

Build the ``<name>_NotYet`` placeholder class for a stub row.

Reuses :func:`arbi_serve.models.base._stub_arch` so a stub arch in
the ``ARCHITECTURES`` map is a class with the same ``from_safetensors``
classmethod surface as the live ones (it raises ``NotImplementedError``
naming the missing feature) — never a bare ``KeyError`` (which reads as
"unknown architecture"). ``arch.cls`` holds the short display name.

The live model class declared by ``model_dir``'s config.json, or ``None``.

Weightless: reads the arch name and resolves the registry row. Stub rows and
unknown architectures return ``None`` — callers treat that as "no support".

The model's output-stream geometry, without loading any weights.

Process mode (``engine_proc``) runs the routes in an API child that holds no
model object, but the speech-out routes still need the ``OmniOutputSpec`` to
gate the request and seed a speech turn. The spec is a property of the
ARCHITECTURE, so resolve it from ``config.json`` + the registry, the same way
the child already resolves its generation defaults.

``None`` for text-only models, stub rows and unknown architectures — every
caller treats ``None`` as "no audio out".

The model's INPUT modality bindings, without loading any weights.

The process-mode API child holds no model object but still has to expand
media placeholders before the request crosses the wire. Bindings are a
property of the ARCHITECTURE — same trick as :func:`omni_output_spec_for_dir`.

``{}`` for text-only models, stub rows and unknown architectures; every
caller treats ``{}`` as "no multimodal input".

Does ``model_dir``'s architecture drive the bespoke NemotronVoiceChat
turn loop (``arbi_serve/runtime/nemotron_voicechat_turn.py``'s
``run_turn``) instead of the standard omni turn streamer
(``arbi_serve/realtime/turn.py``'s ``stream_turn``)?

Process-mode weightless resolution — same "config.json + registry, no
weights" trick as :func:`omni_output_spec_for_dir` /
:func:`mm_bindings_for_dir`, for the realtime WS layer's turn-streamer
dispatch (:mod:`arbi_serve.realtime.nemotron_voicechat_turn`). ``False``
for every other architecture, stub rows, and unknown architectures.

Dispatch to the architecture-specific loader by ``config.json``.

Reads ``architectures`` and looks up the first match in
:data:`_ARCH_TABLE`. Unknown → raise with the supported list; a
recognised-but-stub arch → raise with its specific reason.

Batch-shape-invariant linear for accept-critical numerics.

Greedy/one-hot speculative-decode acceptance compares argmaxes computed
at different batch shapes: c=1 decode runs every projection (and the
lm_head) as a 1-row cuBLAS GEMV while the K+1 verify forward runs the
same weights as a (K+1)-row GEMM. cuBLAS picks a different kernel — a
different reduction order — for the 1-row shape, so the same token can
produce different logits bits depending on which path computed it. On
a bf16 model a near-tie top-2 gap can then flip argmax between the
mtp_k=0 and mtp_k>0 arms.

The fix: run every accept-critical matmul through one pinned GEMM
template. Rows below the pin (``ARBI_ACCEPT_INVARIANT_ROWS``, default
8) are zero-padded up to the pin and the result sliced back. Zero rows
are inert (each output row is an independent dot product), so real
rows are bit-identical to the same rows in any other <=pin-row call —
c=1 decode, K+1 verify, and the MTP draft chain all hit the same
kernel with the same per-row reduction order. This also makes
rank-invariance fall out for replicated drafters: the same pinned
kernel on identical replicated inputs gives identical bits on every
rank.

Gated by ``ARBI_ACCEPT_INVARIANT`` (default OFF — the fallthrough is
the exact previous ``F.linear`` call). Padding never changes
semantics, only the launch shape.

True when the invariant-accept mode is on (``_PIN_ROWS > 0``).

Used by non-linear accept-critical ops (e.g. the GDN verify recurrence)
to select their shape-invariant route. Reads the module global so it is
Dynamo-safe and monkeypatch-able alongside ``_PIN_ROWS``.

``F.linear`` for a 2-D ``x2``, writing the result into ``out``.

``at::linear`` on a 2-D input IS ``addmm(bias, x, w.t())`` (biased) or
``mm(x, w.t())`` (unbiased) — this dispatches to the same two ops with
the caller's destination, so the cuBLAS kernel, its reduction order and
therefore every output bit are the same as ``F.linear``. ``w.t()`` is a
stride view, not a copy: cuBLAS consumes the transpose as an op flag.

INFERENCE ONLY. torch refuses ``out=`` when grad mode is on and an
operand requires grad (writing a caller's storage has no autograd
formula). Every forward in this server runs under
``torch.inference_mode()`` — including the capture warmup and the
recorded region — so the destination path is only ever reached there,
and torch's own refusal is the guard for anyone who calls it elsewhere.

``F.linear(x, weight, bias)`` that is bit-invariant to the row count.

With ``ARBI_ACCEPT_INVARIANT=1`` and ``x`` carrying at most
``ARBI_ACCEPT_INVARIANT_ROWS`` rows (leading dims flattened), the input
is materialized into a contiguous 2-D ``(pin, in_features)`` buffer —
zero-padded when it is shorter — so cuBLAS always sees the same GEMM
shape AND the same operand layout: same kernel, same reduction order,
bit-identical real rows across the {c=1 decode, K+1 verify, draft
chain} shape family. A row count EQUAL to the pin goes through the
buffer too, even though its pad is empty: the pin's whole contract is
that ``K + 1 == pin`` is the largest covered verify shape, and handing
cuBLAS the caller's own (possibly 3-D, possibly non-contiguous) tensor
there is a different call than the padded one the ``c=1`` decode makes.
Everything else (flag off, CPU tensors, rows > pin) is the plain
``F.linear`` — byte-identical to the pre-flag behaviour.

``bias`` rides inside the single ``F.linear`` call (cuBLAS epilogue)
rather than a separate eager add, so biased projections (Qwen2-style
QKV bias) keep the same rounding as the ``nn.Linear`` reference.
Padded rows receive the bias too, but they are sliced off before
return, so real rows are unaffected.

``out`` (destination-passing): write the result into the caller's
``(m, out_features)`` buffer and return it, instead of allocating a
fresh output. The lm_head epilogue uses this to make the captured
decode graph's full-vocab logits land in a pre-allocated persistent
buffer — the single largest term in the cudagraph capture pool is
otherwise one ``(B, vocab)`` allocation baked into EVERY captured
shape. Numerics are unchanged: the un-pinned path is the same
``mm``/``addmm`` ``F.linear`` itself dispatches to (see
:func:`_linear_into`), and the pinned path computes exactly as before
and copies the real rows in (the pad rows must not reach ``out``, so
destination-passing cannot fuse there).

The GEMM row pin for QUANTIZED accept-critical linears.

Returns the active pin (``ARBI_ACCEPT_INVARIANT_ROWS`` when the
invariant-accept mode is on, else ``0``). ``invariant_linear`` covers
the cuBLAS bf16 ``F.linear`` choke points, but the production 27B head
(and drafter GEMMs) are EXL3/FP8-quantized custom ops whose bsz==1 leg
dispatches a DIFFERENT kernel than the bsz>1 leg (EXL3: trellis
``bc.run`` vs ``exl3_gemm``; FP8: an ``M==1`` ``_scaled_mm`` fast path) —
the SAME shape-dependent-reduction divergence class as the GEMV/GEMM
cuBLAS split, which ``invariant_linear`` does not touch. A quant backend
reads this pin and, for a sub-pin row count, zero-pads up to the pin and
takes its MULTI-ROW leg, so the c=1 decode and the K+1 verify share ONE
kernel and produce bit-identical rows. Read as the Dynamo-safe module
global (``_PIN_ROWS``) — safe inside a custom op body (which runs eager,
opaque to Dynamo) and inside a compiled forward alike.

Shared pre-norm decoder-layer skeleton.

A transformer decoder block in the dense / MLA / MoE arches here runs
the same pre-norm residual pattern::

    residual = hidden
    h = input_layernorm(hidden)
    h = mixer(h)                  # attention / MLA
    h, residual = post-attn norm + residual add
    h = ffn(h)                    # dense MLP / MoE
    return residual + h

The arches differ only in which token mixer / FFN ``__init__`` wires,
how the per-step args are threaded into them, and whether the
post-attention norm fuses the residual add. This base captures the
residual bookkeeping once so the accumulation order (which affects bf16
numerics) is identical everywhere; subclasses set the ``_fused_post_norm``
flag, wire their submodules in ``__init__``, and override the small
``_mixer`` / ``_ffn`` hooks that unpack the arch-specific ``extras``.

The attention-DP boundary lives on this same seam. Under
``attn_dp_size > 1`` the mixer runs on this rank's OWN requests while the
FFN is parallel over the full TP group, so the FFN input is gathered from
the attention-DP peers and its output sliced back down — around ``_ffn``
and nothing else, because the residual belongs to the local rows. At the
default ``attn_dp_size == 1`` the gather and the slice are identity
returns that allocate nothing.

Pre-norm transformer decoder block skeleton.

Subclasses wire ``input_layernorm`` / ``self_attn`` /
``post_attention_layernorm`` / ``mlp`` in ``__init__`` and supply the
``_mixer`` / ``_ffn`` hooks (which unpack the arch-specific per-step
``extras`` tuple and call the wired submodules). The public
``forward`` — keeping its arch-specific named signature so existing
callers / tests bind by keyword — lives on the subclass and
delegates the residual skeleton to :meth:`_prenorm_forward`.

Allocate this layer's arena-backed transients, EAGERLY.

Called by the model's per-layer dispatch seam
(:meth:`~arbi_serve.models._layer_stack_model_mixin.
LayerStackModelMixin._dispatch_layer_args`) with the live
:class:`~arbi_serve.runtime.activation_arena.ActivationArena` and
this layer's real (un-traced) ``hidden``. Whatever this returns
takes the arena's place in the per-step ``extras`` tuple, so the
compiled block forward receives tensors and never an allocator.

The attention output buffer is the one transient that qualifies
today: it is large, it is written by the attention kernel and
consumed by ``o_proj`` inside the same layer, and nothing reads
it after the block returns.

Pre-norm residual skeleton shared by every arch.

Residual accumulation order is fixed here: the intermediate
``residual + mixer(norm(hidden))`` is materialized once and
reused as the residual for the final add.

The FFN is bracketed by the attention-DP gather / slice. Both are
identity returns on a trivial attention-DP group, so the
``attn_dp_size == 1`` sequence — every arch's today — is unchanged
op for op.

DeepSeek-V4 blocks: hyper-connections, latent attention, compressor, MoE.

:class:`DSv4DecoderBlock` assembles them into the body every layer of this
arch runs — the main stack's and the speculative stages' alike.

Topology of one layer, with ``hc_mult`` parallel residual streams carried
between layers instead of a single residual::

    h: (N, hc, D)
    ├─ hc_pre(hc_attn_*)  -> x: (N, D) + write weights + combine matrix
    │  attn_norm -> DSv4LatentAttention -> hc_post -> h
    └─ hc_pre(hc_ffn_*)   -> x: (N, D) + write weights + combine matrix
       ffn_norm  -> DSv4MoE            -> hc_post -> h

:class:`DSv4LatentAttention` scores 64 query heads against ONE 512-d latent
per token — the latent is both key and value, so the output is latent-width
and gets un-rotated on its RoPE tail before a grouped low-rank output
projection. Each query reads a per-query id list: the sliding-window ring
always, plus (on a compressed layer) entries of a pooled stream, selected by
:class:`DSv4Indexer` on the fine family and taken whole on the coarse one.

Per-layer state arrives as a :class:`DSv4State` view whose rows ARE the
id space: ``[0, window)`` is the ring, the rest the pooled stream. The
blocks never allocate it and never learn how it is backed — see
:mod:`arbi_serve.cache.dsv4_pool`.

A bias-free linear whose weight is held (and applied) in fp32.

The compressor's projections are bf16 on disk and fp32 in the
reference, which computes the whole pooling — projection, gate softmax,
weighted sum — in fp32. Holding the upcast weight is what makes that
exact; casting per call would allocate an fp32 copy of the weight on
every token batch instead.

Replicated by construction: these projections feed a per-token pooling
that every rank needs whole, so there is no axis to shard.

RMSNorm that scales in fp32, matching the reference's ordering.

The reference holds the (bf16-on-disk) weight in fp32 and multiplies the
normalized fp32 activation by it BEFORE casting back, so the product
never round-trips through bf16 mid-expression. Keeping the weight as an
fp32 buffer preserves both halves of that: the value the checkpoint
stored, and the order the arithmetic happens in.

The per-layer view a DSv4 attention block reads.

Supplied by :class:`~arbi_serve.cache.dsv4_pool.DSv4LayerState`. The
block never allocates it and never learns how it is backed — the
production store reserves virtual address space and maps physical
rows as contexts grow, and a test can hand over plain tensors of the
same shapes.

``kv_flat`` is ``(flat_slots, slot_bytes)`` addressed through
``kv_addr``: the ring at entries ``[0, window)`` and the compressed
stream after it, so a virtual id plus a slot IS an address.
``kv_backed`` is the per-slot count of rows physically backed, which
every id gather clamps into. ``pool_*`` are the fp32 pooling
accumulators.

The read/write mixing over ``hc_mult`` parallel residual streams.

Replaces a plain residual: :meth:`pre` reduces the streams to one
activation for the sublayer to consume (and returns the weights that
say how to write back), :meth:`post` scatters the sublayer's output
across the streams and mixes the old streams through a near-doubly-
stochastic combination matrix.

The three parameter groups are the checkpoint's ``hc_*_fn`` (the
projection producing every mixing logit), ``hc_*_base`` (its bias) and
``hc_*_scale`` (three scalars, one each for the read weights, the write
weights and the combination matrix). All are fp32 on disk and stay
fp32 here — they feed a Sinkhorn iteration whose divisions do not
tolerate bf16.

The FINAL reduction of the parallel streams, before the LM head.

Deliberately not a :class:`DSv4HyperConnection`: there is no write-back
at the end of the stack, so this carries only the READ weights — one
projection row per stream and a single scale, against the full
``(2 + hc) * hc`` rows and three scales a mixing layer needs. The
checkpoint sizes ``hc_head_fn`` accordingly, and building the larger
module here would leave two thirds of it unfilled.

Pools every ``ratio`` tokens into one cached entry by a learned gate.

``wkv`` projects the hidden state to the entry's value and ``wgate`` to
its pooling logits; ``ape`` adds a learned bias per position WITHIN the
group, so the gate can prefer (say) the group's last token. The pooled
entry is RMS-normed, RoPE'd at its group's FIRST position, and then put
through the quantize-dequantize round trip the checkpoint was trained
with.

``ratio == 4`` pools OVERLAPPING windows: ``wkv``/``wgate`` emit twice
the width, the first half feeding the window shared with the previous
group and the second half the current one, so consecutive entries see
``2 * ratio`` tokens. Both projections are bf16 on disk and computed in
fp32 — the softmax over a group is where a bf16 accumulation would
show.

Scores the pooled stream to pick which entries a query may read.

Its own :class:`DSv4Compressor` builds a narrower (``index_head_dim``)
pooled stream, and ``wq_b`` a query per index head. Both sides are
Hadamard-rotated and FP4 round-tripped before scoring, which is how the
checkpoint was trained — the scores are a 4-bit dot product by design,
not an approximation introduced here.

The score is a ReLU-gated, per-head-weighted sum over index heads, so a
head contributes only where it agrees with the entry; ``weights_proj``
supplies the per-head weight per token.

Latent MQA over a sliding window plus a selected pooled stream.

``wq_a``/``q_norm``/``wq_b`` build the per-head query through a low-rank
bottleneck; the query is RMS-scaled per head WITHOUT a learned weight
before RoPE. ``wkv``/``kv_norm`` build the single latent that serves as
both key and value; its NoPE channels take the FP8 round trip and its
tail takes RoPE.

The attention output is latent-width, so its RoPE tail is un-rotated
(the inverse rotation) before the output projection. That projection is
two-stage and GROUPED: ``wo_a`` is read as ``(o_groups, o_lora_rank,
heads_per_group * head_dim)`` and contracts each group of heads on its
own, then ``wo_b`` maps the concatenated ranks back to the hidden size.

Router + FP4 routed experts + one FP8 shared expert.

Two things separate this gate from every other MoE in the repo:

  * the score is ``sqrt(softplus(logits))`` — unnormalized and
    unbounded, so ``norm_topk_prob`` is what turns the selected scores
    into weights;
  * on the first ``num_hash_layers`` layers the expert ids are NOT
    scored. They are read from a per-token-id table (``tid2eid``), so
    routing is a property of the token, not of the hidden state. The
    weights still come from the scores of the ids the table names.

The routed experts are the MXFP4-shaped stack :class:`FusedMoE` already
serves; the shared expert is a per-linear FP8 MLP under this arch's own
``w1``/``w2``/``w3`` names. Both sides clamp their gate and lift
pre-activations at ``swiglu_limit`` — the saturating activation the
checkpoint was trained with, on every expert it routes to.

The shared expert: SwiGLU with the checkpoint's clamps, FP8 linears.

``w1`` gates, ``w3`` lifts and ``w2`` projects back — this arch's names,
which is what lets the FP8 swap bind them by path. The clamp is applied
to both halves before the product, matching the trained activation: the
gate is clamped ABOVE only, the lift on both sides.

Hyper-connected attention + MoE over the ``hc_mult`` parallel streams.

One body for both places this arch runs it: a main decoder layer, and a
speculative stage under ``mtp.N.*``. The stage subclasses this and adds
the entry and exit its checkpoint declares, so a change to the body
reaches both.

Attribute names are the checkpoint's, so a stage's parameters sit at
``mtp.N.attn.*`` / ``mtp.N.ffn.*`` exactly as the tensors are stored.

Assemble every group's pooling window, in one pass.

Groups are assembled by SCATTER rather than by walking sequences:
each token knows its sequence, its group within that sequence's run,
and its row within the group, so one ``index_copy_`` places the
whole batch. Group 0 of each run is seeded from the accumulator
first, so a group split across forwards pools with the tokens the
previous call left behind.

``valid`` is a per-token mask; a masked-off token scatters into a
sentinel cell nothing reads, which is how a partial-accept replay
pools the accepted prefix alone at the shape of the whole block.

Shapes come from ``(num_seqs, max_query_len)`` — the dimensions a
captured graph pins — never from how far any sequence has run.
Completion is a mask, not a branch.

Returns ``(window_kv, window_score, group_abs, completed)`` flattened
over ``(sequence, group)``; the window's rows are the tokens a group
pools over, ``-inf``-scored where it has none. The accumulator is
advanced to what the next call resumes from.

Index-space query per token: ``(N, n_heads, head_dim)``.

On CUDA the rotation, the Hadamard and the FP4 round trip are one
launch with the row resident in registers; off it they are the torch
ops those three stages are specified by.

Attend the whole flat batch in one pass.

Sequence identity is carried per TOKEN (``meta.slot`` /
``meta.run_start``), so every stage is one batched op over the flat
batch: no stage branches on a tensor's value, and the only
Python-level numbers are shapes. That is what makes the step
recordable — and it is also simply faster, since a per-sequence walk
would launch each projection once per request.

The one branch is on the batch's SHAPE — whether every run holds a
single token — which decides when the ring is published and so
which id space the attend addresses. A captured graph holds one
shape, so it holds one side of it.

Per-token id list into the addressable space.

Ids below the store's row count name a cached entry. With
``in_run``, an id at or above it names a row of THIS step's own
latents, which is where a position still in flight lives — the
ring is published after attending, so a run cannot read its own
tail out of it. Without it every id names a cached row, which is
what a run whose ring is already published gets.

Attend each token over the entries its id list names.

``cached_only`` says every id addresses the slab, so the step goes
through the backend seam — turbo-attn's split-K DSv4 kernel on
CUDA. A run that still carries its own in-flight latents has rows
no block table can reach, so it stitches those in and attends in
torch.

Copy the ring rows this block is about to overwrite.

Read BEFORE the publish, so the frame holds the committed latents
of the positions a window earlier — the ones a rejected draft would
otherwise take with it.

Commit each row's accepted prefix and nothing past it.

Puts back every ring row a rejected draft overwrote, and where a
compressor exists restores the accumulator to what the verify step
found and re-pools the accepted prefix through the SAME
:meth:`DSv4Compressor.advance_batched` the forward runs, with every
later token masked off. One pooling implementation, so the replay
cannot drift from the forward.

Put back the ring rows the block's REJECTED positions overwrote.

An accepted position keeps what the forward published. A block is
never wider than the ring (refused at capture), so its positions
occupy distinct rows and the two halves cannot collide — which is
what makes this one scatter rather than a host-side selection of
the rows to restore.

Step metadata describing the retained block as one run per row.

The run keeps its full ``depth`` shape and its LAST position is the
accepted one, which is what makes the replay a fixed-shape launch
set whatever each row accepted.

Publish each sequence's last ``window`` latents into its ring.

After attending, and masked to the tail, for two reasons that both
produce wrong numbers rather than slow ones: an earlier query of
the same run still needs the ring slots a later token would
overwrite, and writing every position of a longer run would hand
``index_put_`` duplicate indices, where which write lands is
undefined. The surviving tail has distinct slots by construction.

Advance one pooled stream over the whole batch.

On CUDA the FP8 stream's pool / norm / rotate / round-trip / pack /
scatter is one launch over one read of the pooling window. The
rotated stream the indexer builds finalizes through a Hadamard and
an FP4 round trip instead, which no fused kernel computes, so it
goes through the torch ops that specify it — as does every stream
off CUDA.

Run one block over ``(N, hc, D)`` streams.

Every stage — the hyper-connections, attention, the MoE — is one
batched op over the flat batch. Sequence identity travels as
per-token tensors on ``meta``, so prefill and decode share this
body and neither walks sequences on the host.

DeepSeek-V4 (``DeepseekV4ForCausalLM``) config parse → dims + layer specs.

Lifts the HF root ``config.json`` into :class:`DeepseekV4Config` once, and
derives :class:`ModelDims` + one :class:`LayerSpec` per layer from it. Every
knob the forward does not implement raises here rather than serving wrong
numbers.

**Three facts the config file does not carry, or carries wrongly.** The
parse takes the tensors' side:

  * ``num_nextn_predict_layers`` reads 1; the checkpoint carries THREE
    DSpark stages (``mtp.0`` / ``mtp.1`` / ``mtp.2``). The stage count is
    derived from ``compress_ratios``, which holds one entry per layer
    INCLUDING the DSpark stages (``len(compress_ratios) -
    num_hidden_layers``), and is cross-checked against the checkpoint's
    own ``mtp.N.`` prefixes when a key list is supplied.
  * ``expert_dtype: "fp4"`` sits OUTSIDE ``quantization_config``, so the
    routed experts are FP4 while everything the quant block describes is
    FP8. Both are read; a checkpoint declaring anything else is refused.
  * whether the speculative head owns its vocab tensors is stated NOWHERE
    in the config — two next-token checkpoints with identical fields
    differ, one shipping ``mtp.0.emb.tok_emb`` + ``mtp.0.head`` and the
    other reading the trunk's. :attr:`DeepseekV4Config.mtp_untied_vocab`
    is derived from the keys, and is ``None`` when none were supplied so a
    stage refuses to guess.

**RoPE is per-layer-family, and neither family takes YaRN's magnitude
term.** Layers with ``compress_ratios[i] != 0`` build their table on
``compress_rope_theta`` under the YaRN NTK-by-parts schedule; the
window-only layers (the first two, and every DSpark stage) build a plain
table on ``rope_theta`` with no scaling at all. The reference applies no
magnitude scaling to either table and keeps the softmax scale at a bare
``head_dim ** -0.5``, so :attr:`DeepseekV4Config.yarn` pins
``attention_factor=1.0`` (the ``0.1·ln(factor)+1`` HF default inflates the
whole cos/sin table by 27.7% at ``factor=16``) and ``mscale_all_dim`` stays
0 so nothing rides the softmax scale.

Does the speculative head ship its own embedding and LM head?

Answerable only from the tensors: a DSpark head always reads the
trunk's, while the two released next-token checkpoints differ — one
ships a full untied pair under ``mtp.0.*``, the other ships neither and
reads the trunk's. ``None`` says the parse was not shown the keys, so
a stage built on it would have to guess.

Per-arch construction bundle for :class:`DeepseekV4Model`.

``compress_ratios`` spans the main layers AND the DSpark stages, in
that order — index ``num_layers + stage`` is stage ``stage``'s ratio,
which is how the reference addresses it.

Build from a DeepSeek-V4 HF root ``config.json`` mapping.

``checkpoint_keys``, when given, cross-checks the derived
speculative stage count against the tensors that are actually
present, and settles whether the speculative head owns its vocab
tensors — which nothing in ``cfg`` states.

Validate the FP8 payload shape; return the activation scale format.

``scale_fmt`` describes the ACTIVATION round trips. The WEIGHT
scales' dtype is whatever the checkpoint stored — E8M0 on
Flash-0731, fp32 on Flash-Base, which declares ``ue8m0`` here
regardless — and the loader binds those at their on-disk dtype, so
neither fact silently overrides the other.

Build the compressed layers' YaRN schedule with no magnitude term.

``attention_factor=1.0`` is the pin that keeps the cos/sin table
unscaled: the reference's ``precompute_freqs_cis`` blends
interpolation with extrapolation and stops there, while the
HF-generic default would post-multiply the table by
``0.1·ln(factor)+1``. ``mscale`` / ``mscale_all_dim`` are refused
rather than honoured — a checkpoint that pins either wants a
magnitude term this forward does not apply.

One spec per main layer, plus one per DSpark stage when enabled.

A DSpark stage carries ``is_mtp_layer=True``: it exists in the list
so the stage's own window state is allocated, and the main decoder
loop skips it.

True iff ``layer_idx``'s gate reads its expert ids from the
per-token-id table instead of scoring a top-k.

Hash routing covers the first ``num_hash_layers`` MAIN layers; the
DSpark stages score their own top-k.

DeepSeek-V4's native KV slot layout — the trained quantization, stored.

The forward puts every cached latent's NoPE channels through a UE8M0-scaled
FP8 round trip, because that is what the checkpoint was trained to read. So
the values in the cache are ALREADY on the FP8 grid with a power-of-two
group scale, and storing them as bf16 stores an fp8 value in 16 bits.

:class:`DSv4KVLayout` stores what the round trip produced instead — the FP8
bytes, their E8M0 exponents, and the bf16 RoPE tail — which is
**bit-exact** to storing the round-tripped bf16 while costing 57% of the
bytes::

    slot = [ nope: head_dim - rope   bytes (E4M3)   ]
           [ scale: (nope / group)   bytes (E8M0)   ]
           [ pad:   to a 4-byte edge                ]
           [ rope:  rope * 2         bytes (bf16)   ]

At this arch's shape that is 448 + 7 + 1 + 128 = 584 bytes, padded to the
alignment the gather wants, against 1024 for a bf16 latent.

:class:`DSv4SlabAddress` is where those slots sit in a per-layer slab:
group-major over the request slots, entry-major inside a group, so each
group of slots is a contiguous byte range that is backed on its own.

This is a lossless STORAGE choice, not a KV-quantization option: nothing
here decides precision. The precision was decided by
:func:`~arbi_serve.models._deepseek_v4_ops.fp8_roundtrip` upstream, whose
output this encodes exactly. A separate bf16 layout exists only as the
parity reference the test suite diffs against.

Byte layout of one cached latent.

``group`` is the FP8 round trip's group width, so one stored exponent
covers exactly the channels one round-trip scale covered — which is
what makes encode/decode an identity on round-tripped input rather
than a re-quantization of it.

Flat addressing of a ``(groups, rows, group_slots, slot_bytes)`` slab.

Slots are cut into groups of :attr:`group_slots`; a group's bytes are
contiguous, so one group is backed independently of the others while the
whole slab keeps ONE base pointer. Entry ``e`` of slot ``s`` sits at flat
slot index::

    (s // group_slots) * rows * group_slots + e * group_slots + s % group_slots

which is arithmetic, not a table. It is not affine in ``s``, so the slab
has no ``(rows, slots)`` strided view; one slot's entries ARE affine, and
:meth:`slot_view` is that view.

Slots per independently backed group.

``min_prefix_bytes`` is what ONE slot of the store the partition is
derived for holds at the shortest context. The group is the smallest
whose one granule fits inside that prefix, so a group's first granule is
fully used and the granularity costs nothing. Group counts divide the
slots as evenly as the count allows, which is what keeps the tail group
from being mostly padding.

Byte offset of the bf16 RoPE tail, aligned to :data:`_ROPE_ALIGN`.

The exponents are one byte each and there is an odd number of them
at this arch's shape, so the tail would otherwise land on an odd
offset — where a device-side ``__nv_bfloat16*`` reinterpret of the
slot is unaligned and a vectorised load of it is a fault, not a
slowdown. Rounding up costs nothing: the slot is padded to
:attr:`alignment` anyway and the pad absorbs it.

Pack ``(..., head_dim)`` latents into ``(..., slot_bytes)`` uint8.

The NoPE channels are quantized with the same rule
:func:`fp8_roundtrip` uses, so encoding an already-round-tripped
latent is VALUE-exact — ``decode(encode(y)) == y`` to the bit.

It is not BYTE-exact, and the difference is not rare. Re-deriving
the scale from an already-quantized group can land on a different
power of two: ``pow2_ceil_scale`` puts a group's scaled amax in
``(224, 448]``, so an amax that lands in ``(224, 232]`` rounds back
down to 224, halving the scale and doubling every payload byte —
the same value written two ways. At this arch's shape that is ~8%
of groups. Anything reproducing this path (a fused device kernel,
say) has to run BOTH stages, not just the encoder.

The RoPE tail is copied as bf16 — it carries position, where a
4-bit-exponent grid loses angles.

Flat indices with ``entry`` clamped into ``slot``'s backed rows.

``backed`` is the per-slot backed row count. A clamped entry is one
the caller's own mask discards; naming an unbacked row would fault
the device.

DeepSeek-V4 speculative stages, stored under ``mtp.N.*``.

Two shapes ship under the same arch, and a checkpoint declares which by
:attr:`~arbi_serve.models._deepseek_v4_config.DeepseekV4Config.mtp_kind`:

  * :class:`DSparkStage` — drafts a whole block per forward and re-ranks it
    with a Markov bigram bias. The stages form a chain: the FIRST projects
    the tapped target hidden states down to the model width, the LAST
    carries the output norm and the ranking heads, and every stage borrows
    the main model's embedding and LM head.
  * :class:`DSv4NextTokenStage` — the classic next-token head. One stage
    predicts the token after the one the trunk just produced, from the
    trunk's hidden state and the embedding of that token. Whether its
    embedding and LM head are its own tensors or the trunk's varies by
    checkpoint, and only the checkpoint's keys say which.

Both are a :class:`~arbi_serve.models._deepseek_v4_blocks.DSv4DecoderBlock`
— window-only attention plus the MoE — with the stage's own entry and exit
around it. Which one a checkpoint gets is a registry row in
:mod:`arbi_serve.models.deepseek_v4`, never a branch at a call site.

The half every speculative stage shares.

A stage's MoE always scores its own top-k — the token-id routing table
covers main layers only — so the block is built with hash routing off
whatever the layer index would imply.

One DSpark speculative stage.

Structurally a decoder layer whose attention is window-only, plus the
stage-specific heads: the FIRST stage projects the tapped target hidden
states down to the model width, and the LAST carries the output norm,
the Markov bigram head and the confidence head. Its embedding and LM
head are the main model's, held by reference rather than re-registered,
so the checkpoint ships neither.

The drafting surface of a DSpark stage chain.

Owns no parameters — the stages, the trunk's embedding and the trunk's
LM head are all held by reference — so it composes onto a built model
without re-registering anything.

One draft proposes a whole BLOCK. Its ``block_size`` positions carry
the anchor (the last committed token) followed by ``noise_token_id``
placeholders, and the chain runs over the block AND over the committed
span the stages have not seen yet, whose entry is the target streams
the config's ``dspark_target_layer_ids`` names. Position ``j`` of the
block never sees the token realized at ``j - 1``, which is what
:class:`DSparkMarkovHead` puts back at sampling time.

A stage is window-only — the config refuses a DSpark stage that
declares a compress ratio — which is what lets a drafted block run
with nothing to roll back: it publishes no ring slot and there is no
pooling accumulator for its tokens to reach.

Where this differs from the reference DFlash drafter, and why: there
every layer draws its context keys and values from the SAME projected
tap, because its layers carry a context-projection seam. A DSpark
stage is a plain decoder block with no such seam — the tap enters at
stage 0's ``main_proj`` and each later stage reads the one before it,
for the context exactly as for the block.

Rank-``r`` bigram logit bias conditioned on the previous token.

The DSpark stages predict a whole block in one forward, so position
``k`` never sees the token sampled at ``k - 1``; this head adds that
dependency back as a factorized transition bias at sampling time.
Parameter names match the checkpoint's ``markov_w1`` / ``markov_w2``.

Scalar confidence per drafted position, in fp32.

Reads the stage's hidden state concatenated with the Markov embedding
of the token sampled there, so the score sees both what the block
predicted and what was actually taken.

The next-token speculative stage.

Entry is the two-stream mix the DeepSeek MTP heads take: the trunk's
residual streams through ``hnorm``, the next token's embedding through
``enorm``, each projected to the model width and summed. ``h_proj`` and
``e_proj`` are the two halves of that mix stored as separate tensors,
which is the same map as one projection over their concatenation.

Exit reduces the parallel streams through its ``hc_head`` and its
``norm`` before the LM head.

The two vocab tensors are the one thing that varies between
checkpoints of this shape: one release ships an untied
``emb.tok_emb`` + ``head`` pair per stage, another ships neither and
reads the trunk's. ``arch.mtp_untied_vocab`` — derived from the
checkpoint's own keys — decides which modules exist here, and a stage
built without that fact refuses rather than binding the wrong pair.

Run the chain and return ``(base_logits, block_hidden)``.

Each stage runs the committed span FIRST and the block second, so
the block's queries read a ring the same stage just published. The
block's own run is marked TRANSIENT, which publishes nothing: a
drafted position takes the ring row of the position a window
earlier, and nothing would put it back — the next draft re-runs the
committed span from wherever the accept left it.

Both returns are one row per BLOCK position, anchor included; the
caller drops the anchor's row, which predicts a token the target
already committed.

Per-position accept LOGIT, from the block hidden and its predecessor.

``markov_first`` flips the concat order the head's single Linear
reads, which is the layout the checkpoint was trained with and is
resolved empirically rather than declared anywhere.

Logits for the token after ``next_ids``, one row per input token.

``streams`` is the trunk's ``(N, hc, D)`` parallel residual streams
at the same positions, NOT their reduction: ``hnorm`` and
``h_proj`` run per stream and keep the stack's stream structure
into this block, while the embedding's projection is broadcast
across them. The reduction happens once, at the exit.

DeepSeek-V4 primitive math — the ops its blocks are built from.

Each function here mirrors one piece of the reference implementation's
numerics (``inference/model.py`` + the tilelang kernels in
``inference/kernel.py``) in plain torch, so the whole arch is executable
and diff-able on CPU:

  * :func:`fp8_roundtrip` / :func:`fp4_roundtrip` — the quantize-dequantize
    round trips the checkpoint was TRAINED with. They are part of the
    forward, not a storage choice: the latent's NoPE channels and the
    indexer's rotated query/stream are seen by the trained weights only
    after passing through them.
  * :func:`hadamard_transform` — the rotation applied before the indexer's
    FP4 round trip, spreading each channel's magnitude across the block so
    a 4-bit grid loses less of it.
  * :func:`hc_split_sinkhorn` — the hyper-connection mixing weights: a
    row-softmax followed by Sinkhorn normalization toward a doubly
    stochastic combination matrix.
  * :func:`sparse_latent_attend` — attention over a gathered, per-query id
    list with a per-head sink in the denominator only.
  * the index builders that say WHICH cached ids each query reads.

``-1`` is the pad id everywhere an id list appears: it contributes nothing
to the numerator and nothing to the denominator, so a list can be padded
to a static width without changing the result.

The UE8M0 scale for a group whose largest magnitude is ``amax``.

``2 ** ceil(log2(amax / grid_max))`` — a power of two, so dividing by
it is exact and the dequantized value carries no scale rounding error
of its own. This is what ``scale_fmt="ue8m0"`` means: the stored scale
is a bare exponent.

Quantize ``x`` to FP8 E4M3 per ``group`` of channels and back.

Returns a tensor of ``x``'s dtype and shape. The scale is the
power-of-two UE8M0 scale of each group (:func:`pow2_ceil_scale`); the
payload is clamped to the E4M3 range and rounded by torch's own cast,
which rounds to nearest-even like the device conversion does.

Round ``a`` onto the E2M1 grid, nearest-even at exact midpoints.

``a`` is already scaled into ``[-6, 6]``. Exact midpoints are common
rather than exotic here — the group scale is a power of two, so
dividing by it preserves the mantissa and lands values like ``0.75``
squarely between two codes; rounding those the wrong way is a
systematic bias, not a rounding detail.

Walsh-Hadamard transform of ``x``'s last dim, times ``scale``.

Sylvester-ordered (the ``H_2n = [[H, H], [H, -H]]`` recursion), which
is the ordering the reference's ``fast_hadamard_transform`` produces.
``scale`` defaults to ``1/sqrt(width)``, making the transform its own
inverse.

Split ``mixes`` into the hyper-connection ``(pre, post, comb)`` weights.

``mixes`` is ``(..., (2 + hc_mult) * hc_mult)``: the first ``hc_mult``
entries drive the read weights over the parallel residual streams, the
next ``hc_mult`` the write weights, and the remaining ``hc_mult ** 2``
the stream-to-stream combination matrix.

``comb`` is row-softmaxed and then Sinkhorn-normalized toward doubly
stochastic: one column normalization, then ``sinkhorn_iters - 1``
row/column pairs. Every division carries ``eps`` in the denominator, so
the iteration count and the epsilon are both part of the result — a
different schedule is a different mixing matrix, not a rounding
difference.

Returns ``(pre, post, comb)`` with shapes ``(..., hc)``, ``(..., hc)``
and ``(..., hc, hc)``, all fp32.

Rotate ``x``'s channel pairs ``(2i, 2i+1)`` by ``(cos, sin)``.

``cos`` / ``sin`` broadcast against ``x`` and hold one entry per
CHANNEL (each angle repeated across its pair), matching
:class:`~arbi_serve.models.layers.RoPECache` built with
``interleaved=True``. ``inverse=True`` rotates by the conjugate, which
is what un-rotates the attention output's tail.

Attend each query to the cached latents its id list names.

Args:
    q: ``(N, H, D)`` — one query row per token, ``H`` heads of width
        ``D``.
    kv: ``(S, D)`` addressable latents shared by every query, or
        ``(N, K, D)`` rows ALREADY gathered in ``ids`` order (which is
        what a per-request slot layout produces — gathering first keeps
        the ``(N, S, D)`` intermediate off the heap). Each latent serves
        as BOTH the key and the value, so the output shares its width.
    ids: ``(N, K)`` int — per-query ids into ``kv`` when ``kv`` is
        shared, or the validity pattern of the pre-gathered rows.
        ``-1`` pads.
    sink: ``(H,)`` fp32 — a per-head logit that enters ONLY the
        softmax denominator, so a query whose keys are all weak decays
        toward zero output instead of amplifying noise.
    scale: softmax scale on the logits.

Returns ``(N, H, D)`` in ``q``'s dtype.

Ids of the sliding window visible to each query position.

A position cached before this run lives in the ring at
``ring_base + (p % window)``. A position INSIDE this run is addressed
relative to the run instead, at ``run_base + (p - run_start)``, because
the ring is written for every token of the run before any of them
attends: token ``p + window`` overwrites token ``p``'s slot, so an
earlier query in the same run would read the later token's latent. The
run's own latents are still live in registers, so reading them there is
both correct and free.

That split is safe in the other direction too: a pre-run position a
query can still see cannot have been overwritten by this run, since the
overwriting token would have to sit more than ``window`` positions
later than the query itself.

Ids are OLDEST-FIRST and padded with ``-1`` before the start of the
sequence.

How many compressed entries exist once ``position`` is cached.

An entry covers ``ratio`` consecutive tokens and is written when the
last of them arrives, so position ``p`` (0-based) has
``(p + 1) // ratio`` entries behind it.

Ids of every compressed entry each query position may read.

Entry ``e`` covers tokens ``[e * ratio, (e + 1) * ratio)`` and becomes
visible only once complete, so query ``pos`` sees entries
``[0, (pos + 1) // ratio)``. ``width`` fixes the id list's extent — the
stream's CAPACITY, a property of the layer's shape rather than of the
batch — so a captured graph holds one id width whatever context its
rows are at.

Ids of the ``k`` highest-scoring entries per row, padded to ``k``.

``scores`` is ``(N, S)`` and ``valid`` the same shape; an invalid entry
can never be selected, and a row with fewer than ``k`` valid entries is
padded with ``-1`` rather than filled with its own repeats.

``sqrt(softplus(x))`` — this arch's router score function.

Unlike softmax it does not normalize across experts, and unlike
sigmoid it is unbounded above, so the gate's own normalization
(``norm_topk_prob``) is what makes the selected weights a
distribution.

Per-step token metadata for the DeepSeek-V4 forward.

Every per-sequence fact the attention needs — which slot a token belongs
to, where its sequence's run starts, how many pooling groups the run can
touch — is derived ONCE per step from data the host already holds, and
handed to the layers as device tensors of fixed shape.

Two properties this exists to guarantee, both of which the layers would
otherwise have to buy per layer and per sequence:

  * **No device→host reads.** Deriving run boundaries from
    ``cu_seqlens_q`` would sync the step and is illegal mid-capture. The
    scheduler knows the per-row lengths on the host, so this is built from
    those (or from :class:`~arbi_serve.engine.batch.HostBatchMirror`) and
    refuses loudly when neither is available — the same contract the MLA
    metadata builder holds for ``total_kv_tokens``.
  * **No per-sequence Python loop.** Sequence identity becomes a per-token
    tensor, so every stage of the layer is one batched op over the flat
    batch whatever the batch's shape.

Per-token sequence identity for one forward.

Args:
    slot: ``(N,)`` — the request slot each token's state lives in.
    run_start: ``(N,)`` — the first POSITION this forward holds for
        each token's sequence. A token is "in this run" iff its
        position is at or after it, which is what lets attention read
        the run's own latents instead of a ring slot a later token of
        the same run has already overwritten.
    run_first_index: ``(N,)`` — the flat index of that first token, so
        an in-run position maps to a row of this step's own tensors.
    run_length: ``(N,)`` — tokens this forward holds for the sequence.
    seq_index: ``(N,)`` — the token's row in the per-sequence tensors.
    slot_of_seq / run_start_of_seq / run_last_of_seq: ``(B,)`` — the
        same facts once per sequence, for the stages that work per
        sequence (pooling groups) rather than per token.
    num_seqs: sequences in the batch.
    max_query_len: longest run in the batch. Bounds every group buffer,
        so shapes depend on the batch's SHAPE and never on its content.
    verify_pass: this step's tokens are PROVISIONAL — a block whose
        tail the target may reject. The attention keeps a copy of every
        ring row the block overwrites (the ring is addressed modulo the
        window, so a rejected position destroys the committed one a
        window earlier and re-running never repairs it) and frames the
        pooling accumulators, so a rollback can undo the rejected tail.
    transient: this step's tokens are NEVER committed — a drafter's own
        proposal, thrown away and re-run from wherever the accept
        lands. The attention publishes nothing to the ring at all;
        there is no rollback for it to undo.

Derive :class:`DSv4StepMeta` from host-side batch metadata.

Everything comes from :class:`~arbi_serve.engine.batch.DSv4Meta`, which
the metadata builder fills from host-side scheduling data. A step
without it is refused rather than served by reading ``cu_seqlens_q``
back off the device.

Derive :class:`DSv4StepMeta` from explicit host-side facts.

The one derivation, for both the engine's step (through
:func:`build_step_meta`) and a drafter that assembles a forward of its
own. Every per-sequence fact is a GATHER of the device positions,
never a read back to the host.

Pooling groups a run can touch at ``ratio``.

A run of ``max_query_len`` tokens spans at most
``ceil(len / ratio) + 1`` groups — the ``+1`` for a run that starts
part-way through one.

Shared per-arch shape-config base for dense GQA decoders.

``Qwen3Config``, ``LlamaConfig`` (and any future plain-dense arch) lift
the SAME 11 shape fields out of HF ``config.json`` and project them into
``ModelDims`` + a uniform list of attention ``LayerSpec``. This base
holds that shared body so a new dense arch's config is ~5 lines:

    @dataclass(frozen=True)
    class FooConfig(DenseDecoderConfig):
        hf_model_type = "foo"          # expected config.json model_type

…and, if the arch carries a knob this generic path hasn't validated,
override :meth:`_validate_hf` to fail loud (e.g. Llama rejects
``rope_scaling`` until a scaling variant is wired).

Per-arch *behavioural* deltas (Qwen3's per-head q/k-norm, sliding
window) live in the DecoderLayer / LayerSpec, NOT here — this is purely
the shape lift. Subclasses that need an extra shape field (e.g. Qwen3's
``sliding_window``) add it as a dataclass field and override
:meth:`_extra_fields` / :meth:`to_layer_specs` minimally.

Shared DFlash hidden-state tap seam for the v2 dense arches.

A DFlash block-diffusion drafter reads the residual-stream output of a
few target decoder layers of the main model. The engine arms the tap via
:meth:`setup_dflash_tap` (called by ``build_dflash_drafter``, which probes
``hasattr(model, "_dflash_tap_ids")`` to accept a target); every forward
then records the tapped layers' residual-stream output into persistent
capture-safe slabs (or an eager dict for CPU unit tests), plus the
per-forward validity metadata.

The seam is byte-identical across the arches that expose it (qwen3,
qwen3_5) EXCEPT for how the residual-stream output is reconstructed at
the tap point: a block that already returns ``residual + h`` needs no
reconstruction (the default), while a cross-layer-fused block returns
only ``mlp_out`` and the true residual-stream output is
``residual_buf[:n] + out`` AT THIS POINT (the next layer's fused
input_layernorm mutates ``residual_buf`` in place). Arches override the
small :meth:`_tap_reconstruct` hook to supply that; everything else is
shared here.

Mix in ALONGSIDE :class:`LayerStackModelMixin` (this seam calls its
``_dispatch_layer_args``). Plain Python — no :class:`nn.Module`
inheritance — so it composes without a metaclass conflict.

CONTRACT — the tap is PER-FORWARD: it holds the tokens of the most
recent forward ONLY, written from row 0, and every forward (each prefill
CHUNK included) overwrites the previous one. Whole-sequence reads go
through :func:`arbi_serve.spec_decode.dflash_tap.collect_tap_full_sequence`,
which validates alignment and fails loud under chunked prefill.

Reconstruct the tapped layer's residual-stream output from a block return.

Default: the block already returns the full residual-stream output
(``residual + h``) so the tap records ``out`` directly. Arches whose
blocks return only the fused ``mlp_out`` override this to add back the
cross-layer ``residual_buf`` at THIS point (before the next layer's
fused norm mutates it in place).

Per-layer dispatcher that also records the residual-stream output
of the DFlash tap layers.

Capture/compile-safe: when persistent tap buffers are armed
(:meth:`setup_dflash_tap`) the residual-stream output is ``copy_``-ed
into a never-freed slab in ``graph_buffers_pool``. Under a captured
decode/verify graph the ``copy_`` is recorded and re-runs on every
replay (stable address → fresh data the driver reads post-step). A
python-dict store would only hold the capture-time tensor, whose
transient storage is freed/reused after capture → garbage drafts. The
dict path stays for the CPU unit tests (which arm ``_dflash_tap_ids``
without buffers). The first tapped layer also records the per-forward
validity metadata (token count + flat positions).

Record the per-forward tap validity metadata (once per forward, at
the first tapped layer): the token count and the flat positions of THIS
forward, via capture-safe device ops on the slab path. Non-flat (M-RoPE)
positions are marked invalid with ``-1`` — the whole-sequence collector
then refuses rather than mis-validating.

Arm the DFlash hidden tap with persistent per-layer buffers.

Allocates ONE ``(max_num_tokens, hidden_size * len(layer_ids))``
slab in ``graph_buffers_pool`` (the same stable-VA pool the
cross-layer ``residual_buf`` uses) so a captured target forward
records ``copy_`` writes into a fixed address, plus the per-forward
validity buffers (token count + flat positions). Must run before the
cudagraph capture sweep. Idempotent at a given shape.

LAYOUT INVARIANT: column block ``j`` of the slab is ``layer_ids[j]``
— the SAME order the drafter concatenates its tap layers in — so a
whole-tap read is a view of the slab rather than a gather per layer
plus a ``cat``. ``_dflash_tap_bufs[lid]`` is that layer's column
view: address-stable, and ``copy_`` into it is capture-safe.

The slabs hold the most recent forward's tokens ONLY (written from row
0, overwritten every forward — ``max_num_tokens`` is the per-STEP token
cap, not a per-sequence context). Reads that span forwards fail loud in
the collection helpers.

Gated Delta Net — production FLA forward paths (prefill + MTP verify).

Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock`. Holds the FLA
custom-op real-impls, the direct FLA kernel call, the chunk-kernel prefill
path, and the chained-T fused-recurrent verify path.

Optional GPU kernel globals + env-gate readers live in
:mod:`arbi_serve.models._gdn_kernels` and are read via the ``kn`` module
alias (``kn._HAVE_FLA`` / ``kn.chunk_gated_delta_rule`` …) so unit tests can
monkeypatch them on the single ``_gdn_kernels`` namespace.

Stage per-seq GDN seed states onto BOS rows for the fused kernel.

The fused-recurrent kernel addresses the initial state by BOS token
position (``h0 + bos * HV*V*K``, ``bos = i_n*T``), not by sequence
index. A ``(B, ...)`` h0 only works at T=1 (``bos == i_n``); at the
verify regime T=K+1>1 the kernel reads row ``i_n*T``, OOB for B>=2.
Return a ``(T_total, ...)`` buffer with each row's seed at its BOS
row so ``bos`` indexing stays in bounds.

The buffer is allocated UNINITIALISED. Read/write coverage is exact on
every caller, so a zero-fill would only initialise rows nothing reads:

  - The only reader of this buffer is the kernel's h0 load, and BOTH
    call paths take its per-BOS branch — ``p_h0 = h0 + bos * HV*V*K``
    with ``bos = cu_seqlens[i_n] = i_n*T`` (``fused_sigmoid_gating.py``
    h0 read). The snapshot path takes it via ``init_state_per_bos=True``;
    the no-snapshot fallback takes it because ``ssm_state_indices=None``
    leaves ``IS_CONTINUOUS_BATCHING`` False. So the read set is exactly
    ``{0, T, 2T, …, (B-1)*T}``.
  - :func:`torch.arange(0, T_total, T)` below is exactly that set, and
    ``index_copy_`` writes every one of those rows.
  - Nothing writes h0 back: the snapshot path sends the kernel's state
    stores to ``final_state_out`` (the snap buffer), and the fallback to
    a fresh ``ht``. This buffer is read-only to the kernel.

Read set == write set ⇒ **byte-identical** to the zero-filled buffer,
so the buffer is left uninitialised and no zero-fill runs.

True under a Dynamo trace OR inside the piecewise cuda-graph
trampoline — the two conditions that require routing the GDN kernel
through the opaque custom op instead of the eager bypass.

Both halves are required: keying on only one re-traces with the
eager bypass active when the runtime ContextVar value flips between
traces and explodes on the ``@torch.compiler.disable``d FLA helpers
(see :meth:`_GDNFLAMixin._forward_prefill_fla` for the full rationale).

Gather per-row GDN recurrent states from the ``(N, HV, V, K)`` slab
into FLA's ``(B, HV, K, V)`` initial-state layout.

Transposes K↔V on the SLAB VIEW first, so ``index_select`` gathers
straight into FLA's layout and its output is already contiguous —
one kernel, one allocation. Gathering first and transposing after
costs a second full-size buffer (the gather result, then the
``contiguous()`` of its transpose) for the same bytes.

``dtype`` is a caller-declared target, NOT an "activation dtype"
default — callers MUST pass the slab's OWN dtype
(``slab.dtype``) to get a true no-op cast. Passing the activation
dtype (e.g. ``v.dtype`` / ``hidden_states.dtype``) silently
downcasts a fp32 recurrent-state slab (``ARBI_GDN_RECURRENT_FP32``,
default on — see ``arbi_serve/cache/_state_view_registry.py``) to
bf16 on every read, throwing away the precision that flag exists to
provide. The decode fast path (``_forward_decode_fla``'s fused
branch, ``_gdn_fla_decode.py``) and the single-launch chunked
MTP-verify path (``_dispatch_verify_chunked``,
``round_h0_to_act_dtype=False``) already gather h0 at the slab's raw
dtype for exactly this reason — every other call site (plain
chunk-prefill, the v2/v3 custom-op real-impls, the eager per-t-step
decode/verify loops) must match that contract.

Scatter FLA's ``(B, HV, K, V)`` final state back into the
``(N, HV, V, K)`` slab at rows ``idx_long``.

Transposes K↔V as a VIEW, casts to ``slab_dtype``, and
``index_copy_`` in place. Returns the ``(B, HV, V, K)`` slab-layout
tensor (pre-cast) so callers can reuse it for a snapshot write.

No ``contiguous()``: ``index_copy_`` reads its source through a
TensorIterator and handles a transposed view directly, so
materializing one first only buys a temporary the size of the whole
batch's recurrent state — allocated per GDN layer, per step, on the
serving hot path. Writing the view straight through folds the
transpose into the same pass as the scatter.

Gated Delta Net — production decode forward paths (FLA + packed-decode).

Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock`. Holds the
per-request-one-token decode paths: the FLA fused-recurrent path
(:meth:`_forward_decode_fla`) and the vendored vLLM packed-decode path
(:meth:`_forward_decode_packed`).

Optional GPU kernel globals + env-gate readers live in
:mod:`arbi_serve.models._gdn_kernels` and are read via the ``kn`` module
alias so unit tests can monkeypatch the single ``_gdn_kernels`` namespace.

Resolve the ``(int64, int32)`` slab-row index aliases for the
decode paths.

The int64 form (for ``index_select`` / ``index_copy_`` on the
recurrent slab) comes from :func:`_resolve_state_idx_long`, which
prefers the persistent ``meta.state_indices_long`` capture buffer.
The int32 form (for ``causal_conv1d_update``'s
``conv_state_indices=`` arg) reuses the persistent
``meta.state_indices`` buffer in place when it is already int32 +
on-device (zero-alloc; same ``data_ptr`` stability under cudagraph);
otherwise it falls back to a per-call ``.to(int32)``.

Shared verbatim by :meth:`_GDNFLADecodeMixin._forward_decode_fla`
and :meth:`_GDNFLADecodeMixin._forward_decode_packed`.

Write the post-decode recurrent state into the MTP rollback
snapshot slot 0.

Re-gathers the just-written rows out of the recurrent slab (the
kernel updated it in place) and ``index_copy_``s them into
``snap_rec_full[0]``. Shared by the legacy-FLA and packed decode
paths; both gather from the slab AFTER the in-place kernel write so
the snapshot captures the committed state.

Fused decode conv update — shared by the FLA and packed-decode
paths.

``causal_conv1d_update`` does gather → roll → depthwise conv →
SiLU → scatter in a single kernel; ``conv_state_indices`` selects
the slab rows. ``has_initial_state[i] = False`` rows are
zero-initialized at admission by ``RecurrentStatePool`` (see
``alloc_recurrent_state`` + ``flush_pending_zero_clears``), so no
explicit ``has_init`` mask is needed — a zero buffer rolled and
convolved with ``x`` matches the previous ``prior * 0`` path.

The pure-PyTorch rolling-buffer fallback covers environments
without the ``causal-conv1d`` wheel (CPU smoke / some fixtures);
it is numerically identical to the kernel path. When a per-step
MTP snapshot is attached (``state_view.snap_conv_state``), the
t=0 buffer is written into ``snap_conv_state[0]``.

Returns ``conv_out`` ``(B, conv_dim_local)``; mutates the conv
slab in place.

Snapshot-less fused conv update — the exact conv step
:meth:`_conv_update_decode` runs, minus the snapshot write.

Byte-identical to a single decode conv step (same
``causal_conv1d_update`` Triton kernel on sm_89, same pure-PyTorch
rolling-buffer fallback on sm_120 / CPU), mutating the ``conv_state``
slab in place. Factored out so the ARBI_ACCEPT_INVARIANT chained
verify path (:meth:`_GDNFLAMixin._dispatch_verify_invariant_chained`)
can advance the conv slab per token with the SAME kernel decode uses
— the fix for the sm_89 decode≠verify divergence (the default verify
conv is a hand-rolled PyTorch unfold that differs from the decode
Triton kernel by a few ulp on sm_89, blowing up at post_mixer). Takes
the raw conv slab (not the layer view) so the custom-op capture route
can call it. The caller owns the per-token snapshot write.

Returns ``conv_out`` ``(B, conv_dim_local)``.

Single-token fused-recurrent decode step (gather → kernel →
scatter → snapshot) on raw slab tensors.

Byte-identical to the inline eager fused decode below — factored out
so the ARBI_ACCEPT_INVARIANT capture route can run the SAME
``fused_sigmoid_gating_delta_rule_update`` kernel through the opaque
``arbi_serve::gdn_decode_fused`` custom op (Inductor must not trace the
kernel directly: its ``b_h`` fp32 accumulator trips an fp32↔fp64
error). Keeping the decode capture path on the fused kernel — rather
than falling back to the legacy FLA op — is what preserves decode ==
verify bit-parity under the flag. Returns ``core_attn_out``
``(B, HV, V)`` and mutates ``recurrent_state`` + ``snapshot0``.

Per-request-one-token decode via :func:`fla.ops.gated_delta_rule.fused_recurrent_gated_delta_rule`.

Stream-capture-safe: every per-row scalar branch / ``.item()``
/ per-row Python loop has been replaced with vectorized
``index_select`` / ``index_copy_`` over ``meta.state_indices``
(the persistent capture buffer the engine pre-binds via
:func:`arbi_serve.runtime.capture.decode.capture_decode`). The
captured graph's launch shape is fully determined by ``B``
(the static capture bucket); slab rows are indirected through
the persistent int32 buffer's ``data_ptr``.

Conv-update fusion: the conv update is a single
``causal_conv1d_update`` call
with ``conv_state_indices=`` so the gather + roll + silu +
scatter happens inside one Triton kernel — no
``index_select`` / ``torch.roll`` / explicit ``index_copy_``
on the conv slab. Net launch reduction in this hot path: 4
kernels → 1.

Per-request-one-token decode via vendored vLLM packed-decode kernel.

Sister to :meth:`_forward_decode_fla` — same conv path, same
slab semantics, same ``meta.state_indices`` contract. The only
difference is the recurrent step:

  * Replaces the FLA call (``q_post``, ``k_post``, ``v_post``,
    ``g``, ``beta`` + ``initial_state=h0``) with the vendored
    :func:`fused_recurrent_gated_delta_rule_packed_decode`,
    which fuses gate + beta computation INTO the kernel and
    reads/writes the recurrent slab in-place via
    ``ssm_state_indices``.
  * Drops the surrounding ``index_select`` (slab→h0) and
    ``index_copy_`` (final_state→slab) — the packed kernel does
    both itself.

Projection fusion. We share :meth:`_project_streams_raw_ba`'s
two fused matmuls (``[Q|K|V|Z]`` and ``[B|A]``) with the FLA
decode path — same ``F.linear`` weights cached on the module,
same persistent ``hidden_states`` source. The fused projection
avoids the projection-stage gap that would otherwise make
packed-decode end-to-end slower than the FLA path despite its
kernel-level win. The ``[B|A]`` chunk returns ``(b, a)``;
we unpack into the kernel arg order ``(a, b)`` (no transpose,
no rename).

Conv path. Routes through ``causal_conv1d_update`` (same single-
kernel gather→roll→silu→scatter as :meth:`_forward_decode_fla`)
when the Dao-AILab kernel is available; the pure-PyTorch
rolling buffer remains as the CPU-smoke / no-kernel fallback.

``has_init==False`` rows. Both the conv slab and the recurrent
slab are pre-zeroed at request admission by
:meth:`MultiStatePool.alloc_recurrent_state` (queued into
:meth:`RecurrentStatePool.flush_pending_zero_clears`, drained
at start-of-step before any kernel reads the slot).
Therefore an explicit ``has_init`` mask multiply on the slab
is redundant: a fresh row has zero state, and the kernel
reading zeros gives the right answer for the first decode
step. The FLA fused path already relies on this invariant
(see :meth:`_forward_decode_fla` — bare ``index_select`` with no
mask). ``decode_pad_cudagraph`` defaults to true, so the shipped
config pads decode up to a captured ``B' > B`` and the padding
rows carry ``has_init==False`` borrowed-slab-row state. Safety
comes not from padding never occurring but from each padding
row being isolated: it borrows a free slab row (never an
in-flight row's row), falls outside every seq's ``cu_seqlens``
span, has its recurrent work bounded to the live
``[:real_num_tokens]`` slice, and has its output/scatter
discarded before the sampler. The
``decode_pad_borrow_alias_consumed`` counter fires iff a
borrowed padding row aliases a consumed ``[:real_num_tokens]``
row's slab, and must read 0 on any serving run.

Stream-capture-safe: every operation is ``index_copy_`` /
``index_select`` / kernel-launch. No ``.item()`` / no Python
branches on a tensor value.

Gated Delta Net — FLA prefill + MTP-verify forward paths.

Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock` via
:class:`~arbi_serve.models._gdn_fla._GDNFLAMixin`.

Optional GPU kernel globals live in :mod:`arbi_serve.models._gdn_kernels`
and are read via the ``kn`` module alias so unit tests can monkeypatch
them on the single ``_gdn_kernels`` namespace.

Return ``(cu_seqlens, cu_seqlens_cpu)`` for the GDN prefill kernel.

``cu_seqlens_cpu`` is the host twin FLA takes so the chunk-index table
can be resolved on the host. Its device branch cannot: FLA's
``_segmented_arange`` runs ``torch.repeat_interleave`` over device
``counts``, whose output length is ``counts.sum()``, so the host blocks
on a device read before the allocation can be sized. FLA's own docstring
names the remedy — "pass host-side counts to avoid it" — so this is the
upstream API used as designed, not a workaround.

The twin is returned only when it provably describes the very tensor
being returned beside it: it is keyed to ``meta.cu_seqlens_q``, so any
path that synthesizes or relocates ``cu`` drops it.

That is deliberately stricter than "when could they still agree". A
mismatched twin does not fail loudly — it yields a correctly-shaped
chunk-index table for the WRONG sequence layout, which is silent output
corruption. Dropping it costs one host sync.

Do NOT cast the dtype of ``cu``: :class:`PiecewiseBuffers` keeps
``cu_seqlens_q`` as int32 and the engine pre-binds captured kernels
against the persistent slice's ``data_ptr``. Casting to int64 here would
re-allocate a fresh tensor every call, which (a) busts FLA's
identity-keyed ``@tensor_cache`` so ``prepare_chunk_indices`` re-runs and
issues its host-blocking device read inside the ``torch.cuda.graph(...)``
region — fatal under stream capture — and (b) breaks the captured-graph
``data_ptr`` contract on replay. Triton's varlen kernels load via
``tl.load(cu_seqlens + i)`` on whatever dtype is provided, so int32 is
fine end-to-end.

NO COUNTER HERE, and that is load-bearing. This function runs inside the
``fullgraph=True`` compiled layer forward, and a flag-truth bump is
``self.refused += 1`` — a READ of a module-global int. Dynamo guards on
the value it read, the guard fails on the very next call, and the layer
recompiles EVERY step until ``accumulated_recompile_limit`` (256), at
which point ``fullgraph=True`` turns the limit into a hard failure.
Measured on the served 27B: ``--prefill-capture=full`` captured ZERO
rungs with ``Hard failure due to fullgraph=True``, last guard
``_CU_SEQLENS_HOST.refused == 490``, and the engine never reached ready.
The census lives in :meth:`~arbi_serve.backends.gdn_backend.
GdnMetadataBuilder._finalize` instead: host-side, deciding the same facts
from the same batch, ONCE per step rather than once per GDN layer.

The verify class's rolling-window conv over ``T`` tokens per row.

One unfold builds every K-wide window from the augmented stream
``[buf[..., 1:K], x_0..x_{T-1}]``; ``snaps[t]`` is the buffer after
token t. The product + K-reduction is one batched op (sum over K
is byte-identical batched vs per-slice). SiLU stays per-t: batched
SiLU rounds differently in fp32, so the per-t launch holds the
output bit-identical. Returns ``(conv_out (B, T, C), final buffer
(B, C, K))``.

Shared by :meth:`_forward_verify_fla` and the row-class-mixed
forward, so the verify class's conv is ONE piece of arithmetic
wherever its rows are batched.

``[verify rows | prefill rows]`` through ONE block forward.

The projections, the gated norm and the out-projection run once over
every row — that is where a mixed step's second weight read lives.
Below them the row class selects an ARITHMETIC, not a forward:

  * the verify rows take the masked-replay recurrence launch — raw
    per-token inputs saved in-kernel, NO state commit, the slab left
    holding the pre-verify h0 the accept path replays from — which
    is the served branch of :meth:`_forward_verify_fla`;
  * the prefill rows take the chunk-kernel chain with its commit,
    which is :meth:`_forward_prefill_fla`;

each over its own slice of the shared streams, each writing its own
slice of one core buffer (the chunk kernel through ``o_out``, so no
concatenation). The side effects are row tables, not branches: the
replay saves address the verify rows' slab rows, the commits address
the prefill rows'. Per class the kernels, the per-row inputs and the
h0 contract are exactly those of the two class-pure forwards, which
is what makes the result bit-identical to them below the projections
(``tools/gdn_row_class/bench_gdn_row_class.py`` asserts it).

Layout contract: rows ``[0, n_v_rows)`` are the verify rows, flat
tokens ``[0, n_v_flat)`` are theirs, uniform ``T = n_v_flat /
n_v_rows`` per row; the remaining rows are prefill rows with
``meta.cu_seqlens_q`` boundaries over the whole flat batch.

Refused, by name, for the shapes whose class-pure forward takes a
different arithmetic than the two composed here: a tree-shaped verify
block, the legacy ``ARBI_GDN_FUSED_RECURRENT=0`` per-token chain,
``ARBI_ACCEPT_INVARIANT``'s chained route, and a pool without the
masked-replay buffers (recompute rollback).

Multi-token prefill via :func:`fla.ops.gated_delta_rule.chunk_gated_delta_rule`.

Pack varlen as ``(1, T_total, ...)`` with ``cu_seqlens``.
FLA returns ``(o, final_state)`` where final_state has
shape ``(N, HV, K, V)`` for ``N == B`` sequences.

Stream-capture-safe: per-row ``int(state_indices[i].item())``
and ``bool(has_init.all())`` host syncs are replaced with
vectorized conv (rolling buffer over the (B, T, C) reshape)
and unconditional ``has_init`` mask multiply. The conv slab
scatter routes through ``index_copy_`` on the persistent
``meta.state_indices`` buffer.

MTP verify pass via chained T=1 ``fused_recurrent_gated_delta_rule`` calls.

T = K+1 tokens per row. Rather than issuing one packed varlen
kernel call (which propagates state with kernel-internal fp32
accumulators across all T tokens before any write-back), we
loop ``t = 0 .. T-1`` and at each step:

  * read each row's recurrent state from
    ``state_view.recurrent_state[slab_row]`` (forces a bf16
    reload — same cast K=1 decode does);
  * batch ALL rows into a single ``(B, 1, H, D)``
    fused-recurrent call (uniform-K invariant: every row
    advances exactly T tokens, so every row is active at
    every step);
  * cast ``final_state`` back to bf16 and write to the view.

This makes the verify path bit-identical to running the K=1
decode kernel T times in sequence — the bonus-slot output and
the persisted recurrent state both match
:meth:`_forward_decode_fla` byte-for-byte, which is the
invariant the bundled Qwen3.5 MTP drafter was trained against.

Stream-capture-safe: ``T = T_total // B`` is shape-derived
(no host sync); per-step token offsets, slab gathers, and
snapshot writes are all vectorized over ``meta.state_indices``
(the persistent capture buffer). Under the uniform-K
contract, no row is "shorter" than another so we don't need
a Python ``[i for i if T_per_row[i] > t]`` mask.

The chunk kernel stays for prefill where T is large; on the
verify-pass T regime (T <= K_max + 1, typically <= 8) the
per-token launch overhead is dominated by the per-layer GDN
compute.

Gated Delta Net — FLA custom-op real-impls + direct kernel call.

Holds the ``arbi_serve::gdn_attention_v2`` / ``_v3`` custom-op
real-impls, the direct FLA delta-rule kernel call, and the single-launch
chunked MTP-verify dispatch. Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock` via
:class:`~arbi_serve.models._gdn_fla._GDNFLAMixin`.

Optional GPU kernel globals + env-gate readers live in
:mod:`arbi_serve.models._gdn_kernels` and are read via the ``kn`` module
alias so unit tests can monkeypatch them on the single ``_gdn_kernels``
namespace.

Real-impl side of ``arbi_serve::gdn_attention_v2``.

Does the slab gather + FLA-kernel + slab scatter
ENTIRELY INSIDE the opaque op so Inductor never traces the
``index_select`` / ``index_copy_`` pair around the kernel call;
the schema's ``mutates_args=("recurrent_state_slab",)`` declaration
is the single source of truth for the in-place mutation.

Slab layout. ``recurrent_state_slab`` is ``(N, HV, V, K)`` (V
outer, K inner — see ``arbi_serve/cache/recurrent_pool.py``).
FLA's kernels consume initial / final state in
``(B, HV, K, V)`` (K outer). We gather the slab rows into
``(B, HV, V, K)``, transpose K↔V, run the kernel, transpose
the returned final_state back to ``(B, HV, V, K)``, and
``index_copy_`` it into the slab. Math is byte-identical to
the v1 path which did the same gather/transpose/scatter
OUTSIDE the op.

``state_indices`` is the int64 slab-row indices. ``cu_seqlens``
is consumed only by ``chunk_gated_delta_rule`` (kernel_mode==0);
for the fused-recurrent path (kernel_mode==1) it's
passed-but-ignored.

``kernel_mode`` (0 = chunk_prefill, 1 = fused_recurrent) is
named to avoid a kwarg-name collision with torch's
``auto_functionalized_v2_fake(mode, _mutable_op, **kwargs)``
signature — see ``_register_gdn_attention_v2``.

The prefill fold: gather, ONE kernel launch, scatter — emitting a
savepoint boundary state when this launch is the step's armed one.

ONE implementation for the compiled path (the real-impl of
``arbi_serve::gdn_attention_v2``, above) and the eager path
(:meth:`_forward_prefill_fla`, :meth:`_forward_mixed_rows_fla`), so a
savepoint is staged by the same code whichever path a step takes —
the seam whose two copies drifted was the whole of #2238.

WHAT THE PLAN IS AND WHERE IT COMES FROM
========================================
Nothing here is read from the metadata or from ``self``: the step's
plan is published host-side by the GDN metadata builder before the
forward runs (:mod:`arbi_serve.cache._fold_emit_staging`) and looked up
by the identity of ``cu`` — the boundary tensor this launch was
handed. No plan, or another tensor's plan: the plain launch, the same
bytes as before this method existed.

WHY THE SAME BITS
=================
The kernel carries each sequence's state in fp32 across its chunks
and stores it in fp32 once at the end; the emit is that same store
at an earlier chunk, into the slab's layout. Nothing about the
fold's blocking, order or rounding changes, so the output and the
final states are the plain call's, and the emitted state is what a
fold cut at that chunk would have ended with. ``offset`` is a
multiple of the fold grid by construction — ``savepoint_fold_split``
refuses rather than rounds — so the chunk index it names is exact.
Measured in ``tools/gdn_row_class/check_emit_state.py``.

BEFORE IT EMITS, THE LAUNCH PROVES IT IS THE PLAN'S
===================================================
The slab row it will scatter into must be the row the snapshot's
host buffers were laid out for (:attr:`BoundarySnapshot.sources`,
compared by device address, no sync), and the launch's token count
must be the plan's. Either mismatch refuses by name on
``gdn_fold_split_staged`` and stages nothing; the commit then refuses
the snapshot whole (``fold_split_incomplete``). A partial snapshot is
never stored. The conv window in the slot is this layer's by stream
order: ``gdn_prefill_conv_flat`` gathered it for this layer just
before this launch, on the same stream.

Under stream capture the copies go to the fixed landing zone instead
of a ring slot, on every replay; the post-replay seam moves them into
the step's ring slot when a plan was armed
(:meth:`~arbi_serve.cache._fold_emit_staging.FoldEmitStaging.finish_replayed`).

Real-impl side of ``arbi_serve::gdn_attention_v3``.

Mirrors :meth:`_dispatch_through_custom_op_v2` but additionally
writes the per-step MTP snapshot in the same dispatch.  Saves
one ``slab.index_select(0, state_indices)`` per t-step in the
MTP verify path.

``snapshot_out`` is ``snap_rec_full[t]`` — the per-step slice
of the snapshot buffer, shape ``(N, HV, V, K)`` (same layout
as the slab).  We ``index_copy_`` ``final_state_vk`` (cast to
the snapshot dtype) into rows ``state_indices`` of
``snapshot_out``.  The slab scatter happens in the same
sequence (both mutations are declared on the schema, so
Inductor respects in-place semantics on both tensors).

Single-launch chunked MTP-verify recurrence with IN-KERNEL
per-token snapshot scatter, shared by the eager path and the
chunked verify kernel.

One ``fused_sigmoid_gating_delta_rule_update`` launch walks the
recurrence varlen-style over all T*B tokens
(``cu_seqlens=[0, T, 2T, ..., B*T]``). We stage ``h0`` per-BOS
from the COMPACT main recurrent slab (1 slot/row — read via the
kernel's per-BOS path, ``init_state_per_bos=True``) so the main
slab stays compact and ``num_spec_plus_1`` stays 1 on every GDN
layer view (the packed-decode fast path is never disqualified).

The IN-KERNEL improvement vs main's Python ``index_copy_``: the
kernel scatters each token's state DIRECTLY into the SEPARATE
per-token snapshot buffer (``snapshot_out``, ``(T, N, HV, V, K)``)
via ``ssm_state_indices`` — slot ``(t, slab_row)`` is the
flattened linear index ``t * N + slab_row`` over the
``(T*N, HV, V, K)`` view. This eliminates the Python per-t
``index_copy_`` (no separate scatter launch, no
rollback_batch/broadcast_rollback path needed for the snapshot
write). The main slab's committed state is the state after the
LAST token, which the kernel already wrote to snap slot
``(T-1, slab_row)``; we commit it to the compact slab with one
``index_copy_`` (the slab is NOT addressed by the kernel).

No K↔V transpose — the kernel consumes / emits the slab's native
``(N, HV, V, K)`` layout.

Returns ``core_flat`` ``(T_total, HV, V)``; mutates
``recurrent_state_slab`` and (when present) ``snapshot_out`` in
place at rows ``state_indices``.

MASKED-REPLAY verify forward — single chunked launch, no ladder.

Identical recurrence math + h0 contract to the ladder path's
in-kernel gather (`_dispatch_verify_chunked` with ``snapshot_out``):
raw fp32 slab h0, ``init_state_mask`` honored in-kernel, same
varlen walk — so the per-token outputs are byte-identical to
ladder mode. Differences are stores only:

  * per-token state goes NOWHERE (the T·N·HV·V·K fp32 ladder
    traffic is deleted);
  * the raw per-token (k, v, b, a) inputs are saved at
    ``row * T_buf + t`` of the replay buffers (in-kernel);
  * NO state commits at all (``commit_final_state=False``), so the
    slab still holds the h0 this launch read.

Rollback recomputes the accepted-prefix state bit-for-bit from that
slab h0 + the saved inputs
(``RecurrentStatePool._dispatch_masked_replay``).

Shared by the eager path and the ``arbi_serve::gdn_verify_replay_save``
custom op (the captured-graph / Dynamo route — the fused kernel must
stay inside an opaque boundary under ``fullgraph=True``).

Direct FLA delta-rule kernel call — eager-path entry.

``o_out`` (chunk mode, vendored entry only) — a caller-owned output
buffer the kernel writes instead of allocating; on the other legs
the result is copied into it so the contract holds either way.
``emit_chunk`` / ``emit_out`` (chunk mode, vendored entry only) — the
in-kernel fold-state emit; refused by name on the other legs, which
cannot produce it.

Called from the eager hot path WITHOUT going through a custom
op. The wrapper costs a custom_op schema validate + fake-impl
shape check + dispatcher dispatch + side-channel publish on
``self._call_state_view`` / ``self._call_meta`` per call, which
adds up across the GDN layers run per step.

The Inductor compile path routes through
``arbi_serve::gdn_attention_v2`` instead — the op definition is
the stable boundary Dynamo / Inductor needs. Only the
live-runtime call site bypasses it.

Gated Delta Net — ARBI_ACCEPT_INVARIANT chained verify + recompute rollback.

Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock` via
:class:`~arbi_serve.models._gdn_fla._GDNFLAMixin`.

Optional GPU kernel globals live in :mod:`arbi_serve.models._gdn_kernels`
and are read via the ``kn`` module alias so unit tests can monkeypatch
them on the single ``_gdn_kernels`` namespace.

ARBI_ACCEPT_INVARIANT verify recurrence: T chained decode calls.

Bit-parity contract with :meth:`_forward_decode_fla`'s fused path —
per token, the EXACT decode invocation:

  * same kernel (``fused_sigmoid_gating_delta_rule_update``),
  * same launch shape (``cu_seqlens = [0, 1, .., B]`` — one token
    per row, from the same persistent buffer decode uses),
  * same h0 handling (bf16 slab ``index_select``, no mask — the
    slab pre-zero invariant covers fresh rows exactly as in
    decode),
  * same bf16 final-state write-back between tokens (the chunked
    kernel's in-kernel fp32 state propagation is precisely the
    shape-dependent numeric the invariant mode removes).

The per-t snapshot writes preserve the MTP rollback contract
(``snap[t][slab_row]`` = state after committing tokens [0..t]).
Costs T-1 extra kernel launches per GDN layer per verify step
relative to the chunked route — accept-invariance tax, paid only
under the flag.

Partial-accept rollback by recompute-from-base — path dispatch.

Production (CUDA + FLA verify) replays via the ``fused_sigmoid_-
gating_delta_rule_update`` kernel so the replayed accepted state
tracks the chunked verify kernel it must match (accept-parity, not
bit-parity — the per-token vs chunked fp32 accumulation differs).
CPU / no-FLA falls back to the pure-PyTorch reference replay.
Driven by ``RecurrentStatePool.rollback_batch`` (GDN runs EAGER, so
this data-dependent replay has no cudagraph-capture constraint).

FLA recompute rollback: replay conv + ``fused_sigmoid_gating_-
delta_rule_update`` from the T=1 base over the accepted prefix.

Batched over the partial-accept rows: one kernel launch per token
step ``t = 0 .. max(n_accepted)``; each row's slab is committed at
its own ``t == n_accepted`` step. Fully-accepted rows (``n ==
k_uniform``) keep the forward's post-token-K state (skipped). The
conv is re-run from the base buffer (byte-identical to the verify's
PyTorch unfold conv); the recurrence uses the SAME per-token kernel
+ raw ``b``/``a`` inputs the verify forward projected.

This is also the rollback the ARBI_ACCEPT_INVARIANT verify path pairs
with (``gdn_mtp_rollback_mode`` forces ``recompute`` under the flag —
it is the only per-token-chained rollback, so it tracks that path's
chained recurrence far better than the in-graph masked replay's one
chunked launch would). Fidelity is accept-parity, not bit-parity —
the repo's standing recompute contract: the replay carries the state
in fp32 across tokens (no per-token bf16 slab round-trip) and re-runs
the PyTorch unfold conv rather than the invariant forward's ``causal_-
conv1d_update``, so the committed accepted-prefix state is a few ulp
off the forward's intermediate. The accepted-token OUTPUTS were
already emitted bit-exactly by the (chained) forward; only the state
seeding the NEXT tick carries this drift, at the same tolerance the
default (chunked) verify's recompute rollback already accepts.

Commit each row's accepted-prefix state by MASKED replay.

``nacc_by_row`` is the per-SLAB-ROW accepted offset (−1 = leave the
row untouched); ``route_rows`` maps inactive rows to the
zero-sentinel slot (kernel ``state_idx <= 0`` skip — zero writes);
``cu_seqlens`` = ``[0, T, 2T, …, N·T]``; ``step_mask`` is the
``(N, T)`` bool ``t <= nacc`` prefix mask (all shared across
layers, built once per tick by the pool). Per layer this issues:

  * two input-masking writes over the saved buffers
    (``k → 0`` kills the rank-1 update and the v-correction;
    ``a → -inf`` makes ``softplus → 0`` so the decay is
    ``exp(-0.0) = 1.0`` exactly — fp-exact identities for every
    masked step);
  * or, with ``n_steps``, neither: the kernel's trip count stops at
    the accepted prefix instead. It reaches the same state without
    executing the identity steps, so the writes that made them
    identities have nothing left to do — and a row with nothing to
    commit returns before the state gather rather than after it;
  * one ``fused_sigmoid_gating_delta_rule_update`` launch with the
    same wrapper flags (so the same compiled binary) the
    replay-mode verify forward used, recursing from the SLAB — which
    that forward left holding the pre-verify state — over the saved
    raw bytes and committing only the final state back to the row's
    slab slot. ``initial_state`` and ``final_state_out`` alias, which
    is safe because each program reads and writes only its own
    ``(row, head, v-block)`` tile;
  * a pure-gather conv window advance: the conv state after
    ``nacc + 1`` tokens is window ``nacc`` of
    ``[conv_slab[..., 1:] ‖ x_0..x_{T-1}]``. The gather materialises
    that concatenation before reading, so the slab being both source
    and destination is safe.

Fixed shapes, fixed trip count, no host reads — stream-capture-
safe; the pool captures the per-layer launch set into a side
cudagraph so a tick's rollback is one graph launch. The in-place
input masking destroys only slots that are dead after this tick
(the next verify forward fully re-saves its slate's rows).

CPU / no-FLA falls back to the pure-PyTorch reference replay
(:meth:`_replay_masked_commit_reference`).

Pure-PyTorch masked-replay recurrence (CPU / no-FLA pools).

Math-equal (not byte-gated — the byte gate runs the CUDA kernel
pair) replay of the accepted prefix from the slab's pre-verify state
over the saved post-conv inputs, using the same reference helpers as
:meth:`_rollback_recompute_reference`. Conv is handled by the
shared gather in :meth:`replay_masked_commit`.

ARBI_ACCEPT_INVARIANT verify: the per-token conv+recurrence loop.

T chained EXACT decode calls. Bit-parity contract with T sequential
:meth:`_forward_decode_fla` fused calls — per token, the SAME conv
kernel AND the SAME recurrent kernel decode runs, advancing BOTH the
conv slab and the recurrent slab in place between tokens:

  * conv via :meth:`_conv_update_step` → ``causal_conv1d_update``
    (the Dao-AILab Triton kernel decode uses on sm_89), not the
    hand-rolled PyTorch unfold conv the default verify path runs;
  * recurrence via ``fused_sigmoid_gating_delta_rule_update`` with
    a bf16 slab ``index_select`` h0 reload and a bf16 final-state
    write-back per token — identical to decode.

The default verify recurrence (:meth:`_forward_verify_fused_chained`)
already chains the recurrent kernel, but the conv is computed once
upfront by a hand-rolled PyTorch unfold (``_forward_verify_fla``
rolling-buffer block). On sm_120 the Dao-AILab kernel has no cubin
so decode also falls to a pure-PyTorch conv, so both match and the
invariant holds. On sm_89 decode runs the Triton
``causal_conv1d_update`` kernel while verify ran the PyTorch
unfold — the two conv outputs differ by a few ulp, and that
difference propagates through the delta-rule recurrence into a
large divergence at post_mixer. Forcing the decode conv kernel on
the verify side too makes the whole mixer bit-identical to T
sequential decode steps on sm_89.

Operates on raw slab tensors (not the layer view) so it can run both
on the eager route and inside the opaque ``arbi_serve::gdn_verify_-
chained`` custom op on the capture route (Inductor must not trace
``fused_sigmoid_gating_delta_rule_update`` directly — its ``b_h`` fp32
accumulator trips an fp32↔fp64 InductorError, the same reason decode
gates its fused path off under capture). Mutates ``conv_state`` and
``recurrent_state``, committing the final post-token-(T-1) state.
Returns ``core_flat`` ``(T_total, HV, V)`` (the pre-``_materialize_-
out`` mixer output).

Rollback carries no per-position ladder: the partial-accept
rollback recomputes the accepted prefix from the 1× base frame +
the per-token raw inputs the caller captured, host-side via
:meth:`rollback_recompute`. The recompute pairs with this chained
forward because ``gdn_mtp_rollback_mode`` forces ``recompute`` under
the invariant flag — see :meth:`_GDNFLAForwardMixin._forward_verify_-
fla`.

GDN verify forward for a TREE-shaped block, on the shipped kernels.

The chain verify hands its block to one varlen scan per layer: ``B`` rows
of ``K+1`` consecutive tokens, state carried in registers. That IS a
sequence, and a tree's block is not — node ``j`` sits after its siblings
in block order, so a chain scan folds their tokens into ``j``'s state and
the target's logits at ``j`` stop being its prediction from ``j``'s own
prefix.

Nothing here is a new kernel. The tree enters through index tensors:

* the block, tree-shaped, is ``TreeSpec((1, *widths))`` — the committed
  token is node 0 and the drafted tree hangs beneath it, so block row and
  node id are the same integer and every table in
  :mod:`arbi_serve.models._gdn_tree_scan` addresses the block directly;
* the conv is one gather along each node's own ancestor chain
  (:func:`~arbi_serve.models._gdn_tree_scan.tree_conv_forward`);
* the recurrence is the chain's own
  ``fused_sigmoid_gating_delta_rule_update`` call over the PATH-PACKED
  arrangement — every root-to-leaf path is its own varlen row, so each row
  is an ancestor chain and the kernel needs no tree awareness at all.

Two things the chain gets for free and a tree must pay for by hand, both
because the packed launch has ``L`` rows where the slab has one:

**The forward commits no state at all.** A committing scan writes each
packed row's last state through ``ssm_state_indices``, which for a tree
names the SAME slab row ``L`` times, so one row's store could land before
another row's ``h0`` load and seed that path from a sibling's state. The
scan therefore runs with ``commit_final_state=False``: no store is
emitted, the slab is left holding the pre-verify state, and there is no
race to arbitrate and no spare slab to hold the losers. That is sound
because a tree ALWAYS reconciles at accept —
:meth:`_GDNFLAVerifyTreeMixin.replay_masked_commit_tree` recomputes the
accepted path's state from that slab state on every step, where a chain
may skip the replay on a full accept.
:func:`~arbi_serve.spec_decode.tree_boot.assert_tree_supported` refuses at
boot unless the replay-mode buffers are attached, and
``_rollback_recurrent_offloaded`` drops the full-accept skip for a tree by
name.

**The per-token replay inputs are staged by hand.** The kernel's
``save_inputs_kvba`` also addresses by ``ssm_state_indices``, so on the
packed launch all ``L`` rows would race for one slab row's replay slots
with DIFFERENT tokens. The forward instead stages the NODE-ORDER streams
— one row per block row, no duplication — and the accept-time replay
gathers the accepted path out of them. ``save_base_state_out`` stays
in-kernel: every packed row writes the same effective ``h0``, so those
writes are self-consistent.

One layer's verify forward over a tree-shaped block.

``hidden_states`` is ``(B * N_block, hidden)`` in BLOCK order — row 0
the committed token, rows ``1..N`` the drafted nodes breadth-first
— which is node order for ``tables.spec``. Returns the mixer
output in that same order, so the caller's block layout, the
attention layers' block layout and the accept walk's node ids all
stay one indexing.

A batch is ``B`` blocks laid out request-major, exactly as the
chain's is. The packed launch then has ``B * L`` rows and its 1-D
``ssm_state_indices`` names each row's OWN request's slab row, so
every path is seeded from the state it belongs to.

The packed scan's single kernel launch, on the chain's own call.

Shared by the eager path and the ``arbi_serve::gdn_tree_packed_scan``
custom op. Every argument is a tensor and the geometry is carried by
``cu_seqlens`` alone, which is what lets the op be opaque.

``commit_final_state=False``, so the scan stores NO final state
anywhere. A tree's launch has one packed row per root-to-leaf path
and all of them name the same slab index, and the kernel's grid runs
them as independent programs — so a committing scan would let one
row's final store land before another row's ``h0`` load and seed that
path from a sibling's state. Suppressing the store removes the race
at its source and needs no spare slab to absorb it. The state is not
lost: the accepted path's is recomputed at accept time from the slab
h0 and the saved per-token inputs, on every tree step.

``initial_state`` is the slab and nothing writes it, so it still
holds the pre-verify state the accept-time replay recurses from.

Commit the accepted PATH's state, by gather + the shipped replay.

Same launch set as :meth:`GDNBlock.replay_masked_commit` and the
same masking identities; the one difference is where the replay's
tokens come from. A chain's accepted tokens are block rows
``0..nacc``, already contiguous in the staged buffers, so it
replays them in place. A tree's are scattered through node order,
so ``path_index`` — ``(N_rows, D)``, the accepted path's node ids
padded past ``nacc`` — gathers them into an ancestor chain first.
The gathered slate is then a PREFIX, which is exactly what the
masked replay already commits.

The conv side needs no gather at all: the window after the
accepted path is ``Kw`` columns along the accepted LEAF's own
ancestor chain, which is one row of the same static table the
forward's conv read (:func:`advance_conv_window_tree`).

Centralized optional-kernel globals + env-gate helpers for the Gated
Delta Net block.

This is the single home for every GPU-only kernel import the GDN block
(:class:`~arbi_serve.models.gdn_block.GDNBlock` and its mixins) routes
through, plus the per-call env-flag readers. The block — split across
``gdn_block.py`` + the ``_gdn_reference`` / ``_gdn_fla`` /
``_gdn_fla_decode`` mixins — reads these as module attributes of this
module (``import ... as kn`` then ``kn._HAVE_FLA`` / ``kn.causal_conv1d_update``),
never as ``from``-imported names. That keeps a single namespace that
unit tests monkeypatch (``_gdn_kernels._HAVE_FLA = True``,
``_gdn_kernels.causal_conv1d_update = stub`` …) and have every method —
in whichever mixin module it is defined — observe the patch.

The block constructs fine without any of these wheels; the GPU forward
paths raise / fall back to the pure-PyTorch reference when a kernel is
absent.

Whether the fused Dao-AILab ``causal_conv1d_update`` kernel can run
on ``device`` for the GDN decode conv update.

``assume_constant_result`` bakes the answer into the compiled graph at
trace time — same pattern as the three flag accessors below.
Without it, Dynamo traces into ``runtime_flags()``'s dataclass
construction and every hybrid decode capture fails with
``Unsupported: dataclass fields failure`` (the engine silently serves
eager decode). The kernel/arch answer cannot change post-boot for a
fixed device; test monkeypatches that flip ``_HAVE_CAUSAL_CONV1D`` or
the env flags invalidate Dynamo's specialization guards and recompile,
exactly as documented on ``_gdn_fused_recurrent_enabled``.

The GDN decode gate calls this instead of reading ``_HAVE_CAUSAL_CONV1D``
directly so that the Blackwell (sm_120) wheel-cubin gap auto-routes to
the numerically-identical pure-PyTorch fallback (no
``cudaErrorNoKernelImageForDevice`` during decode-graph capture).

Returns ``False`` (→ pure-PyTorch fallback) when any of:
  * the wheel is not importable (``_HAVE_CAUSAL_CONV1D`` is False, or a
    test monkeypatched it off);
  * ``device`` is not CUDA (CPU smoke / shape tests);
  * ``ARBI_GDN_CONV_FALLBACK=1`` (forced fallback, any arch);
  * ``device`` is Blackwell (cc major >= 12) AND
    ``ARBI_GDN_CONV_FORCE_KERNEL`` is not set.

Live read of ``ARBI_GDN_FUSED_RECURRENT`` (per-call).

Tests flip this env var at runtime to exercise the legacy branch;
the per-call read is the contract they depend on.

Compile-on integration: ``@torch._dynamo.assume_constant_result``
tells Dynamo to specialize this call's return value at trace time
and bake it into the compiled graph as a constant. The ``os.environ``
read happens once during tracing; subsequent compiled forwards
re-use the baked value. If a test flips the env at runtime, the
compile cache will recompile (Dynamo's specialization guards
notice the constant changed). Production never flips it post-boot,
so steady-state has zero env reads on the hot path.

Read the env flag at call time (NOT module-load time).

Reading at call time means tests can flip the flag inside a single
process to compare both paths without reimporting the module — the
numerical-equivalence test relies on this.

Live read of ``ARBI_GDN_VENDORED_FLA`` (per-call).

**Default ON.**  When ON, the GDN prefill path calls
:func:`arbi_serve.kernels.fla_vendored.chunk_gated_delta_rule_inference_fwd`
instead of FLA's public ``chunk_gated_delta_rule`` — bypassing
the upstream ``@torch.compiler.disable`` decorator and the
``torch.autograd.Function.apply`` indirection.  Set
``ARBI_GDN_VENDORED_FLA=0`` to opt out.

Performant options default ON.

Numerics: byte-identical when the same inputs are passed.  The
``tests/test_fla_vendored_parity.py`` regression locks this in.

Compile-on integration: ``@torch._dynamo.assume_constant_result``
bakes the env-read into the compiled graph at trace time, same
pattern as :func:`_gdn_fused_recurrent_enabled`.

Gated Delta Net — pure-PyTorch reference forward paths (CPU + small-GPU smoke).

Mixed into :class:`~arbi_serve.models.gdn_block.GDNBlock`; these methods are
the numerically-faithful oracle the unit tests run against when CUDA / FLA
are absent. They use only ``self``-attributes/methods defined on the block.

Multi-token prefill — pure-PyTorch reference (CPU / no-FLA).

Walks each request's token range from ``cu_seqlens`` (or
falls back to splitting evenly when ``cu_seqlens`` is
absent — only valid for B==1 or equal-length tests). Updates
the per-request conv + recurrent state in place.

Partial-accept rollback by recomputing from the pre-verify base
(pure-PyTorch reference — CPU / no-FLA verify path, uses saved
``beta``/``g``).

For each ``(slab_row, n)`` pair, restore the layer's recurrent +
conv slab rows to the post-token-``n`` verify state by replaying
the recurrence from the captured base
(``snap_{recurrent,conv}_state[0, row]``) over the retained
accepted-prefix inputs ``saved_*[0..n, row]``. Rollback is a
single-pass forward over the accepted prefix, with the
``T=(K+1)`` snapshot dimension collapsed to 1.

Reuses the same forward helpers (:meth:`_conv_prefill_pertoken`,
:meth:`_delta_rule_step`, :meth:`_l2norm`) the verify pass ran, so
on a shared-math device the replayed state is byte-identical to
the old snapshot read.

GDN-only, host-driven (loops the per-row accepted length); this
method backs the eager / sync rollback path.

Gated Delta Net — tree-order recurrent scan and conv-window gather.

A recurrent mixer consumes the verify block as a SEQUENCE. A tree's block
is not one: node ``j``'s row sits after its siblings in block order, and a
chain scan would fold their tokens into ``j``'s state. This module is the
geometry that removes that folding, expressed entirely as index tensors
over the SAME kernels the chain verify already runs.

Three pieces, one for each thing a tree needs from a recurrent layer:

``tree_conv_window_index``
    The conv window of node ``i`` is the last ``Kw`` tokens along ``i``'s
    OWN ancestor chain. That is a gather with a per-node index VECTOR
    where :func:`~arbi_serve.models.recurrent_common.advance_conv_window_masked`
    uses a per-row scalar offset — the same augmented array
    ``[base[..., 1:] ‖ x_0 .. x_{N-1}]``, indexed by ancestor id instead of
    by ``n + j``. Depth does not enter: one gather covers the whole tree
    because a conv window is finite, not recurrent.

``tree_scan_levelwise``
    Breadth-first: at depth ``d`` every node starts from ITS PARENT's
    state, so the level is one batched step of ``level_size`` independent
    rows. ``depth`` sequential steps for the whole tree, whatever its node
    count. The parent gather is ``repeat_interleave`` — node ids are
    assigned breadth-first and children are contiguous, so a level's
    parent map is ``arange(level) // widths[d]``, never a general gather.

``tree_scan_paths``
    Every root-to-leaf path packed as its own varlen row. Each row IS an
    ancestor chain, so the chain verify's kernel call, its conv unfold and
    its accepted-prefix rollback all apply verbatim — the tree needs no
    per-node index math at all, at the price of recomputing an interior
    node once per leaf beneath it.

The two scans are alternative realizations of one relation, and
:func:`tree_scan_per_path_independent` is the definition both are checked
against: node ``i``'s state is what you get by running ``i``'s ancestor
chain alone, from the committed state, touching no other node.

Static index tables for one ``(TreeSpec, conv_kernel, device)``.

Every field is a pure function of the geometry, which is a
process-global flag, so these are built once and handed to every layer
of every step — the same stability contract
:func:`~arbi_serve.spec_decode.tree_spec.active_tree_mask` keeps, and
for the same reason: a fresh allocation per step moves a pointer a
captured graph baked.

The verify BLOCK, as a tree in its own right.

The block a verify forward consumes is ``[committed_token, node_0 ..
node_{N-1}]``: one more row than the tree has nodes, and that extra row
is the committed token every depth-0 node descends from. Prepending a
width-1 level makes it a node — block row ``i`` becomes node ``i`` of
``TreeSpec((1, *widths))``, breadth-first ids and all — so the conv
window table, the packed leaf paths and the node-slot map address the
BLOCK directly and no ±1 shift is carried anywhere.

The recurrent state this block's scan starts from is then the state
BEFORE the committed token, which is exactly the ``h0`` the chain
verify seeds from. Same base frame, same rollback contract.

``(N, Kw)`` index into ``[base_conv[..., 1:] ‖ x_0 .. x_{N-1}]``.

Column ``j`` of node ``i``'s window is the token ``m = Kw - 1 - j``
steps back along ``i``'s ancestor chain (``m = 0`` is ``i`` itself).
When that walk runs off the top of the tree it lands in the committed
prefix, which the base buffer's own columns already hold.

The chain is the ``widths == (1,) * K`` case and reduces to the scalar
form :func:`~arbi_serve.models.recurrent_common.advance_conv_window_masked`
uses: node ``i`` sits at depth ``i`` with ``ancestor_m(i) = i - m``, so
the index is ``i + j`` — that function's ``nacc + j`` at ``nacc = i``.

``(paths, true_lengths)`` — one row per leaf, root first, RECTANGULAR.

A product tree's leaves all sit at the last depth, so its paths are
the same length and the packing is exact. A spine tree's tails end
early, so a short path is padded to :attr:`TreeSpec.depth` by
repeating its own leaf.

Padding at the END is what makes it free of consequence: those
repeats evolve only states that come AFTER the leaf, which the packed
scan never commits, and ``node_slot`` addresses a node's FIRST
occurrence, which is always the real one. Keeping the packing
rectangular in exchange keeps ``cu_seqlens`` uniform — the same
tensor the chain-shaped launch builds — so no kernel sees a shape it
has not already been served.

``device`` with its index filled in, so one device has one key.

A boot names its device both ways — the pool carries ``cuda`` and a
tensor carries ``cuda:0`` — and two keys for one device is two table
sets, which is the one thing these must never be.

Build (or reuse) the static index tables for ``spec``.

Every caller — the verify forward, the accept-time replay, the
capture fingerprint — must come through here, because what they
share is the tensor IDENTITY and not just the values: a captured
graph bakes the address it gathered through, and a second tensor
holding the same numbers is a different address.

Materialize the active tree's block tables ahead of any forward.

Called from the boot path that attaches the recurrent replay buffers,
which the tree forward cannot run without. It has to happen there and
not on first use: a first use inside a cudagraph capture allocates the
tables from the capture's PRIVATE pool, and that pool's memory is
handed out again to the next graph captured into it — so the pointer
every captured tree verify baked would then be reading another graph's
activations. Building them here puts them in the default pool, where
nothing reclaims them for the life of the process.

Tables for the active tree, IF this call's block is that tree's.

``None`` for every call a tree does not own, and the three reasons are
not interchangeable:

  * no tree is configured — the chain route, unchanged;
  * the tree is width-1 — a chain written as geometry. Its accepted
    path IS a prefix of block order and no node has a sibling to
    absorb, so the chain scan is already exact for it. Routing it here
    would trade a proven path for an equivalent one and cost the
    width-1 lane its role as the equivalence gate;
  * the block does not have the tree's width. ``mtp_block_m`` is set
    per verify step but a prefill or a plain decode can land between
    two of them, and matching on the CALL's shape rather than on a
    flag is what keeps a stale width from selecting tree geometry for
    a chain-shaped block.

Per-node causal-conv output, one gather for the whole tree.

``base_conv`` ``(B, C, Kw)`` is the committed rolling buffer,
``x_nodes`` ``(B, N, C)`` the pre-conv stream in node-id order,
``weight`` ``(C, Kw)`` the depthwise filter. Returns ``(B, N, C)``.

Depth-independent: a conv window reaches back ``Kw`` tokens and the
ancestor chain supplying them is a static table, so the whole tree's
windows are one advanced index — no per-level loop, no recurrence.

Breadth-first recurrent scan — ``depth`` sequential steps, total.

``h0`` ``(1, *S)`` is the committed state every depth-0 node starts
from. ``step(h_in, node_ids)`` runs ONE token for each of ``R``
independent rows and returns ``(out_rows, state_rows)``; it is the
caller's single-token kernel invocation, unchanged — the tree enters
only through which ``h_in`` rows it is handed.

Returns ``(out_nodes, last_level_states)`` with ``out_nodes`` in node-id
order. The sequential step count is ``depth``, NOT ``num_nodes``: a
level's rows are mutually independent, so they go in one call.

Path-packed recurrent scan — ONE call, ``depth`` tokens per row.

Every root-to-leaf path is its own varlen row, so ``step_seq(h0_rows,
node_ids)`` is the chain verify's existing call with ``B = num_leaves``
and ``T = depth``. It returns ``(out (L, D, *O), final (L, *S))``.

Interior nodes are recomputed once per leaf beneath them; every copy is
the same arithmetic on the same inputs, so ``tables.node_slot`` may pick
any one. The redundancy factor is ``leaves * depth / num_nodes`` — 1.71
for ``2x2x2``, and exactly 1.0 for a chain, where there is one path.

The DEFINITION: each node's state from its own ancestor chain alone.

``num_nodes`` separate walks, each starting from the committed state and
touching only that node's ancestors — no batching, no shared work, no
opportunity for a sibling to leak in. Deliberately the slow form: it is
the oracle the two production-shaped scans above are checked against,
and it is correct by construction rather than by argument.

Returns ``(out_nodes, per_node_states)``.

Shared :class:`LayerStackModelMixin` for v2 per-arch model files.

The seven per-arch models (qwen3, qwen3_5, deepseek_v3, lfm2,
lfm2_moe, gemma4, nemotron_h) all share the same construction +
forward-loop boilerplate:

  * build a :class:`LayerStack` from ``layer_specs`` + a parallel
    list of per-spec block converters;
  * route every per-layer call through
    :func:`arbi_serve.runtime.capture.dispatch.model_dispatch` via a
    per-forward closure that captures ``model``, ``num_seqs``,
    ``is_prefill``;
  * materialize meta-device parameters into empty CUDA tensors of the
    target dtype after a quant-backend swap (``apply_quant_if_present``);
  * filter the ``weight_map`` to drop ``.weight`` entries for linears
    the quant backend swapped out.

This mixin lifts those repeated bodies to one shared site. Per-arch
files declare the spec → block dispatcher (the single ``if kind ==``
switch) and the per-step extras they thread through ``**extra``; they
do not re-implement the dispatch wrapper, the meta-materialize walk,
or the quant filter every time.

Following the "no back-compat shims" rule, this is the **single
canonical** boilerplate path — no legacy variants survive.

The mixin is plain Python (no :class:`nn.Module` inheritance) so per-
arch classes mix it in alongside ``nn.Module`` without a metaclass
conflict.

Shared construction + forward-loop boilerplate for v2 models.

Mix into the per-arch ``nn.Module`` subclass. The mixin assumes the
subclass exposes:
  * ``self.dims`` — :class:`ModelDims`;
  * ``self.layer_specs: list[LayerSpec]`` after ``__init__``
    finishes wiring the arch-specific submodules.

Build the ``batch_meta`` / ``state_views`` dicts the
:class:`LayerStack` iterator threads to each block.

Walks ``specs`` (defaults to ``self.layer_specs``) and
populates only the kinds the model uses, pulling each kind's
metadata off the live :class:`ScheduledBatch`. The
:class:`LayerStack` orchestrator pre-resolves each layer's
per-layer view from the supplied :class:`MultiStatePool` (via
:meth:`MultiStatePool.per_layer_views`) and threads the
resolved entry through ``block_args[2]`` — block forwards
consume ``state_view`` as the per-layer slab handle directly,
no inner ``layer_view(self.layer_idx)`` call. Keeping
``layer_idx`` out of the traced block is what stops Dynamo
from specializing (and recompiling) per layer index.

``StateKind.NONE`` (MLP-only layers) is intentionally absent
so the dict ``.get`` returns ``None`` and the MLP adapter
discards both metadata and slot view.

Build a :class:`LayerStack` from spec / block lists.

``blocks`` is the per-arch dispatcher's output (one inner
decoder layer per spec, in spec order). The mixin does not see
the ``if kind ==`` switch — that's the per-arch ``_make_*``
function.

Type-erased to ``Sequence[Any]`` because :class:`LayerStack`
types blocks as :class:`nn.Module` but its own tests pass
:class:`MagicMock`s; the runtime contract is callable, not
module.

Build the per-forward dispatch closure.

Routes every :class:`LayerStack` block call through
:func:`piecewise.model_dispatch`. ``batch.num_seqs`` /
``batch.is_prefill`` are bound by closure so the LayerStack
iterator's wrapper signature stays uniform across kinds.

Legacy kwargs path — kept for tests and any caller that still
threads through :meth:`LayerStack.forward` (the kwargs API).
Production model forwards prefer :meth:`_dispatch_layer_args`
below, which is closure-free and Dynamo-traceable.

Closure-free per-layer dispatcher.

Bound method on the model — Dynamo traces the call as a
method invocation against ``self`` (a stable key in the
symbolic registry), not a freshly-defined Python closure.
``self._step_num_seqs`` / ``self._step_is_prefill`` are
scalar attribute writes set once at the top of every model
``forward`` before the LayerStack walk begins.

Block args are spread positionally:
``block(hidden, positions, meta, view, *extras)``. Recurrent
blocks that ignore some of those (``positions`` for Mamba /
ShortConv, ``rope_cache`` / ``attn_ops`` for everything that
isn't a paged-attention block) accept ``*_unused`` and
discard.

The model-level decorator installs a marker trampoline (eager
passthrough), so this dispatcher always runs outside any Dynamo
trace. The pool branch fires when the engine has set
``_piecewise_pool`` (boot capture sweep + chunked-prefill replay);
otherwise falls through to a direct ``block(hidden, *block_args)``
call. The layer-level trampolines do their own Inductor compile
when the engine invokes ``layer(...)`` from inside the dispatcher.

Per-layer attn-op de-specialization: when ``block`` carries the
``_arbi_attn_op_extras_idx`` marker, the named ``extras`` slot
holds the full per-layer ``attn_ops`` list. Indexing it by
``self.layer_idx`` inside the (compiled) block forward forces
Dynamo to recompile once per ``layer_idx`` value (the list is an
opaque Python object; the int index becomes a guard). We resolve
``attn_ops[layer_idx]`` here — this dispatcher always runs eager,
outside the trace — and thread the single resolved op into the
block instead, collapsing those per-layer specializations to one
cache slot (mirrors the pre-resolved per-layer ``view``).

The activation arena rides the same seam: when
:attr:`arena_extras_idx` names a live arena, that slot is
rewound and replaced with the per-layer buffers the block asked
for (``block.alloc_arena_buffers``), so the compiled block sees
tensors and never the allocator.

The one lm_head dispatch point: ``lm_head(last_hidden)`` (+ epilogue).

EVERY arch's forward calls the head through here so a single place
honours ``batch.logits_out`` — the capture-time destination buffer
that keeps a full ``(B, vocab)`` row out of the cudagraph capture
pool for each captured shape (see
:attr:`arbi_serve.engine.batch.ScheduledBatch.logits_out` and
:func:`arbi_serve.models.linear.linear_into`). ``batch.logits_out``
is ``None`` on every live step, where this is exactly the previous
``logits = self.lm_head(last_hidden)``.

``scale`` (granite's ``logits_scaling``) and ``softcap`` (gemma's
``final_logit_softcapping``) are applied IN PLACE when writing into
a destination buffer, so the epilogue does not re-introduce the
vocab-sized allocation the destination was there to avoid. In place
or not, the arithmetic — and hence the result — is the same.

``lm_head`` over an ARBITRARY set of post-final-norm rows.

The logprobs surface scores prompt positions, which
:meth:`_finalize_logits` discards when it gathers the per-sequence
last token. This is the same head and the same epilogue applied to
a caller-chosen ``(rows, hidden_size)`` slice instead.

``batch=None`` so ``batch.logits_out`` never binds: that
destination buffer is pinned to a captured shape, and these rows
are a tile of the caller's choosing. The granite ``scale`` and
gemma ``softcap`` epilogues resolve off the model, so an arch that
declares either gets it here without a per-arch call site.

This head's declared logit epilogue, as
:func:`~arbi_serve.models.linear.apply_logit_epilogue` keywords.

ONE naming of the three attributes an arch may declare, so a caller
that projects the head ITSELF — the sharded verify's per-shard GEMM,
which cannot go through :meth:`logits_from_hidden` — composes the same
epilogue as :meth:`_lm_head` without restating where the terms live.
Every term is ``None`` on an arch that declares none, which
``apply_logit_epilogue`` skips.

Final RMSNorm (optional) + last-token gather + lm_head epilogue.

Shared tail of the per-arch model ``forward``. When ``norm`` is
given it is the final-norm module applied to ``hidden`` first;
pass ``None`` when the caller already normed (e.g. it needs the
normed hidden for an earlier return). ``softcap``, when truthy,
applies ``tanh`` logit-softcapping (gemma) after ``lm_head``.

Logits stay in the model's working dtype (bf16 in production);
the sampler promotes where it needs fp32.

Return the hidden representation consumed by ``lm_head``.

Most bundled draft heads consume the same final hidden tensor that
feeds the language-model head. Architectures with a wider recurrent
draft state override this seam and collapse it only for logits.

Boot-time post-construction steps shared across arches.

Sequence:
  1. Quant-backend swap (no-op on dense). Runs BEFORE any meta
     tensor is materialized so we never allocate dense tensors
     that the swap would immediately discard — at 9B+
     GPU.
  2. Replace remaining meta tensors with empty CUDA tensors of
     the target dtype so :func:`load_model_weights` can
     ``copy_`` into them. Linears the swap rebound have no
     ``.weight`` Parameter so the walk skips them naturally
     (they expose backend-specific buffers instead).

Cold-boot peak control: when the engine's dense cold path has
installed an
:class:`arbi_serve.loader.flat_loader.AllocateThenFillBinder`
(dense checkpoints only — no quant backend swapped anything),
step 2 binds every remaining meta tensor to a flat-pool slab
view instead of a standalone ``torch.empty``. The weight loader
then fills the slabs in place, so the boot peak stays at ≈ 1×
model size and ``weights_pool`` ends frag~0 with no post-load
compaction pass.

Mamba-2 MTP verify forward + partial-accept rollback.

Mixed into :class:`~arbi_serve.models.mamba2_block.Mamba2Block`. Mirrors
the GDN construct pair in
:mod:`arbi_serve.models._gdn_fla_verify_chained`: a per-token verify
forward that leaves a 1-deep base frame plus the retained per-token
recurrence inputs, and the two rollback modes those buffers select
(masked replay / recompute-from-base).

Optional GPU kernel globals live on the block module and are read through
:func:`_block_module` so unit tests keep monkeypatching ONE namespace.

The ``mamba2_block`` namespace holding the optional-kernel globals.

Resolved lazily so this module and its host can be imported in either
order, and so the globals are read at CALL time — the single surface a
test monkeypatches.

MTP verify pass — T = K+1 tokens per row through the PER-TOKEN
decode kernels.

The varlen chunk scan (:meth:`_forward_kernel_prefill`) commits
only the post-all-K+1 state: nothing observes the intermediate
positions, so a partial accept has nothing to return to. This
route walks ``t = 0 .. T-1`` and at each step advances BOTH slabs
one token with the same ``causal_conv1d_update`` +
``selective_state_update`` pair the K=1 decode path runs, which is
what makes the accepted-prefix state reachable. vLLM reaches the
same place by raising its ``reorder_batch_threshold`` so verify
rows land in the decode bucket.

Alongside the forward it captures the rollback inputs: the 1-deep
pre-verify BASE frame (``snap_{ssm,conv}_state[0, slab_row]``) and,
per token, either the ``(N, T, …)`` masked-replay buffers or the
``(T, N, …)`` recompute buffers — whichever
:meth:`RecurrentStatePool.attach_mtp_snapshot_buffers` allocated.

``T = T_total // B`` is shape-derived (no host sync); a
non-uniform ``T_total`` is a misrouted caller and raises.

Partial-accept rollback by recomputing from the pre-verify base.

For each ``(slab_row, n)`` pair, restore the layer's conv + ssm
slab rows to the post-token-``n`` verify state by replaying the
recurrence from ``snap_{ssm,conv}_state[0, row]`` over the
retained ``saved_{x_conv,dt}[0..n, row]`` inputs. Batched across
the partial-accept rows: one step per ``t = 0 .. max(n_accepted)``,
each row committed at its own ``t == n``. Fully-accepted rows
(``n == k_uniform``) keep the forward's post-token-K state — a
provable no-op — and are skipped.

Driven by :meth:`RecurrentStatePool.rollback_batch`; host-driven,
so it carries no cudagraph-capture constraint. Pure-torch on both
CPU and CUDA: the Mamba-2 step is elementwise over
``(P, H, head_dim, state_dim)``, with no kernel to match.

Commit each row's accepted-prefix state by MASKED replay.

``nacc_by_row`` is the per-SLAB-ROW accepted offset (−1 = leave
the row untouched); ``step_mask`` is the ``(N, T)`` bool
``t <= nacc`` prefix mask the pool builds once per tick. Per
layer this issues:

  * one ``dt → -inf`` input-masking write over the saved raw
    ``dt``. ``softplus(-inf) == 0`` exactly, so a masked step has
    ``dA = exp(0 · A) = 1`` and ``dB = 0 · B = 0`` — the exact
    identity, the Mamba-2 analogue of GDN's ``a → -inf`` /
    ``k → 0`` pair;
  * the base frame written onto the selected slab rows, then the
    ``T`` per-token :func:`selective_state_update` steps the verify
    forward itself ran, replayed over the saved POST-conv stream
    (:meth:`_replay_masked_steps_kernel`);
  * the shared pure-gather conv window advance
    (:func:`advance_conv_window_masked`).

Fixed shapes, fixed trip count, no host reads — stream-capture-
safe; the pool captures the per-layer launch set into a side
cudagraph. The in-place ``dt`` masking destroys only slots that
are dead after this tick (the next verify forward re-saves its
slate's rows).

CPU / no-``mamba-ssm`` falls back to the pure-PyTorch batched
replay (:meth:`_replay_masked_scan_reference`), which applies the
same ``-inf`` mask through the same identity.

``n_steps`` is accepted for dispatch parity and unused: this replay is
a per-token :func:`selective_state_update` over the whole slab with
``state_batch_indices=None``, which has no per-sequence trip count to
shorten. The pool refuses to stage one for a Mamba view.

Advance ``ssm_slab`` in place over the T replay steps with the SAME
per-token :func:`selective_state_update` the verify forward ran.

The forward (:meth:`_forward_verify`) advances the slab one token at a
time through ``selective_state_update``; a rollback that reproduces the
post-token-``n`` state must therefore replay those steps through the
same kernel, in the same order, from the same base — anything else
commits a state the plain K=1 decode path would never hold, and greedy
speculative decoding stops being lossless.

Runs over ALL ``N`` slab rows with ``state_batch_indices=None``: the
per-row recurrence is independent, and a row the caller did not select
carries ``dt = -inf`` at every step (``softplus(-inf) == 0`` ⇒
``dA = 1``, ``dB = 0``), so it is advanced by an exact identity and its
slab bytes are unchanged. Fixed shapes and a fixed trip count, no host
reads — stream-capture-safe.

Pure-PyTorch batched masked replay (CPU / no-``mamba-ssm`` pools).

Math-equal replay of the accepted prefix from the base frame over
the saved post-conv stream, using the same batched step helper the
recompute path runs. Masked steps carry ``dt = -inf`` and are the
identity, so the whole ``(N, T)`` slab is walked unconditionally.

One rule for how a checkpoint's routed experts are served.

Every MoE arch faces the same question — can these experts be stacked into
:class:`~arbi_serve.models.moe.FusedMoE`, or must they stay per-expert
modules? — and the answer belongs to the quant backend, never to the arch.
An arch that decides this locally either names a backend (and refuses every
other pack, including plain bf16) or re-derives half the rule and silently
loses the stacked path for quants it could have fused.

Resolve the expert representation for ``expert_paths``.

Fused when the pack is dense (no quant backend at all) or when its
backend answers :meth:`~arbi_serve.weight_quant.base.QuantBackend.stacked_moe_spec`
with a spec. Per-expert otherwise — a payload the fused kernels cannot
read (an EXL3 trellis, a W4A8 pre-scale) keeps the per-linear modules.

Binding for checkpoints that ship ONE 4-D tensor per expert stack.

:func:`bind_stacked_expert_stacks` reads each stack tensor once and
scatters it across every local expert slot of a
:class:`~arbi_serve.models.moe.FusedMoE`.

Source layouts:

  * ``quant_kind="mxfp4"`` — ``{stack}_blocks`` ``(E, R, K // 32, 16)``
    uint8 packed E2M1 and ``{stack}_scales`` ``(E, R, K // 32)`` uint8
    E8M0, plus ``{stack}_bias`` ``(E, R)``.
  * dense — ``{stack}`` with the contracted and output axes declared by
    ``dense_transpose``, plus optional ``{stack}_bias`` ``(E, R)``.

``gate_up_interleaved`` declares that the fused ``2I`` axis alternates
gate and up per row; ``FusedMoE.w13_weight`` stores ``[gate | up]``, so
such a source is split by parity. The expert axis takes the EP slice and
the intermediate axis the ``moe_tp`` slice.

Stacked-expert weight-map construction — arch-agnostic.

:func:`add_stacked_moe_expert_keys` binds a checkpoint's per-expert
tensors straight into the two STACKED
:class:`~arbi_serve.models.moe.FusedMoE` buffers (``w13_weight`` /
``w2_weight``, plus their ``_scale`` grids on a block-wise fp8 pack).
Every rule in here is a property of the fused layout, not of any model
family:

  * **EP slice** — each rank binds only its contiguous
    ``n_experts // expert_shard_count`` experts, remapping the global
    checkpoint expert id (src) onto the local stack slot
    (``dst_expert``). At ``expert_shard_count == 1`` the map is the
    identity over all experts. ``expert_shard_count`` is ``ep_size``
    (expert shards ACROSS TP groups) times ``moe_ep_size`` (expert
    shards WITHIN a TP group, i.e. ``--enable-expert-parallel``) — one
    number, so this builder is blind to which of the two supplied it.
  * **TP cut** — the intermediate dim shards across the ranks that SHARE
    this rank's expert shard (``moe_tp_size``, == ``tp_size`` whenever
    ``moe_ep_size == 1``): gate/up lose rows, down loses columns. The
    cut is declared ONLY when real so a transform-free spec keeps the
    loader's direct-DMA path eligible (at ``moe_tp_size == 1`` a no-op
    narrow would cost a host staging copy for every one of the thousands
    of expert tensors — which is exactly the full-expert-parallel case,
    so that topology gets the DMA path even at ``tp_size > 1``).
  * **Gate/up seam** — ``gate_proj`` occupies rows ``[0, N)`` of the
    ``w13`` slot and ``up_proj`` rows ``[N, 2N)``: the order the fused
    kernel's SiLU-and-multiply assumes.
  * **Block-wise fp8** — the ``weight_scale_inv`` grid lives on the same
    axes as the weight, so it shards and offsets by the same rule in
    units of scale blocks, and the seam must land on a block boundary.
  * **NVFP4** — the modelopt ``{weight, weight_scale, weight_scale_2}``
    triple is already in the fused kernel's axis order, so it binds raw;
    the per-tensor ``weight_scale_2`` is a 0-d scalar that lands in the
    ``gscale`` row of its own gate/up half. See
    :func:`_add_nvfp4_expert_keys`.
  * **AWQ int4** — the checkpoint's ``{qweight, scales, qzeros}`` triple
    is K-major and packed along the OUTPUT axis, which is neither of the
    two axis orders the fused kernel reads, so each spec carries a
    ``transform`` and its cut is declared on the TRANSFORMED tensor in
    the destination's storage units. See
    :func:`_add_awq_int4_expert_keys`.

The only arch-shaped inputs are the dst/src prefixes and ``proj_names``
(the checkpoint's per-expert linear names — HF-standard
``{gate,up,down}_proj`` by default; LFM2 stores ``w1``/``w3``/``w2``).
Router gate, shared experts, and expert-score bias keys stay arch-side:
their naming and presence genuinely vary per family.

View a raw quant payload as uint8 without touching its bits.

An MXFP4 pack stores E2M1 pairs and E8M0 exponents in whatever signed
or float8 container the exporter chose; the kernel indexes bytes. This
is a reinterpretation, not a conversion — a ``.to(torch.uint8)`` would
round the values and destroy the payload.

Bind this rank's expert slice into the stacked FusedMoE buffers.

``dst_prefix`` is the :class:`~arbi_serve.models.moe.FusedMoE`
module path (e.g. ``model.layers.3.mlp.experts``); ``src_prefix`` is
the checkpoint's per-expert prefix WITHOUT the expert index (expert
``E``'s tensors are ``{src_prefix}.{E}.{proj}.weight``).
``proj_names`` is ``(gate, up, down)`` in that order — gate and up
stack into ``w13_weight`` (gate rows first), down binds to
``w2_weight``.

Entries are keyed by a binding id
(``{dst_prefix}.w13_weight#e{slot}.gate``) and carry ``dst_param`` —
a dict cannot key ``3 × n_local`` checkpoint tensors under two
parameter names.

Bind a legacy-AWQ int4 expert slice into the stacked buffers.

The checkpoint triple (``qweight`` / ``scales`` / ``qzeros``) is
K-major and packed along the OUTPUT axis; the fused kernel wants an
N-major stack packed along K (weights) and along N (zero-points). So
each spec carries a ``transform`` from
:mod:`arbi_serve.weight_quant.awq.moe_repack`, and its ``shard_dim``
indexes the TRANSFORMED tensor: the intermediate dim is dim 0 for
gate/up and dim 1 for down, in the destination's own storage units
(halved for a nibble-packed axis, divided by the group for a scale
or zero-point axis).

Bind a modelopt NVFP4 expert slice into the stacked buffers.

The on-disk triple is already N-major and packed along K, which is the
order the fused kernel reads, so weight and block scale bind raw at
their checkpoint dtype. Their TP cuts differ only in UNIT: the gate/up
cut is on the output rows, shared by both; the down cut is on K, which
is ``n_local_rows // 2`` bytes of weight and ``n_local_rows //
NVFP4_BLOCK_SIZE`` scale entries — each derived from the source's own
extent, so one ``shard_dim`` covers both. The scale grid's ROW axis is
the output rows, so the gate/up seam offsets by ``n_local_rows`` in
both buffers. ``weight_scale_2`` is a scalar per source linear and is
replicated (it rides no axis).

Per-arch construction bundle for :class:`Qwen3_5Model`.

Collects the ctor args that don't fit on :class:`ModelDims` (which
is shared across all archs) but are model-arch-specific knobs
derived from ``text_config``. The arch's :meth:`from_safetensors`
builds one of these and the constructor consumes it (rather than
threading 4 positional args through the call site).

Args:
    gdn_cfg: per-layer GDN config — same instance shared across
        every linear-attention layer (Qwen 3.5 / 3.6 use one
        consistent GDN shape across the whole stack).
    rotary_dim: number of head channels carrying partial RoPE
        (``head_dim * partial_rotary_factor``). Qwen 3.5 / 3.6
        default to ``head_dim * 0.25``.
    mtp_attn: shape of the bundled MTP draft head's attention,
        mirroring the main-model full-attention layers. ``None``
        when the safetensors carries no ``mtp.*`` keys.

Parsed ``text_config`` bundle produced by
:meth:`Qwen3_5Model._parse_text_config`.

Carries the raw HF config dicts plus the derived
:class:`ModelDims`, per-layer :class:`LayerSpec` list, GDN/RoPE
knobs, and the attention head shape — everything the rest of
:meth:`Qwen3_5Model.from_safetensors` needs to append the MTP
marker, resolve the vision tower, and construct the model. The
``layer_specs`` list is mutated in place (the MTP marker is
appended) by :meth:`Qwen3_5Model._maybe_append_mtp_spec`.

Qwen 3.5 / 3.6 gated MLP — the shared :class:`GatedSiLUMLP` in
``lazy_fuse_forward`` mode: separate ``gate_proj`` / ``up_proj`` params
(quant-swappable + per-projection LoRA) with a lazily-built ``[gate|up]``
fused forward matmul on the CUDA non-LoRA hot path. The fusion logic
lives once in the shared MLP; this subclass only pins the mode + the
name the Qwen3.5 decoder layers and their tests reference.

One Qwen 3.5 / 3.6 attention decoder layer.

Implements the uniform :class:`LayerStack` block signature directly.
Body (cross-layer fusion ON, ``residual_buf is not None``):
``input_layernorm(hidden, residual_buf) → self_attn → +residual_buf →
post_attention_layernorm → mlp → return mlp_out``.

The ``residual_buf`` is the model-owned persistent slab living in
``graph_buffers_pool`` — its ``data_ptr`` is baked into captured
layer-graphs at capture time and reused across replays. Threading it
through the LayerStack ``extras`` tuple keeps the captured-graph
dispatch surface 1-tensor-in / 1-tensor-out (the dispatched
``hidden`` is now the ``pending_add`` from layer N → layer N+1, NOT
the running residual). Fusion fires per layer:

  1. ``input_layernorm(pending_add, residual_buf)`` — fuses the
     cross-layer ``residual + h`` from layer N-1 into layer N's
     input norm. ``residual_buf`` is mutated in place to
     ``residual_buf + pending_add``.
  2. ``self_attn`` runs over the normed hidden.
  3. ``post_attention_layernorm(attn_out, residual_buf)`` — the
     per-layer post-attn fusion.
  4. ``mlp`` runs over the normed hidden, returns ``mlp_out``.

The returned ``mlp_out`` becomes ``pending_add`` for the next
layer's input_layernorm fusion. The cross-layer add ``residual_buf
+= mlp_out`` happens INSIDE the next layer's input_layernorm
fusion, NOT at the exit of this layer — that's the launch saving.

Backwards-compatible: when ``residual_buf is None`` (eager path /
legacy callers / non-fused tests), falls back to the pre-fusion
body (``residual = hidden; h = input_layernorm(hidden); ...; return
residual + h``). Bit-equivalent to the old path under that gate.

Split-attn. The forward chains three explicit Python
methods: :meth:`_pre_attn` (input_layernorm + fused QKV proj +
Q/K norm + RoPE) → :meth:`_attn_eager` (the ``arbi_serve.attn``
op call only) → :meth:`_post_attn` (sigmoid gate + o_proj + post
layernorm + MLP). When ``compile_config.split_attn`` is ON the
decorator's ``methods=("_pre_attn", "_post_attn")`` kwarg compiles
each non-attn piece independently and the orchestrator captures
them as TWO cuda graphs per layer per bucket with the eager attn
call running between replays — attention kernel scratch (FA
varlen / paged_attention) lives outside the captured-graph
mempool. When ``split_attn`` is OFF the methods are still defined
+ chained from ``forward`` but the decorator falls through to
:func:`install_torch_compiled` and compiles the whole forward as
one captured graph (byte-equivalent to the non-split surface).

One Qwen 3.5 / 3.6 GDN decoder layer.

The slab-mutation contract is correct end-to-end under compile:
the production hot paths route through
``arbi_serve::gdn_attention_v2`` which declares
``mutates_args=("recurrent_state_slab",)`` and does the slab
gather + FLA-kernel + slab scatter INSIDE the opaque op so
Inductor cannot reorder around the in-place mutation. See
``arbi_serve._custom_ops::_register_gdn_attention_v2`` and
:meth:`GDNBlock._dispatch_through_custom_op_v2`.

Decorated with ``@support_torch_compile(level="block",
methods=("_pre_gdn", "_post_gdn"))`` so the GDN body — which
carries the FLA chunk_gated_delta_rule kernel and its scratch —
runs eager between two compiled+captured pieces when
``compile_config.split_attn`` is ON. With split_attn OFF the
decorator falls through to the legacy whole-forward install
path; ``forward`` chains the triple either way so the numerics
are identical (split_attn only changes which set of callables
Dynamo / the orchestrator see).

Compile-path shape and kernel-surface notes:

  1. **Symbolic-shape view in ``_forward_prefill_fla``.**
     Under layer compile, ``T_total`` (s59) and ``B`` (s13)
     come in as independent SymInts.  ``T_per_row = T_total //
     B`` is computed and ``torch._check(T_per_row * B ==
     T_total)`` registers the divisibility relation with
     Dynamo's symbolic-shape solver, which then accepts the
     ``view(B, T_per_row, conv_dim_local)`` reshape.  Same
     pattern in :meth:`GDNBlock._forward_verify_fla`.

  2. **Packed-decode kernel surface.**
     Wrapping the packed-decode Triton JIT as
     ``arbi_serve::gdn_packed_decode_v2`` lets the compile path
     also take packed-decode (Dynamo otherwise refuses to trace
     ``triton.knobs`` and routes to the slower
     ``_forward_decode_fla``).

Implements the uniform :class:`LayerStack` block signature directly.
Body (cross-layer fusion ON): identical structure to
:class:`_Qwen3_5AttentionDecoderLayer` — fused input_layernorm at
entry, fused post-attn norm before MLP, return ``mlp_out`` (the
pending-add for the next layer). See that class's docstring for
the mutation contract; this class swaps the attention block for
:class:`GDNBlock` (linear attention).

Split-gdn. Three standalone methods —
:meth:`_pre_gdn` (input_layernorm fusion → ``h``) →
:meth:`_gdn_eager` (the ``GDNBlock`` forward — projections, conv,
FLA chunk_gated_delta_rule / fused_recurrent_gated_delta_rule,
materialize) → :meth:`_post_gdn` (post layernorm fusion + MLP) —
are exposed for the **piecewise prefill capture orchestrator**
only. Under split_attn=True the orchestrator captures pre/post
as TWO cuda graphs per layer per bucket and runs ``_gdn_eager``
between replays; the FLA chunk-fwd workspace, conv buffer,
projection intermediates, and ``core_attn_out`` allocation all
live in the regular allocator pool, NOT the captured-graph
mempool — same captured-pool reduction the attention split
delivered, applied to the 18 GDN layers (vs 6 attention layers)
on Qwen3.5-0.8B.

Runtime callers (decode whole-forward replay, prefill fallback,
eager smoke) go through the MONOLITHIC :meth:`forward` body —
NOT the chained ``_pre_gdn`` / ``_gdn_eager`` / ``_post_gdn``
triple. Chaining the triple from ``forward`` would force
Inductor to compile each piece as a separate callable and lose
cross-method fusion opportunities at decode (chaining the triple
in ``forward`` regresses B=1 decode; the monolithic body
avoids that — the orchestrator's piecewise capture path still
calls ``layer._pre_gdn`` / ``layer._gdn_eager`` /
``layer._post_gdn`` directly so the captured-pool reduction is
preserved).

Allocate this layer's arena-backed attention output buffer.

Called eagerly from the model's per-layer dispatch seam, which
puts the result into ``extras`` slot 4 in the arena's place, so
the compiled block forward receives a tensor and never the
allocator (``ActivationArena.alloc`` is not traceable).

The GDN layers of this same stack have no such method, so the
dispatcher blanks their slot: a recurrent layer's transients are
written into the persistent state slab and outlive the step,
which is exactly what the arena must not hold.

Pre-attn: input_layernorm fusion → AttentionBlock._pre_attn.

Returns ``(q, k, v, gate_flat)`` — the input shape for
:meth:`_attn_eager`. ``AttentionBlock`` implements the matching
``_pre_attn`` primitive and we delegate to it.

``lora_state`` (per-step adapter state, ``None`` for non-LoRA
traffic — short-circuits inside the projections) is threaded to
``AttentionBlock._pre_attn`` so q/k/v LoRA applies.

Eager attn op call — runs OUTSIDE captured graphs under split_attn.

``out_buf`` is the activation arena's slot for this layer when one
is live (``None`` under split_attn, whose replay path calls this
from the captured seam without one, and whenever the arena is off).

Post-attn: AttentionBlock._post_attn → post layernorm fusion + MLP.

When ``residual_buf is not None`` the fused
``post_attention_layernorm`` mutates ``residual_buf`` in place
to ``residual_buf + o_proj_out`` and we return ``mlp_out`` as
the pending-add for layer N+1's input_layernorm fusion. The
legacy ``residual_buf is None`` path reconstructs ``residual +
h`` from ``self._eager_residual_buf`` stashed by
:meth:`_pre_attn`.

``lora_state`` is threaded to ``AttentionBlock._post_attn``
(o_proj LoRA) and to the MLP (gate/up/down LoRA). ``None`` for
non-LoRA traffic — short-circuits in each projection.

Run the gated-attention block then MLP, each with a pre-norm
residual (fused into ``residual_buf`` when split-attn is active).

``lora_state`` rides the LayerStack ``extras`` tuple (slot 3,
after ``residual_buf``); ``None`` for non-LoRA traffic, in which
case every projection short-circuits and the path is
bit-identical to the pre-LoRA-wiring forward.

``attn_out_buf`` (slot 4) is this layer's attention output
buffer, allocated from the activation arena at the model's eager
dispatch seam — see :meth:`alloc_arena_buffers`. ``None`` when
the arena is off, which is the default and is byte-identical to
allocating it inside :meth:`_attn_eager`.

Pre-gdn: input_layernorm fusion → ``h: (N, hidden_size)``.

Returns the post-input-layernorm hidden state — the input to
:meth:`_gdn_eager`. Under cross-layer fusion (``residual_buf
is not None``) the layernorm mutates ``residual_buf`` in
place (``residual_buf += hidden``) AND returns the normed
view; same contract as the attention layer's ``_pre_attn``.

Eager GDN op call — runs OUTSIDE captured graphs under split_attn.

Calls :class:`GDNBlock`.forward over ``h`` and returns the
per-token GDN output ``(N, hidden_size)``. Under
:data:`compile_config.split_attn`, this method runs OUTSIDE
the captured-graph blocks — its conv buffer, FLA chunk-fwd
workspace, projection intermediates, and the returned
``core_attn_out`` allocation all land in the regular
allocator pool, not the captured-graph mempool. That is the
captured-pool reduction.

Post-gdn: post layernorm fusion + MLP → ``mlp_out: (N, hidden_size)``.

When ``residual_buf is not None`` the fused
``post_attention_layernorm`` mutates ``residual_buf`` in
place to ``residual_buf + gdn_out`` and we return ``mlp_out``
as the pending-add for layer N+1's input_layernorm fusion.
The legacy ``residual_buf is None`` path reconstructs
``residual + mlp_out`` from ``self._eager_residual_buf``
stashed by :meth:`_pre_gdn`.

``lora_state`` is threaded to the MLP (gate/up/down LoRA). The
GDN mixer itself uses fused qkvz projections that have no
standard PEFT target, so no LoRA applies to the mixer. ``None``
for non-LoRA traffic.

Run the Gated Delta Net mixer then MLP, each with a pre-norm
residual (fused into ``residual_buf`` for cross-layer fusion).

``lora_state`` (slot 3 of the LayerStack ``extras``) is threaded
to the MLP only — the GDN mixer's fused qkvz projections have no
standard PEFT target. ``None`` for non-LoRA traffic (the MLP
projections short-circuit, leaving this path bit-identical).

Boot-time config parsing + vision-tower decision logging for Qwen 3.5 / 3.6.

Non-hot-path helpers: ``config.json`` / ``text_config`` parsing into a
:class:`_ParsedTextConfig`, the bundled-MTP marker append, the optional
vision-tower config resolution, and the one boot line for the vision
load/skip decision.

Emit one boot line for the vision-tower load/skip decision.

Three cases:

  * No ``vision_config`` in the checkpoint → text-only model, nothing to
    report (no line).
  * Checkpoint has a tower and ``enable_vision`` is on → "vision tower:
    loaded" so the path is visible.
  * Checkpoint has a tower but ``enable_vision`` is off (the default /
    ``--text-only``) → "vision tower: skipped (text-only)" with the count
    of ``visual.*`` tensors and the GB not loaded. The byte count is the
    exact on-disk sum of the

Log the vision-tower load/skip decision and record the flag's effect.

The one seam every Qwen-VL-family arch calls at boot, so the boot line
and the ``ARBI_ENABLE_VISION`` flag-truth accounting say the same thing
whichever arch is being served.

Parse ``config.json`` + its nested ``text_config`` into a
:class:`_ParsedTextConfig`.

Validates ``model_type`` (``"qwen3_5"`` dense by default; the MoE
arch passes ``"qwen3_5_moe"`` — same attention/GDN interleave, the
FFN width lives in ``moe_intermediate_size`` instead of
``intermediate_size``), ``attn_output_gate``, the ``layer_types``
length, and TP-divisibility of the attention head counts, then
derives :class:`ModelDims`, the per-layer :class:`LayerSpec` list
(GDN/attention interleave), the shared GDN config, and the
partial-RoPE ``rotary_dim``.

Detect a bundled MTP draft head and, when present, append its
marker :class:`LayerSpec` to ``layer_specs`` (mutated in place)
and bump ``dims.num_layers`` by one.

Returns ``(mtp_attn, dims)`` — ``mtp_attn`` is the head's
attention shape (``None`` when the checkpoint carries no
``mtp.*`` keys) and ``dims`` is the (possibly slab-bumped)
:class:`ModelDims`.

Resolve the optional multimodal vision tower config.

Bound only when the checkpoint has a ``vision_config`` and
``ARBI_ENABLE_VISION`` is set; default off keeps text-only
boot/VRAM bit-identical. Emits one boot line either way.
Returns ``(vision_cfg, image_token_id, mrope_section)`` — all
``None`` on the text-only path.

Partial-RoPE caches for Qwen 3.5 / 3.6 (partial_rotary_factor = 0.25).

The text-only :class:`_PartialRoPECache` and the multimodal
:class:`_MRoPEPartialCache`. Self-contained ``nn.Module`` cos/sin
tables; the model holds one instance as ``_rope_cache`` and every
attention layer calls ``.apply``.

NEOX-style RoPE applied to the first ``rotary_dim`` channels only.

Qwen 3.5 / 3.6 set ``partial_rotary_factor = 0.25``: with
``head_dim = 256`` only the first 64 channels per head carry
rotation. The rest pass through unchanged. The cos / sin tables
have shape ``(max_seq_len, rotary_dim)``.

This is structurally identical to :class:`RoPECache` but indexed
by ``rotary_dim`` rather than the full head dim, and ``apply``
splits Q / K and only rotates the leading half of channels.

Interleaved M-RoPE over the leading ``rotary_dim`` channels.

Qwen-VL spreads RoPE across three position channels (temporal,
height, width). When ``positions`` is 1-D this is bit-identical to
:class:`_PartialRoPECache` (so text-only decode steps and the
no-image prompt reuse the precomputed-table fast path); when
``positions`` is ``(3, N)`` the leading-channel frequencies are
interleaved per ``mrope_section`` — indices ``i % 3`` select T / H
/ W — exactly matching ``Qwen3_5TextRotaryEmbedding``.

Only instantiated when the vision tower is active, so the
production text-only path keeps the base class untouched.

Rotate the leading ``rotary_dim`` channels of Q and K at ``positions``,
passing the remaining channels through unchanged.

**MUTATES ``q`` and ``k`` IN PLACE** and returns them. Only the leading
``rotary_dim`` channels are written; the rest are already the values a
rebuild would have copied.

The rebuild is what this avoids. ``partial_rotary_factor = 0.25``
means 192 of 256 channels per head pass through BYTE-IDENTICAL, so
``torch.cat([q_rot_out, q_pass], -1)`` allocated a full Q and a full K
per layer per forward in order to copy three quarters of them onto
themselves. Writing the rotated quarter back into ``q`` deletes both
allocations and three quarters of the writes.

CALLER CONTRACT, and it is the whole safety argument -- ``q`` and ``k``
must be tensors the caller owns and does not read again afterwards:

  * ``attn.py:353`` rebinds (``q, k = rope_cache.apply(q, k, ...)``)
    and returns immediately. Its ``q`` is fresh, because ``use_qk_norm``
    defaults True (``attn.py:51``) and ``RMSNorm.forward`` returns
    ``x32.to(in_dtype) * self.weight`` -- a new tensor, not a view
    (``layers.py:86``). That freshness is what makes this a SAVING
    rather than a RELOCATION: were ``q`` a view, the cost would simply
    reappear at ``attn.py:500``'s ``.contiguous()``.
  * :meth:`_MRoPEPartialCache.apply` forwards 1-D positions here.

This is why the method lives on the PARTIAL caches and NOT on
:class:`~arbi_serve.models.rope.RoPECache`. That base class is reached
from ``mla_block.py:234`` with ``q_pe`` -- a SLICE of the caller's
``q``, whose sibling slice is re-read at ``:236``. Mutating there is a
real contract break, so the base class keeps the rebuild.

``(N, rotary_dim)`` cos/sin from ``(3, N)`` 3-D positions.

Builds per-channel ``freqs`` then folds the three channels into
one with the interleaved layout ``[T H W T H W … T T]`` over the
``rotary_dim // 2`` frequency indices, mirroring
``apply_interleaved_mrope``.

Rotate Q/K — 1-D positions reuse the base table; ``(3, N)``
positions take the interleaved M-RoPE path.

**MUTATES ``q`` and ``k`` IN PLACE** and returns them, on both
branches. See :meth:`_PartialRoPECache.apply` for the caller contract
and why the base :class:`~arbi_serve.models.rope.RoPECache` must keep
the rebuild.

Per-parameter shard spec.

The HF ``Qwen3_5ForConditionalGeneration`` checkpoint nests
text-tower weights under ``model.language_model.*``; this map
binds the model's ``model.*`` namespace (see
:class:`_Qwen3_5TextInner`) to that source prefix. Vision tower
(``model.visual.*``) and MTP draft layers (``mtp.*``) are
intentionally ignored — text-only by design.

Per-layer norms + token-mixer keys shared by the dense and MoE archs.

The attention/GDN interleave is IDENTICAL between
``Qwen3_5ForConditionalGeneration`` (dense FFN) and
``Qwen3_5MoeForConditionalGeneration`` (sparse MoE FFN) — only the
FFN keys differ, so both weight-map builders call this for the
mixer half of each layer.

Per-parameter shard spec for :class:`Qwen3_5MoeModel`.

Same ``model.language_model.*`` source namespace and mixer keys as
the dense arch (:func:`add_qwen3_5_mixer_keys`); the FFN half binds
the sparse-MoE keys instead of a dense gated MLP:

  - ``mlp.gate.weight`` — the router (dense bf16;
    ``modules_to_not_convert`` in the FP8 pack).
  - ``mlp.experts.{w13,w2}_weight`` (+ ``_scale`` on an fp8 pack) —
    the STACKED routed-expert tensors. Each checkpoint expert's
    ``{gate,up,down}_proj`` binds directly into its slot of the stack
    never materialized a second time to be concatenated.
    **EP-aware**: each rank binds only its contiguous slice of
    ``n_experts // ep_size`` experts, remapping the global checkpoint
    expert id (src) onto the local slot (dst). At ``ep_size == 1``
    the map is the identity over all experts.
  - ``mlp.shared_expert.{gate,up,down}_proj.weight`` — the always-on
    shared expert (replicated on every rank).
  - ``mlp.shared_expert_gate.weight`` — the sigmoid gate on the
    shared expert's output (dense bf16, replicated).

The bundled MTP draft head (``mtp.*``) binds wholly through
:meth:`~arbi_serve.spec_decode.mtp_head.Qwen3_5MtpHead.weight_map`,
whose FFN half runs the ``mlp_weight_map`` this arch injected at
construction — the same :func:`_add_qwen3_5_moe_ffn_keys` the main
layers use, so the head's decoder block (a full sparse-MoE FFN) is
EP-sliced per rank exactly like a main layer's. The vision tower
(``model.visual.*``) is intentionally ignored — text-only.

MoE FFN keys for one layer: router + local expert slice + shared.

On the fused dispatch the routed experts bind through the arch-agnostic
:func:`~arbi_serve.models._moe_weight_map.add_stacked_moe_expert_keys`
(EP slice, TP cut, gate/up seam, fp8 scale grid — all fused-layout
rules, not Qwen rules); otherwise each expert keeps its own three
projections, keyed by GLOBAL expert id so the quant swap resolves the
checkpoint path of every rank's slice. Qwen-shaped here: the
``mlp.gate`` router key and the gated shared expert + its sigmoid gate.

Bind this rank's expert slice as per-expert gated MLPs.

Module slot and checkpoint key both carry the GLOBAL expert id; this
rank simply declares fewer of them. At ``moe_tp_size == 1`` the expert
owns its whole intermediate dim, so no projection is TP-sharded — under
expert parallelism the rank holds WHOLE experts.

Shared weight-map builder helpers for per-arch model files.

Every model's ``weight_map()`` method constructs a
``dict[str, WeightShardSpec]`` mapping parameter names to safetensors
source keys + TP shard specs.  The embed/norm/lm_head init, standard
GQA attention binding, and gated-MLP binding are identical across
4-8 architectures — these helpers eliminate that duplication.

Per-arch files call the helpers for the common blocks and add
arch-specific entries (MLA projections, Mamba state, MoE routing,
GDN linear-attention, …) inline.

Add standard GQA attention projections (q/k/v/o) to *wm*.

``shard_dim=0`` for q/k/v (column-parallel), ``shard_dim=1``
for o (row-parallel).  When *with_qk_norm* is True, also adds
``q_norm.weight`` and ``k_norm.weight`` (Qwen3 style).
When *qkv_bias* is True, also adds ``q/k/v_proj.bias`` sharded
along dim 0 like their weights (Qwen2 / Step-Audio-2 style).
*o_proj_bias* adds the replicated ``o_proj.bias`` (gpt-oss
``attention_bias: true``). *attention_sinks* adds the per-head
``sinks`` vector, sharded along dim 0 with the query heads.

The module-side key always nests the projections under
``self_attn`` (that is where :class:`AttentionBlock` lives). Only
the *source* (checkpoint) key uses *src_attn_prefix* — most HF
checkpoints store the projections under ``self_attn`` too, but some
``LlamaForCausalLM`` derivatives (e.g. MiniCPM5) flatten them to
``model.layers.N.q_proj.weight`` with no ``self_attn`` segment.
Pass ``src_attn_prefix=""`` for those.

*skip_kv_proj* drops ``k_proj``, ``v_proj`` (and ``k_norm`` when
*with_qk_norm*) — Gemma 4 kv-shared layers reuse another layer's
K/V and carry no own projections.  *tie_v_to_k* drops only
``v_proj`` (V aliased from K) — Gemma 4 ``attention_k_eq_v``
full-attention layers.

Add gated-MLP projections (gate/up/down) to *wm*.

``shard_dim=0`` for gate/up (column-parallel), ``shard_dim=1`` for down
(row-parallel). *prefix* defaults to ``"mlp"``; Qwen3.5-MoE passes
``"mlp.shared_expert"``.

``merged=True`` binds the checkpoint's separate ``gate_proj`` + ``up_proj``
into ONE :class:`MergedColumnParallelLinear` ``gate_up_proj`` weight via
``shard_id`` (``"gate"`` / ``"up"``) — same numerics, one fused GEMM per
step. Keys carry a ``#gate`` / ``#up`` suffix to stay unique while
``dst_param`` points at the single fused parameter; the loader dispatches
one ``weight_loader(full_sub, shard_id=…)`` per shard, each tp-slicing
itself. MUST pair with ``GatedSiLUMLP(..., use_merged_gate_up=True)``.

Backend-agnostic attention block.

Owns the projection layers + per-head Q / K norms; the attention math
(compress + attend, paged FA, MLA, …) lives behind the per-step
:class:`AttnOp` instance the engine threads in.

The constructor takes a :class:`LayerSpec`; projections are
:class:`ColumnParallelLinear` / :class:`RowParallelLinear` (TP=1
short-circuits the all-reduce). A per-head sink bias materialises
when ``LayerSpec.attention_sinks is not None``.

Single attention layer — backend-agnostic.

The engine owns the per-layer :class:`AttnOp` for whatever backend
is active at this layer's :class:`StateKind` and threads it in via
``attn_op=``. ``attn_op.forward`` writes the per-token output into
a 3-D buffer; we then run ``o_proj`` on the flattened view.

Drop the cached fused QKV weight; next forward rebuilds it.

Call when the source projection weights' ``data_ptr``s change
(weight reload, sleep/resume re-bind). Kept off the hot path so
no per-instance ``data_ptr`` guard is emitted into the compiled
layer graph.

Lazily build (and cache) the ``[Q | gate | K | V]`` fused weight.

Returns ``None`` when fusion does not apply: the gate is off,
a projection has been swapped to a quantized linear (``.weight``
is not a plain dense tensor), or a projection is still on the
meta device. The fused tensor is ``torch.cat`` along dim 0 of
``[q_proj.weight | k_proj.weight | v_proj.weight]`` — bit-
equivalent (modulo cuBLAS algo selection on the wider GEMM) to
the three F.linears run separately and concatenated.

Run one forward step for this layer.

Returns ``(N_tokens, hidden_size)`` post-output-projection
residual contribution.

``out_buf`` is the pre-allocated attention output destination —
:meth:`alloc_arena_out_buf`'s result when the activation arena is
live, ``None`` otherwise (then :meth:`_attn_eager` allocates on
the heap exactly as before).

Allocate this layer's attention output buffer from ``arena``.

Runs EAGERLY, at the model's per-layer dispatch seam — never
inside a compiled block (:meth:`ActivationArena.alloc` refuses a
traced call). The result is threaded into the block as an
ordinary tensor argument, so the compiled trace sees a buffer,
not an allocator.

The dtype must match what ``q_proj`` will emit, because the
attention kernel asserts ``output.dtype == q.dtype``. A float
projection weight fixes it (bf16 in production); a quantized
projection restores its input's dtype on the way out, so
``hidden.dtype`` — read here off the real eager tensor, not a
traced one — is the answer for that case.

Project + per-head Q/K norm + RoPE → ``(q, k, v, gate_flat)``.

``gate_flat`` is ``None`` unless the output gate is on, in which
case it is the per-token gate (second half of ``q_proj``'s
output) flattened to ``(N, num_heads * head_dim)``; the gate is
applied to the attention output in :meth:`_post_attn`.

Gated QKV projection → ``(q, k, v, gate_flat)`` (pre-norm).

``q_proj`` emits ``(N, num_heads, head_dim * 2)``; the second
half per head is the gate. On GPU the q/k/v projections fold
into one matmul against the lazy ``[Q | gate | K | V]`` fused
weight (saves two launches per layer). On CPU — or when a
projection is quantized / LoRA is wired — the unfused path
runs so tests that monkey-patch the individual linears keep
working.

Run ``attn_op.forward`` into an output buffer → ``(N, H*D)``.

``out_buf`` is the caller's destination — the activation arena's
slot for this layer, allocated eagerly at the dispatch seam. It
is used only when it matches the shape and dtype the kernel
demands; anything else (no buffer, a stale width, a dtype the
projections did not in fact produce) falls back to a heap
allocation, so a mis-sized hand-off costs the arena's saving and
never correctness. The attention output is one of the two
largest per-step transients (the other is the MLP intermediate),
which is why it is the buffer worth routing.

The buffer MUST match ``q.dtype`` — the bf16 attention kernel writes
its result into ``output`` and asserts dtype-equality with q
(torchgen out-wrapper, "Expected out tensor to have dtype
c10::BFloat16, but got float instead"). We deliberately bind to
``q.dtype`` (the projection-output dtype) NOT ``hidden.dtype`` —
under ``ARBI_COMPILE_ON=1`` + Inductor + per-layer compile, the
piecewise-prefill capture path can trace ``hidden`` as fp32.
``q`` is freshly produced by ``q_proj`` (+ optional ``q_norm`` /
``rope_cache.apply``); its dtype is locked by the projection
weight (bf16 in production) regardless of upstream tracer drift.

Apply the output gate (when on) then ``o_proj`` → ``(N, hidden)``.

With the gate on, the attention output is multiplied element-
wise by ``sigmoid(gate)`` before ``o_proj`` (RowParallelLinear
all-reduces internally when tp_size > 1).

Audio towers — weight-bearing encoders for audio-input models.

The audio-side sibling of :mod:`arbi_serve.models.vision`: a tower
turns preprocessed audio features (log-mel chunks) into LM-embedding-
space tokens, registered on the host model via an ``mm_bindings``
entry (see :mod:`arbi_serve.multimodal.registry`). Weights load through
the host text model's ``weight_map`` (the tower contributes its
sub-map); the tower never owns a loader.

NemotronLabs-VoiceChat-11B TTS — Mixture-of-Gaussians CFG sampling head.

Ported from NVIDIA's ``MoGHead`` (reference:
``nemo/collections/speechlm2/modules/ear_tts_model.py``, lines ~495-668)
plus the RVQ quantization-snap logic (``depthsum_encoding_step``, lines
~472-492) and the per-frame CFG + iterative-unmasking sampling loop
that composes them (``RVQEARTTSModel.generate_step``, lines ~1546-1678).

Checkpoint prefix: ``tts_model.tts_model.mog_head.*`` — verified shapes:
``num_predictions=1024`` mixture components, ``low_rank=64``,
``out_size=latent_size=512`` (``low_mat`` ``[1024, 512, 64]`` confirms
``[num_predictions, out_size, low_rank]``), 3 gated-MLP blocks
(``mlp_stack.0/1/2``, hidden 1152<->4608) + 1 final RMSNorm
(``mlp_stack.3`` — a norm, not a 4th MLP layer), ``proj_logits``
``[1024,1152]``, ``proj_mus`` ``[65536,1152]`` (``= 1024*64, 1152``),
``proj_logs`` ``[1,1152]`` (one shared scalar log-std across every
component), ``proj_else`` ``[512,1152]``. Inference config (real
``config.json``): ``inference_guidance_scale=0.2``,
``inference_noise_scale=0.001``, ``inference_top_p_or_k=0.95``.

EOS classifier (reference ``RVQEARTTSModel.lm_head``, a 2-class linear)
is NOT ported: the real checkpoint's ``config.json`` already carries
``tts_config.disable_eos_prediction=True`` (and NeMo's
``checkpoint_utils/import_utils.py:_apply_tts_inference_config`` forces
it True at inference regardless), so the reference itself never
constructs ``lm_head`` for this checkpoint — EOS is driven by the text
channel instead. There is correspondingly no ``lm_head.*`` key under
``tts_model.tts_model.*`` in the real safetensors header. This module's
:func:`sample_frame_codes` therefore has no EOS-prediction branch and
no ``num_delay_speech_tokens`` handling (also forced to 0 at inference).

Depends on :mod:`arbi_serve.models.audio.nemotron_ear_tts` for the
shared ``RMSNorm``/``GemmaGatedMLP`` building blocks and
``depthsum_embedding`` (one-directional dependency: this module never
imports the backbone/composite classes).

Keep-rate -> masking-rate power schedule; its own inverse.

``rate``: desired fraction of quantizer depths already revealed,
in ``[0, 1]``. Returns the fraction still masked. With
``exponent=3`` (the checkpoint's ``tts_config.exponent``) this
concentrates most of the unmasking work in the LAST few
iterations — reveal-rate ramps slowly at first, then accelerates.

Nucleus filter over the LAST dim: mask all but the smallest prefix
of sorted-descending entries whose cumulative softmax mass reaches
``top_p`` to ``-inf``. Mirrors HF's ``TopPLogitsWarper`` (and the
equivalent per-row rule in
``arbi_serve.sampler.sampling_masks._vectorized_top_p_mask``,
specialized here to one scalar ``p`` shared by every row instead of
a per-request tensor).

Per-row weight-matrix gather + matvec: ``out[i] = w[y[i]] @ x[i]``.

``x``: ``[B, d_in]``. ``w``: ``[num_weights, d_out, d_in]``. ``y``:
``[B]`` long index into ``w``'s first dim. Returns ``[B, d_out]``.

Pure-PyTorch port of the reference's ``batch_matmul_pytorch``
fallback (the reference also has a Triton fast path,
``batch_matmul_triton``, gated on CUDA + triton availability; not
ported here — this head runs once per unmasking iteration, not
per token, so the extra kernel-launch cost is not the bottleneck
this reference optimized for a full-vocab MoE-style gather).

Predicts + samples a mixture-of-Gaussians over a continuous target.

``num_predictions`` mixture components, each with a low-rank mean
(``low_mat[k] @ proj_mus_k`` when ``low_rank`` is set) and ONE
shared scalar log-std across every component and every output
dimension (``proj_logs`` projects to a single scalar, not
``num_predictions`` nor ``out_size`` values).

``mlp_stack``: ``num_layers`` :class:`MLPLayer` blocks followed by
one final :class:`RMSNorm` (checkpoint's ``mlp_stack.{num_layers}``
entry — index ``3`` for ``num_layers=3`` — is a bare norm, not
another :class:`MLPLayer`).

Greedy per-depth nearest-neighbor RVQ encoding of a residual vector.

For each of the ``k`` quantizer depths in ``[depth_start,
depth_start + k)``: find the nearest codebook entry to the current
residual ``r`` (squared-L2 nearest neighbor, computed via the
standard ``||e||^2 - 2*r.e`` expansion — the constant ``||r||^2``
term is dropped since it doesn't affect the argmin), subtract that
entry from ``r`` (residual-quantization), and write the chosen
index into ``code[..., i]``. ``code`` is mutated AND returned.

``embs_sq``: optional precomputed ``[num_quantizers, codebook_size]``
per-entry squared norms (``embs[:, :codebook_size].pow(2).sum(-1)``).
They are a pure function of a frozen codebook, so a caller running
this per frame should hoist them out; computed here per depth
otherwise. Bit-identical either way.

``embs``: ``[num_quantizers, codebook_size + 1, latent_size]`` —
the codec's per-depth table, ALREADY padded with a sentinel row at
``codebook_size`` (see this module's docstring / the interface note
in ``nemotron_ear_tts.depthsum_embedding``). That sentinel row is
EXCLUDED from the nearest-neighbor search here (sliced off) so a
partially-revealed code can never accidentally snap TO the
not-yet-revealed sentinel — mirrors the reference, which searches
over the UNPADDED (real checkpoint) table directly.

``@torch._dynamo.disable`` (matching the reference): the loop
bound ``k`` is a Python int synced off a per-step ``.item()`` call
by the caller and varies iteration to iteration (see
:func:`sample_frame_codes`), and the body mutates ``code`` in place
via an indexed write — both are exactly the shapes of code Dynamo
either graph-breaks on or recompiles once per distinct ``k``/
``depth_start`` value. Forcing eager here avoids a recompile storm
without losing anything: this is an O(k * codebook_size) argmin +
gather per call, not a hot per-token kernel, so there is no
compiled-fusion win being left on the table. arbi-serve's own
compile story (Dynamo, same as upstream) needs the identical guard
for the identical reason — kept as-is rather than re-litigated.

Per-outer-step count of quantizer depths newly revealed.

``num_iter`` evenly spaced keep-rates in ``[0, 1)`` through
:func:`get_masking_rate`'s power schedule give a monotonically
DEcreasing per-step masking count; this returns the discrete
difference between consecutive steps, so ``sum(...) ==
num_quantizers`` and every depth is revealed exactly once, in one
monotonically growing prefix. For the real checkpoint's
``num_iter=8, num_quantizers=31, exponent=3.0`` this is
``(0, 0, 0, 1, 1, 3, 4, 22)``.

Memoized on its three scalar arguments and computed on CPU as plain
Python ints: it is a fixed property of the sampler's configuration,
not of any frame's data, and :func:`sample_frame_codes` needs the
values on the host to drive its loop — computing it per frame on the
device costs a ``linspace``/``pow``/``ceil``/``pad`` chain plus one
host sync per outer step, every 80 ms, for a constant.

One frame's worth of RVQ codes via CFG + iterative unmasking.

Ports ``RVQEARTTSModel.generate_step`` (reference lines ~1546-1678)
MINUS the ``lm_head``/EOS branch, which this checkpoint's inference
config never constructs (see module docstring) — so unlike the
reference, the per-step scale lists here have length ``num_iter``
exactly (no reserved index-0 slot for the removed lm_head step).

``hidden_states``: ``[B, 1, H]``, or ``[2B, 1, H]`` == ``[cond;
uncond]`` concatenated on batch dim 0 when ``guidance_scale`` is
given (mirrors :meth:`nemotron_ear_tts.NemotronEarTTSBackbone.forward`'s
``guidance_enabled`` batch-doubling convention). ``embed_code``: the
backbone's ``nn.Linear(latent_size, hidden_size, bias=False)`` —
passed in rather than imported so this module never depends on the
composite backbone class (see module docstring).

Unmasking schedule: see :func:`unmasking_schedule`.

``rvq_embs_sq``: optional precomputed codebook squared norms, handed
straight to :func:`depthsum_encoding_step` — see there.

Returns ``[B, 1, num_quantizers]`` long codes in
``[0, codebook_size)``. With ``return_debug=True`` also returns the
per-outer-step count of REVEALED depths (length ``num_iter``,
non-decreasing, ending at ``num_quantizers``) for testing the
unmasking schedule directly.

CFG-guided single-sample draw from the predicted mixture.

``x``: ``[B, T, H]`` (or ``[2B, T, H]`` == ``[cond; uncond]``
concatenated on batch dim 0 when ``guidance_scale > 0``).
Returns ``(mu, logs)`` each ``[B, T, out_size]`` /
``[B, T, 1]`` — the SELECTED component's (guided) mean and the
shared log-std; caller adds the Gaussian noise
(``mu + exp(logs) * randn * noise_scale``, see
:func:`sample_frame_codes`).

Training-mode forward: full mixture parameters, no sampling.

Returns ``(logits, mus, mu_res, logs)`` — ``logits``
``[B,T,num_predictions]``, ``mus`` ``[B,T,num_predictions,d]``
(``d = low_rank or out_size``), ``mu_res`` ``[B,T,d]``, ``logs``
``[B,T,1]``. Not exercised by the inference-only sampling path
(:func:`sample_frame_codes`); ported for completeness/fidelity
with the reference class.

Squared distance from each mixture mean to ``mu``, in the
(possibly low-rank-expanded) output space. ``mus``:
``[B,T,n,d]``, ``mu``: ``[B,T,d]``. Returns ``[B,T,n]``. Not
exercised by the inference-only sampling path; ported for
completeness with the reference class (used by its training
loss).

NemotronLabs-VoiceChat-11B — TTS backbone half (generation side).

The real checkpoint (``NVIDIA-NemotronLabs-VoiceChat-11B``) is a
composite speech-to-speech model: a perception tower + STT backbone
(``stt_model.*`` — see ``arbi_serve/models/nemotron_voicechat.py``,
owned by a different workstream), an RNNT decoder, and this TTS stack
(``tts_model.tts_model.*``). This file implements the TTS stack's
Gemma3-text-shaped transformer backbone plus its embedding-fusion
modules (character-aware subword encoder, gated audio/text fusion, BOS
marker). The Mixture-of-Gaussians sampling head lives in
``mog_head.py`` (imported here only to assemble the composite module).
The audio codec (``arbi_serve/audio/nemotron_audio_codec.py``, RVQ
embedding table + decode) is a third workstream's file — this module
never imports it; every function that needs the codec's per-quantizer
embedding table (``rvq_embs``) takes it as an explicit argument.

NOT wired into
``LayerStackModelMixin`` / the paged-KV ``ScheduledBatch`` engine path
those take discrete ``input_ids`` + a fixed per-token forward, while
this backbone consumes CONTINUOUS embeddings and runs one call per audio
FRAME inside an iterative CFG+unmasking sampling loop (see
``mog_head.sample_frame_codes``). A later integration pass composes
this with the STT backbone and the audio codec and decides how frame-
by-frame decoding maps onto the engine's batch scheduler. Until then
this is plain ``nn.Module`` code with a simple list-based KV cache.

=== Architecture (verified against the real checkpoint's safetensors
header, ``tts_model.tts_model.*``, 418 tensors, and the NeMo Hydra
``config.json`` at
``model.speech_generation.model.tts_config`` — see per-class docstrings
for exact key/shape citations) ===

``backbone`` is architecturally a literal HF ``Gemma3TextModel``
(``AutoConfig.for_model("gemma3_text", ...)`` in the reference
``RVQEARTTSModel.__init__``), confirmed by the checkpoint's QK-norm +
4-norm sandwich tensor layout: hidden_size=1152, 28 layers, 16 heads ==
16 kv_heads (MHA, not GQA), head_dim=72, intermediate_size=4608.
``embed_tokens`` does not exist in the checkpoint — the reference code
deletes it (``find_and_delete_module``) because this backbone is fed
``inputs_embeds`` only, never token ids, so Gemma3's
``embed_scale = hidden_size**0.5`` word-embedding scaling never
applies here.

Sliding-window pattern: not stated anywhere in this checkpoint's
config (only ``sliding_window=7500`` is given), so it is the
``transformers`` library's ``Gemma3TextConfig`` DEFAULT alternation —
confirmed by reading the installed ``transformers`` package source
(``configuration_gemma3.py``): ``sliding_window_pattern=6``,
``layer_types[i] = "full_attention" if (i+1) % 6 == 0 else
"sliding_attention"``. For 28 layers (0-indexed) that makes layers
{5, 11, 17, 23} full-attention and the other 24 sliding-attention
(window=7500, causal). Gemma3 also uses TWO rope thetas — global
1e6 for full-attention layers, local 1e4 for sliding-attention layers
(``Gemma3TextConfig.default_theta``) — and attention scaling is
``query_pre_attn_scalar ** -0.5`` (default 256, NOT ``head_dim**-0.5``)
per the real ``Gemma2Attention``/``Gemma3Attention`` source both
inherit from.

Gemma-family RMSNorm: ``y = rmsnorm(x) * (1 + weight)``, fp32-internal.

Weight is zero-init (checkpoint stores the zero-centred offset, not
the scale directly) — matches ``Gemma3RMSNorm``/``Gemma2RMSNorm``
in the reference ``transformers`` source and the NeMo reference's
own ``RMSNorm`` class in ``ear_tts_model.py`` bit-for-bit (both
compute the norm in fp32 THEN multiply by ``(1 + weight)`` before
casting back down — "Llama does x.to(fp16) * w whilst Gemma is
(x * w).to(fp16)").

Compute ``(cos, sin)`` tables for ``position_ids`` at the given theta.

``position_ids``: ``[B, T]`` long. Returns each of shape
``[B, T, head_dim]`` (the half-frequency table concatenated with
itself, matching HF's ``emb = cat([freqs, freqs], dim=-1)``
convention that :func:`arbi_serve.models.rope.rotate_half` pairs
with).

Boolean ``[q_len, kv_len]`` allow-mask: causal, optionally windowed.

``True`` = attend. ``kv_len >= q_len`` when a KV cache prefix is
present; the causal diagonal is anchored so the LAST ``q_len`` KV
positions align with the query positions (standard incremental-
decode convention: ``kv_idx <= q_idx + (kv_len - q_len)``).

Physical sliding-window KV length after trimming this step's cache.

A sliding-attention layer's cache never needs to hold more than
``sliding_window`` positions: :func:`build_causal_mask` already masks
any key more than ``sliding_window`` steps behind its query to
``-inf``, so a key past that horizon contributes nothing and can be
dropped from the cache with no effect on the attention output.
``past_len`` is therefore clamped to the window before adding this
step's ``t`` new positions. Returns the untrimmed total when ``t``
alone exceeds the window (a single step cannot attend to fewer than
its own query tokens). :meth:`Gemma3SelfAttention.forward` trims its
own cache to this exact length — the two must always agree, since a
caller may share one precomputed mask (sized off this function) across
every sliding-window layer's SDPA call.

Exact worst-case byte size of :class:`Gemma3SelfAttention`'s own
``past_k``/``past_v`` KV cache after ``num_ticks`` sequential
single-token decode steps (one per 80 ms TTS frame tick).

Closed-form, not measured: unlike a per-tick activation transient
(:func:`~arbi_serve.runtime.profile_peak.profile_tts_codec_peak`,
which needs a real forward pass because kernel/allocator behavior
isn't practical to hand-derive), this cache's size after N sequential
steps is EXACTLY determined by GPU call needed.

Of ``cfg.num_hidden_layers`` layers, the sliding-attention ones (all
but the 4 full-attention layers — see
:meth:`Gemma3TextBackboneConfig.is_full_attention`) hold at most
``cfg.sliding_window`` tokens once trimmed; the full-attention layers
hold every one of ``num_ticks`` tokens, This is why ``num_ticks``
must be a REAL bound (the duplex connection-duration hard cap,
``BatchConfig.duplex_max_session_s`` / ``DuplexTickPump.max_ticks``) for this to
be a meaningful reserve: a session with no duration cap has no finite
worst case to reserve for, full-attention layers included — see
:func:`arbi_serve.engine.inprocess_capture.serving_floor_for_grow`'s
own docstring for how the caller resolves that bound.

``batch``: the backbone's own batch dim at TTS-driving connections when CFG
(``guidance_scale > 0``) doubles the batch, which is this
checkpoint's real default (``generate_tts_frame``'s own
``guidance_scale=0.2``) and therefore the worst case a caller should
reserve for, matching
:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.profile_worst_case_tts_frame`'s
own choice to always exercise CFG.

Physical byte capacity of the bounded TTS paged-KV slab.

Each connection owns a conditional row and a CFG-unconditional row.
Rows include the speaker-conditioning prefix and are rounded to whole
pages. Page zero remains reserved by the shared page-table convention.

MHA self-attention with per-head QK-RMSNorm and RoPE.

Checkpoint keys (per layer): ``self_attn.{q,k,v,o}_proj.weight``
(all ``[1152, 1152]``, bias-free) and ``self_attn.{q,k}_norm.weight``
(``[72]``). ``num_key_value_heads == num_attention_heads`` here (16
== 16) so this is plain MHA — no ``repeat_kv`` needed.

The 28-layer Gemma3-text stack. Checkpoint prefix: ``backbone.*``.

Consumes ``inputs_embeds`` only (no ``embed_tokens`` — see module
docstring). Two RoPE thetas alternate by layer per
:meth:`Gemma3TextBackboneConfig.is_full_attention`. Runs causal
attention throughout (``use_bidirectional_attention`` is not set in
the checkpoint config, so it defaults to ``False``).

One layer of the char-encoder's tiny bidirectional t5gemma tower.

Checkpoint keys (``embed_subword.backbone.encoder.layers.0.*``):
``pre_self_attn_layernorm`` / ``post_self_attn_layernorm`` /
``pre_feedforward_layernorm`` / ``post_feedforward_layernorm`` +
``self_attn.{q,k,v,o}_proj`` + ``mlp.*``. No QK-norm (checkpoint
carries no ``q_norm``/``k_norm`` under this prefix) and full
(bidirectional, non-causal) attention — this is a 1-layer text
ENCODER, not a decoder.

``embed_subword.backbone.encoder.*`` — 1-layer bidirectional tower.

Two details verified against the REAL installed ``transformers``
``T5GemmaEncoder``/``T5GemmaRotaryEmbedding`` source (not just the
checkpoint's tensor names, which don't reveal either of these):

* ``T5GemmaEncoder.forward`` scales ``inputs_embeds`` by
  ``hidden_size ** 0.5`` UNCONDITIONALLY, every call — unlike
  ``Gemma3TextScaledWordEmbedding`` (the main 28-layer backbone's own
  embed-scale), where the scaling lives INSIDE the ``embed_tokens``
  module and is only applied when ``embed_tokens(input_ids)`` actually
  runs, so it's correctly skipped when ``inputs_embeds`` is supplied
  directly (this checkpoint's case, since ``embed_tokens`` is deleted
  — see this file's module docstring). T5Gemma's encoder does NOT
  gate the scaling on which input path was used, so omitting it here
  (as an earlier version of this port did, by analogy with the main
  backbone) was a real bug, not a matching simplification — confirmed
  by a live diff against the real checkpoint + real HF T5Gemma module:
  the 1-layer char-encoder's output was off by ~7% relative norm
  without this scale, and became bit-identical (max abs diff 0.0)
  with it restored.
* The encoder's own RoPE theta is 10000 (``T5GemmaRotaryEmbedding``'s
  ``rope_parameters`` for this checkpoint's ``cas_config`` — a
  single shared theta, not the main backbone's global/local SPLIT).
  10000 happens to equal :attr:`Gemma3TextBackboneConfig.rope_theta_local`,
  but that's this dataclass being reused for BOTH the main 28-layer
  backbone and this unrelated 1-layer char-encoder, not a real
  conceptual link — ``rope_theta_local`` is used here only because it
  already holds the right number for an unrelated reason. Using
  ``rope_theta_global`` (1e6, the main backbone's OWN default) was
  the bug: verified silently correct on the earlier single-character
  probe this port's tests used (RoPE at position 0 is the identity
  rotation regardless of theta) but wrong on any real multi-character
  subword — confirmed via a live 2-character diff against the real
  HF module (bit-identical, max abs diff 0.0, only once both this fix
  AND the embed-scale fix above were applied together).

Matches the checkpoint's ``embed_subword.backbone.encoder.*`` nesting.

The reference builds the char-encoder via
``AutoModelForTextEncoding.from_config(AutoConfig.for_model("t5gemma",
encoder=...))``, whose HF class nests the actual layer stack one
level under ``.encoder``. This thin wrapper reproduces that extra
attribute hop so state-dict keys line up exactly.

Continuation-token flag embedding.

Checkpoint keys (``embed_subword.subword_flag_emb.*``):
``cont_emb.weight`` ``[2, 1152]``, ``is_continuation``
``[subword_vocab_size + 1]`` (int64 buffer, index 0 = word-start,
1 = continuation; row ``subword_vocab_size`` is the custom pad
slot appended by the reference's ``build_vocabs``), ``pad_tensor``
(scalar int64 = ``subword_vocab_size``).

``is_continuation``/``pad_tensor`` are DATA derived from a live
subword tokenizer (which continuation tokens don't start a new
word) — this standalone module cannot rebuild that table without a
tokenizer, so it is a plain buffer the checkpoint loader fills;
:meth:`set_flags` lets a caller (the later integration pass, which
owns the tokenizer) populate it explicitly.

Independent BOS/EOS marker embedding added into the subword stream.

Checkpoint keys (``embed_subword.bos_eos_emb.*``): ``special_emb.weight``
``[3, 1152]`` (row 0 = regular/zeroed, 1 = BOS, 2 = EOS),
``special_flags`` (int64 buffer, per-vocab-id flag), ``pad_tensor``
(scalar int64). NOT the same mechanism as the top-level
``NemotronEarTTSBackbone.bos_emb`` Parameter, which is a single
audio-side marker vector added into the codec-embedding stream at
BOS-mask transitions — two distinct mechanisms (see reference
``ear_tts_model.py`` module docstring / class comments).

Subword embeddings built from mean-pooled character embeddings.

Checkpoint prefix ``embed_subword.*``: ``embed_tokens.weight``
``[257, 1152]`` (256-char vocab + 1 pad, ``padding_idx=256``),
``backbone.encoder.*`` (1-layer bidirectional t5gemma mini-tower,
see :class:`_CASEBackboneWrapper`), ``proj_embedding.weight``
``[1152, 1152]`` (bias-free), ``subword_flag_emb.*``,
``bos_eos_emb.*``.

``subword_id_to_char_ids`` (subword id -> tuple of char ids) is
tokenizer-derived data the reference builds once via
``build_vocabs(tokenizer, vocab_dir)``; this standalone module
takes it as an injectable mapping (set via :meth:`set_char_map`)
rather than depending on NeMo's ``AutoTokenizer`` — the later
integration pass supplies it from arbi-serve's own tokenizer.

Gated fusion of the audio (codec) and text (subword) embedding streams.

Checkpoint keys (``gated_fusion_audio_text.*``): ``audio_proj``
``[1152,1152]``+bias, ``text_proj`` ``[1152,1152]``+bias, ``gate``
``[1152]`` (fp32), ``residual_scale`` scalar (fp32), ``final_norm.weight``
``[1152]``.

``h = sigmoid(gate)*audio_proj(audio/num_codebooks) +
(1-sigmoid(gate))*text_proj(text)``, then ``h *= sigmoid(residual_scale)``,
then ``RMSNorm(h)``. Gate/residual_scale are computed in fp32 and
cast back to the projection dtype before multiplying, matching the
reference's ``fp32_precision()`` context (kept explicit here via
``.float()`` rather than an autocast-disabling context manager).

``num_codebooks``: the reference call site
(``GatedProjectedSumRMSNorm(hidden, hidden, hidden,
self.config.num_quantizers)``) passes ``num_quantizers`` into the
constructor's 4th POSITIONAL slot, which is actually ``final_norm:
bool`` — not the `num_codebooks` parameter (a wiring quirk in the
reference, not something to replicate). ``num_codebooks`` therefore
silently keeps its class default of 31, which happens to equal the
real ``num_quantizers`` anyway. This port uses an explicit,
correctly-named keyword argument instead of reproducing that
accidental-default footgun; the numeric behavior (divide by 31) is
identical.

Sum per-quantizer-depth embedding lookups into one continuous vector.

``rvq_embs``: ``[num_quantizers, codebook_size + 1, latent_size]`` —
the codec's per-quantizer-depth embedding table, ALREADY padded
with one extra "masked/unrevealed" row at index ``codebook_size``
(this is the interface contract with the audio-codec module; see
module docstring). ``code``: ``[B, T, num_quantizers]`` long, each
entry in ``[0, codebook_size]`` (``codebook_size`` == not-yet-
revealed sentinel).

NOTE — real checkpoint disagreement: the actual checkpoint stores
``tts_model.tts_model.rvq_embs`` UNPADDED, ``[31, 1024, 512]`` (no
extra sentinel row), and the reference's own
``RVQEARTTSModel.depthsum_embedding`` pads it with a zero row via
``F.pad(self.rvq_embs, [0,0,0,1])`` on every call. This port takes
the ALREADY-padded table per this task's stated codec interface
(``codebook_size + 1`` rows) instead, so no per-call pad is needed
here — see ``nemotron_ear_tts.py``/``mog_head.py`` module
docstrings for the full disagreement note.

Returns ``[B, T, latent_size]``.

Composite TTS backbone: embedding fusion + Gemma3 decoder stack.

Checkpoint prefix: ``tts_model.tts_model.*``. Wires
``audio_prompt_projection_W``, ``backbone``, ``bos_emb``,
``embed_code``, ``embed_subword``, ``gated_fusion_audio_text``,
``mog_head`` (from ``mog_head.py``), ``null_emb`` — every top-level
entry EXCEPT ``rvq_embs``, which the codec module owns (see module
docstring; every method that needs it takes it as an argument
instead of storing it as a buffer here).

``forward`` mirrors the reference's INFERENCE path only (frame-by-
frame decode with an externally supplied ``code`` tensor); the
reference's training-time masked/dropped-code construction
(``prepare_training_inputs``) and loss computation are out of
scope for this generation-only port.

Index of a layer whose KV cache is never window-trimmed.

A full-attention layer's cache always holds every past position
(see :func:`sliding_cache_len`), so its ``shape[2]`` is the
reliable source of the true absolute sequence length once
sliding-window layers' caches are bounded. Falls back to layer 0
for a degenerate config with no full-attention layer at all (e.g.
``num_hidden_layers < sliding_window_pattern``) — such a config has
no layer that can anchor absolute position past its own window,
which is a property of the config, not of this method.

Frozen pseudo-random projection ``audio_prompt_projection_W``
(``use_audio_prompt_frozen_projection=True`` in the checkpoint
config): QR-orthogonalized random matrix scaled by singular
values drawn uniformly in ``[0.4, 2.5]``, matching the
reference's construction verbatim. Skipped (empty placeholder)
under ``torch.device("meta")`` construction — ``torch.linalg.qr``
of a meta tensor has no real values to compute and this buffer
is checkpoint-loaded anyway.

Text-side conditioning: subword embedding, or ``null_emb`` for
the CFG unconditional branch. Mirrors the reference
``_prepare_conditioning`` minus the ``embed_context`` path (the
real checkpoint's ``context_hidden_size`` is ``None`` — no such
module exists to port).

One backbone step over ``code`` (``[B, RVQ codes).

When ``guidance_enabled``, the batch is doubled internally
(conditional first half, unconditional second half via
``null_emb``) so a single backbone call serves both CFG
branches — mirrors the reference's ``guidance_enabled`` path in
``RVQEARTTSModel.forward``. Returns ``(hidden_states,
present_key_values)``.

Nemotron VoiceChat perception tower — a Fast Conformer encoder.

Faithful port of the checkpoint's ``stt_model.perception.*`` tensors (the
NeMo ``AudioPerceptionModule`` = ``preprocessor`` + ``encoder``
(``nemo.collections.asr.modules.ConformerEncoder``) + ``modality_adapter``
+ ``proj``), configured exactly as this checkpoint's
``model.stt.model.perception`` config subtree:

  * subsampling: ``dw_striding``, factor 8, causal (``causal_downsampling
    =True``) — a plain Conv2d(1→256,k3,s2) then two depthwise-separable
    Conv2d stride-2 stages, each padded causally (left=k-1, right=s-1,
    applied to BOTH time and freq dims — this is NeMo's actual
    ``CausalConv2D``, not a time-only causal pad) — then a Linear
    (4352→1024) over the flattened (channel·freq) axis.
  * 24 Conformer layers, d_model=1024, 8 heads (head_dim=128),
    ff_expansion_factor=4 (macaron: two half-step-residual FF modules
    per layer), conv module kernel_size=9 with GLU + depthwise-separable
    conv, **causal** (``conv_context_size="causal"`` → pad
    ``[kernel-1, 0]``, zero look-ahead).
  * self-attention: Transformer-XL relative-position (``pos_bias_u``/
    ``pos_bias_v``, ``linear_pos``) — genuinely different math from
    everything else in this codebase (see :class:`_RelPositionMultiHeadAttention`).
  * ``att_context_style="chunked_limited"`` with ``att_context_size=[70,0]``
    — a causal, 70-frame-left-context attention mask (right context 0),
    applied identically to every batch row (the mask depends only on
    position indices, not on per-item length — see the causality note
    below).
  * ``use_bias=False`` and ``xscaling=False`` on the encoder — BOTH are
    deviations from the ``ConformerEncoder``/``ConformerLayer`` defaults
    (``True`` for both). The checkpoint has NO bias tensors on any of
    ``self_attn.linear_{q,k,v,out}``, ``feed_forward{1,2}.linear{1,2}``,
    or ``conv.{pointwise_conv1,depthwise_conv,pointwise_conv2}`` — only
    the subsampling module's plain ``torch.nn.Conv2d`` layers (unaffected
    by ``use_bias``) keep bias. ``xscaling=False`` means the pre-encoder
    features are NOT scaled by ``sqrt(d_model)`` before the rel-pos
    encoding is added — the plain ``ConformerEncoder`` default does scale.
  * ``conv_norm_type="layer_norm"``: the conv module's norm submodule is
    named ``batch_norm`` in NeMo (checkpoint-compat historical name) but
    is actually an ``nn.LayerNorm`` here — confirmed by the checkpoint
    carrying no ``running_mean``/``running_var``/``num_batches_tracked``
    buffers for it, only ``weight``/``bias``.
  * ``modality_adapter`` is ``IdentityConnector`` (no weights; the
    checkpoint has no ``modality_adapter.*`` tensors) — the encoder's
    output feeds ``proj`` directly.
  * final ``proj``: Linear(1024 → 4480) — the composite model's hidden
    size. This IS part of the tower's output (``encode()`` returns
    already-projected embeddings), per ``docs/adding-a-model.md`` §7's
    ``MediaBinding.encode(features, device) -> (n_tokens, hidden)``
    contract.

Two entry points, both exercising the SAME weights and the same math:

  * :meth:`NemotronVoiceChatPerceptionTower.encode` — batch/offline. One
    forward call per request over the whole buffered utterance's mel
    spectrogram. The reference has no separate "chunk a long utterance"
    path for offline inference distinct from this — ``att_context_style
    ="chunked_limited"`` IS the chunking, expressed as a banded causal
    attention mask inside a single forward, not a manual split into
    multiple forward calls.
  * :meth:`NemotronVoiceChatPerceptionTower.encode_stream` — genuine
    conv state tick to tick in a :class:`PerceptionStreamState`, O(1)
    work per new frame. See ":ref:`streaming`" below.

.. _streaming:

=== Cache-aware streaming ===

This checkpoint's encoder is not merely "causal enough to chunk after
the fact" — it IS NVIDIA's cache-aware streaming Fast Conformer, reused
verbatim: the model card names its source as ``nvidia/nemotron-speech
-streaming-en-0.6b`` ("Cache-Aware FastConformer with 24 encoder
layers"), whose four documented streaming modes are exactly
``att_context_size`` ``[70,0]``/``[70,1]``/``[70,6]``/``[70,13]`` — and
this checkpoint pins ``[70, 0]``, the 80 ms / one-frame / ZERO-lookahead
mode. The three config knobs that make cache-aware streaming valid were
therefore all in force at TRAINING time, not retrofitted at inference:
``att_context_style="chunked_limited"``, ``causal_downsampling=True``,
``conv_context_size="causal"``.

Consequently ``lookahead_steps = att_context_size[1] = 0`` and
``cache_drop_size = 0`` (NeMo ``ConformerEncoder.setup_streaming_params``),
i.e. NOTHING about a frame's output depends on any later frame, at any
depth — the receptive field compounds only to the LEFT, and the caches
below carry that left context exactly rather than recomputing it.

Per step, with ``subsampling_factor=8`` (one encoder frame per 80 ms):

  * ``mel_carry`` — the trailing ``2*subsampling_factor`` mel frames.
    Encoder frame ``j`` has a bounded, exact 15-mel-frame subsampling
    receptive field ``[8j-14, 8j]``, so a 17-frame window ``[8j-16, 8j]``
    (NeMo's ``pre_encode_cache_size = subsampling_factor + 1``, plus the
    8 new frames) reproduces it; the first ``n_window - n_new``
    pre-encode outputs are discarded (NeMo's ``drop_extra_pre_encoded``).
  * ``cache_channel`` ``(n_layers, 1, att_left=70, d_model)`` — each
    layer's PRE-projection self-attention input for the last 70 frames
    (NeMo caches the query, not K/V: ``MultiHeadAttention.update_cache``).
  * ``cache_time`` ``(n_layers, 1, d_model, kernel_size-1=8)`` — each
    layer's causal depthwise-conv left context (NeMo ``CausalConv1D
    .update_cache``).

error against the KV cache of the backbone this feeds.

Warm-up is EXACT here, deliberately unlike the reference: NeMo's
``CacheAwareStreamingAudioBuffer`` left-pads the pre-encode cache with
literal zero mel frames (its ``zeros_pads``) before real history exists,
which is NOT the same as the offline forward's stage-wise conv zero
padding once the subsampling convs' biases are in play. Feeding the real
(shorter) mel prefix instead — ``mel[0:1]`` for frame 0, ``mel[0:9]`` for
frame 1, the steady-state 17-frame window from frame 2 on — gives the
convs the identical padding the offline path gives them. Measured
against this checkpoint's real weights and real audio, that makes
streaming agree with the already-reference-verified offline path to
float32 round-off (max |diff| 6e-7, per-frame cosine 0.999999999995);
see ``tests/test_nemotron_voicechat_perception_streaming.py``.

A streamed utterance yields one FEWER frame than the offline encode of
the same finite waveform, and that is correct, not a shortfall: offline's
last frame is the one ``center=True`` STFT padding + the subsampling
convs synthesize from the utterance's trailing zero pad — an artifact of
knowing where the audio ends, which a live stream by definition does not.

Causality invariant (why batch padding needs no per-item numeric
special-casing beyond length masking): both the subsampling conv stack
and the encoder attention are causal with a bounded, position-independent
lookahead (``stride-1`` per causal subsampling stage; zero for the
attention, since ``att_context_size[1] == 0``). A VALID query position
can therefore only ever depend on positions at or before it (plus the
tiny constant subsampling lookahead), so the padded tail of a shorter
item in a batch never contaminates another item's or its own valid
outputs — this module still threads exact per-layer masks (mirroring
NeMo's ``MaskedConvSequential`` + ``_create_masks``) rather than relying
on this invariant alone, since a bias-carrying layer (the subsampling
convs have bias=True) would otherwise leak nonzero values into the
padded region between layers.

Module attribute names mirror the checkpoint keys 1:1 (``encoder.pre_encode
.conv.{0,2,3,5,6}``, ``encoder.layers.N.*``, ``proj``) so ``weight_map()``
is a mechanical walk over ``named_parameters()`` — no remapping.

Shape fields lifted from ``model.stt.model.perception`` in config.json.

Every field below is read from the REAL checkpoint's config.json
(``NVIDIA-NemotronLabs-VoiceChat-11B``); the guards in
:meth:`from_hf_config` raise ``NotImplementedError`` on any knob this
file has not verified against that checkpoint, per the "fail loud on
what you did not wire" rule (``docs/adding-a-model.md`` §3).

Plain-Python twin of :func:`_calc_length`, used at ``__init__`` time.

``_ConvSubsampling.__init__`` needs a concrete int (``nn.Linear``'s
``in_features``) to size the post-subsampling frequency axis — it
cannot go through a tensor ``.item()`` call, because construction
commonly happens inside ``with torch.device("meta")`` (see
``docs/adding-a-model.md``'s meta-device build step), under which a
freshly-created tensor would itself be a meta tensor and ``.item()``
would raise.

``chunked_limited`` attention mask for one cached streaming step.

The key axis is ``[cache (left_ctx) | new (n_new)]``, so new row ``i``
sits at key index ``left_ctx + i``; the queries are the new rows only
(NeMo ``ConformerEncoder.forward_internal``'s ``att_mask[:, cache_len:]``
slice). ``cache_len`` is how much of the cache is real history —
anything older is still-zeroed padding and must be masked out (NeMo's
``offset = -cache_last_channel_len + cache_len``).

Returns ``(1, n_new, left_ctx + n_new)`` bool, True where attention
must NOT look.

``chunked_limited`` causal attention mask + padding mask.

Mirrors ``ConformerEncoder._create_masks`` restricted to the
non-cached, non-local-attention, ``att_context_style="chunked_limited"``
branch this checkpoint uses.

Returns:
    pad_mask: ``(B, T)`` bool, True at INVALID (padded) positions —
        fed to the conv module's ``masked_fill``.
    att_mask: ``(B, T, T)`` bool, True at positions the attention
        must NOT attend to — fed to ``masked_fill(-INF_VAL)``.

NeMo's ``CausalConv2D``: pad ``(k-1, s-1)`` on BOTH time and freq
dims (via manual ``F.pad``), then a plain ``nn.Conv2d`` with padding=0.

Bias defaults True — the subsampling module is NOT covered by the
encoder's ``use_bias`` knob (confirmed: every ``pre_encode.conv.*.bias``
tensor is present in the checkpoint).

``dw_striding``, causal, factor-8 subsampling — the checkpoint's
``encoder.pre_encode``.

``self.conv`` is an ``nn.Sequential`` of 8 entries so its parameter
names land at indices 0/2/3/5/6 exactly like the checkpoint
(activations at 1/4/7 carry no parameters) — mirrors NeMo's own
``ConvSubsampling`` layer list construction 1:1.

Transformer-XL sinusoidal relative positional encoding.

``xscaling=False`` on this checkpoint's encoder: unlike NeMo's default
``RelPositionalEncoding`` (which scales ``x`` by ``sqrt(d_model)``
before returning it), this module returns ``x`` UNCHANGED.

The ``pe`` table is a lazily-built plain attribute, not a registered
buffer — it is a pure function of position/``d_model`` (not a
checkpoint tensor; the checkpoint carries no ``pos_enc.*`` key), so
keeping it out of the module's state avoids the meta→ones
buffer-materialization loop clobbering it (``docs/adding-a-model.md``
§7's rotary ``inv_freq`` caveat, same failure mode).

Transformer-XL relative-position self-attention (rel_pos).

Distinct from every other attention in this codebase: two learned
per-head bias vectors (``pos_bias_u``/``pos_bias_v``, shape
``(n_heads, head_dim)``) combine with a projected relative position
embedding (``linear_pos``) via the "matrix bd" relative-shift trick
(Dai et al., Transformer-XL, https://arxiv.org/abs/1901.02860 §3.3).
All four QKVO linears AND ``linear_pos`` are bias-free on this
checkpoint (``use_bias=False``; ``linear_pos`` has no bias in NeMo
regardless of ``use_bias``).

GLU depthwise-separable conv module, causal (zero look-ahead).

``pointwise_conv1`` (d_model -> 2*d_model) -> GLU -> causal
``depthwise_conv`` (kernel_size, groups=d_model) -> ``batch_norm``
(an ``nn.LayerNorm`` despite the name — see module docstring) ->
Swish -> ``pointwise_conv2``. All three convs are bias-free on this
checkpoint.

One Conformer block: macaron-FF / self-attn / conv / macaron-FF.

``x + 0.5*FF1(LN(x))``, then ``+ SelfAttn(LN(.))``, then
``+ Conv(LN(.))``, then ``+ 0.5*FF2(LN(.))``, then a final ``LN``
(``norm_out``) — the standard Conformer block structure, attribute
names mirroring the checkpoint's ``encoder.layers.N.*`` keys exactly.

One instance per duplex connection, threaded through
:meth:`NemotronVoiceChatPerceptionTower.encode_stream`. Mutated in
place each step — never shared between connections. See the module
docstring's ":ref:`streaming`" section for what each cache holds and
why the sizes are what they are.

Args:
    mel_carry: ``(1, <= 2*subsampling_factor, feat_in)`` — trailing mel
        frames retained so the next step's subsampling window can cover
        encoder frame ``j``'s full ``[8j-16, 8j]`` span.
    cache_channel: ``(n_layers, 1, att_left, d_model)`` — per-layer
        self-attention left context.
    cache_time: ``(n_layers, 1, d_model, conv_kernel-1)`` — per-layer
        causal depthwise-conv left context.
    cache_len: how many of ``cache_channel``'s ``att_left`` slots are
        real history rather than still-zeroed warm-up padding.
    mel_seen: mel frames consumed so far (the streaming schedule's
        clock — encoder frame ``j`` is due once mel frame ``8j`` lands).
    frames_out: encoder frames emitted so far.

Fast Conformer perception encoder + final projection.

``encode()`` implements the ``MediaBinding.encode(features, device)
-> (n_tokens, hidden)`` contract from ``docs/adding-a-model.md`` §7:
the output is already projected into the composite model's hidden
space (``proj``, 1024 -> 4480), ready to scatter onto placeholder
tokens. ``mm_bindings`` wiring here; a later
integration pass registers this tower on the composite model.

``x``: ``(B, T, feat_in)``, every frame valid. Returns ``(B, T', d_model)``.

The length masking :meth:`forward` threads is what keeps a padded
batch tail from leaking between items; a streaming window has no
padding at all, so it is skipped rather than made a no-op.

``cache_len`` extends the position span by the cached left context
a streaming step attends over (NeMo ``RelPositionalEncoding.forward``'s
own ``cache_len`` argument): the relative distances a query can see
run over ``x.size(1) + cache_len`` positions, not ``x.size(1)``.

``cache`` (streaming): ``(B, left_ctx, d_model)`` of this layer's
own PREVIOUS inputs — NeMo caches the pre-projection query, not K/V
(``MultiHeadAttention.update_cache``), so K/V are re-projected from
``[cache | x]`` each step. Returns ``(out, next_cache)`` when given
a cache, a bare ``out`` otherwise (mirrors NeMo's own signature).

``cache_drop_size`` is 0 for this checkpoint (``att_context_style
="chunked_limited"`` with zero right context), so the next cache is
simply the trailing ``left_ctx`` rows of ``[cache | x]``.

``cache`` (streaming): ``(B, d_model, kernel_size-1)`` of the
post-GLU activations immediately left of ``x`` — exactly what the
offline path's causal ``F.pad`` supplies as zeros at ``t=0`` (NeMo
``CausalConv1D.update_cache``). Returns ``(out, next_cache)`` when
given a cache, a bare ``out`` otherwise.

One cached streaming step through this block.

Identical module order and residual structure to :meth:`forward` —
only the self-attention and conv modules gain their caches, and
there is no ``pad_mask`` (a streaming window is entirely valid).

Consume ``new_mel`` ``(1, n, feat_in)``; return ``(1, n_new, d_model)``.

``n_new`` is however many encoder frames the newly-arrived mel
completes — commonly 1 per 80 ms of audio, 0 if not enough has
arrived yet, more when catching up on a burst. All of them are
emitted in ONE forward (the attention mask and both caches handle a
multi-row query block), so a backlog costs one pass, not one pass
per frame.

Run the tower over a request's packed mel-spectrogram batch.

Returns ``(sum_item_tokens, output_dim)`` — the valid projected
outputs of every audio item concatenated in prompt order. Raises
when an item's produced length disagrees with the prompt's
expanded placeholder count for it (mirrors the audio-tower
convention in ``models/audio/step_audio2.py``).

Incrementally encode newly-arrived mel frames for one live stream.

The streaming counterpart of :meth:`encode` — but driven by a
live, still-arriving audio stream rather than a whole buffered
utterance, carrying per-layer attention/conv state in ``state`` so
each new frame costs O(1) rather than a re-encode. Returns BOTH
stages of the tower's output, mirroring the reference's own
``perception()`` return shape (``(encoded, lens, asr_emb)``): the
projected embedding the LLM fusion consumes, and the raw encoder
output the bundled RNNT endpoint-detector head
(:class:`~arbi_serve.models.audio.nemotron_voicechat_rnnt
.NemotronVoiceChatRnntHead`) decodes — the reference's
``asr_emb``, computed inside this forward either way.

Args:
    mel_frames: ``(n, feat_in)`` or ``(1, n, feat_in)`` — the mel
        frames that have arrived since the last call. May be empty.
    state: this stream's :class:`PerceptionStreamState`, mutated
        in place.

Returns:
    ``(proj, enc)`` — ``proj``: ``(n_new, output_dim)``, the
    projected frames (identical rows to :meth:`encode`'s output);
    ``enc``: ``(n_new, d_model)``, the same frames BEFORE the
    final ``proj``. Often ``n_new == 0`` (not enough new audio
    yet); one frame per 80 ms of audio in steady state.

NemotronVoiceChat's bundled RNNT endpoint-detector head.

The VoiceChat-11B checkpoint ships, alongside the perception tower and
the LLM backbone, a small streaming RNNT ASR head over the perception
encoder's PRE-projection output: ``stt_model.rnnt_decoder.*`` (a
1025-row prediction-network embedding + a 2-layer LSTM) and
``stt_model.rnnt_joint.*`` (encoder/prediction projections + the joint
network), 15 tensors total, with its own 1024-token SentencePiece
vocabulary under ``rnnt_tokenizer/``. NVIDIA's production backend runs a
greedy decode step of this head on every 80 ms encoder frame and derives
its ENTIRE turn-taking evidence from the blank/non-blank outcome:
``rnnt_user_speaking`` (BOU), ``rnnt_silent_frames`` (EOU), rolling
speech density, and the noise reset all read this head's per-frame
verdict, never an acoustic energy/probability gate. That is the
"semantic VAD" property: a frame counts as user speech only when the
ASR head can actually transcribe content from it — breathing, mouth
noise, and non-speech background decode to blank.

This module is a faithful port of exactly the per-frame greedy step the
reference's EOU path runs (``_rnnt_eou_decode_frame`` in the extracted
NIM backend, specialized to one stream), on the same weights:

  * **Prime**: the prediction network is primed by running the LSTM over
    TWO all-zero input rows (NeMo ``RNNTDecoder.predict(y=None,
    add_sos=True)``: a zero row emulating the pad-token embedding plus a
    prepended zero SOS row) and keeping the last output.
  * **Per frame**: one joint step ``joint_net(enc(f) + pred(g))`` over
    the frame's 1024-d encoder output and the current prediction-network
    output; ``argmax`` picks the token; ``blank`` (id 1024,
    ``num_classes - 1``) means "no transcribable speech in this frame".
    The blank/non-blank verdict is taken from this FIRST prediction.
  * **Label drain**: while the argmax is non-blank (capped at
    ``max_symbols=10`` per frame, the reference's ``RNNT_MAX_SYMBOLS``),
    the emitted label advances the LSTM and the joint re-evaluates —
    standard greedy RNNT, keeping the prediction network aligned with
    what the head has recognized so far.

The reference's EOU path also carries punctuation-bias machinery
(``_punct_word_acc`` / ``_punct_bias_val``), but in that path it is
inert: its word accumulator is only ever appended to from an ``emitted``
list that the EOU drain loop never appends non-punctuation tokens to, so
the bias stays 0.0 and the injection branch never runs (it is live only
in the separate transcript-DISPLAY decoder, which this repository does
not run). This port therefore implements the plain greedy step — the
shipped EOU behavior — and none of the display-only biasing.

Per-tick cost: one ``(1024→640)+(640→640)`` projection pair, a ReLU, a
``(640→1025)`` linear, and — only on non-blank frames — one 2-layer
LSTM(640) step per emitted label. Negligible against the 24-layer
Conformer step that produces the frame it reads.

Shape fields lifted from ``_rnnt_merge_info`` in the checkpoint's
``config.json`` (the NeMo RNNT decoder/joint construction record).

Every field is read from the REAL checkpoint's config; the guards in
:meth:`from_hf_config` raise ``NotImplementedError`` on any knob this
file has not verified against that checkpoint, per the "fail loud on
what you did not wire" rule (``docs/adding-a-model.md`` §3).

``stt_model.rnnt_joint.*`` — NeMo ``RNNTJoint``'s split-projection
joint: ``joint_net(enc(f) + pred(g))`` with the additive expansion
specialized here to the streaming ``T=U=1`` case (one frame, one
prediction-network output), where it degenerates to a plain sum.

One live audio stream's RNNT decode state — per connection,
mutated in place by :meth:`NemotronVoiceChatRnntHead.step_frame`.

``pred_out`` is the prediction network's current output ``(1, 1, H)``
and ``pred_hidden`` the LSTM's ``(h, c)`` state; both ``None`` until
the first frame primes them.

NeMo ``RNNTDecoder.predict(y=None, state=None, add_sos=True)``:
run the LSTM over two all-zero input rows (the pad-token-embedding
stand-in plus the prepended zero SOS row) and keep the LAST output
row — exactly the reference's priming call, which slices
``pred_out[:, -1:, :]`` from the two-row result.

Decode ONE 80 ms encoder frame; return ``(is_blank, emitted)``.

``asr_frame``: ``(enc_hidden,)`` or ``(1, enc_hidden)`` — one
frame of the perception encoder's PRE-projection output (the
reference's ``asr_emb``), NOT the 4480-d projected embedding the
LLM fusion consumes.

``is_blank`` is the frame's turn-taking evidence, taken from the
FIRST joint prediction (the reference keeps exactly this one as
blank/non-blank evidence and uses the drain loop only to advance
the prediction network). ``emitted`` is the frame's greedy label
sequence — the actual recognized subword ids, empty on a blank
frame; the per-tick gate ignores it, diagnostic tooling decodes
it against ``rnnt_tokenizer/``.

Step-Audio-2 audio tower — Whisper-style encoder + Conv1d/MLP adapter.

Faithful re-implementation of the checkpoint's ``AudioEncoder`` +
``Adaptor`` (``encoder.*`` / ``adapter.*`` tensors):

  encoder: Conv1d(128→1280, k3 s1 p1)+GELU → Conv1d(1280→1280, k3 s2
  p1)+GELU → +frozen learned pos-emb (1500×1280) → 32 × pre-LN blocks
  (MHA 20 heads, q/k each scaled by d^−0.25, **no key bias**, fp32
  softmax; MLP 1280→5120→GELU→1280) → AvgPool1d(2,2) over time →
  LayerNorm.

  adapter: Conv1d(1280→1280, k3 s2 p1)+GELU → Linear(1280→2048) → ReLU
  → Linear(2048→llm_dim 3584).

Module attribute names mirror the checkpoint keys exactly, so the
tower's ``weight_map`` is a mechanical walk over
``named_parameters()``. The tower is replicated across TP ranks (never
sharded) and runs eagerly, once per request, off the KV-cached decode
path.

Length accounting (must agree with the prompt's placeholder expansion):
``mel_lens`` (mel frames − 2) → encoder ``(len+1)//2//2`` → adapter
``(len−1)//2+1``; the mismatch guard raises rather than mis-scattering.

Run the tower over a request's packed mel chunks.

Returns ``(sum_chunk_tokens, llm_dim)`` — the valid adapter
outputs of every chunk concatenated in prompt order, ready to
scatter onto the ``<audio_patch>`` positions. Raises when a
chunk's produced length disagrees with the prompt's expanded
placeholder count for it.

Replicated spec per tower parameter.

Module attribute paths mirror the checkpoint keys (``encoder.*``,
``adapter.*``; the MLP ``nn.Sequential`` yields the reference's
``mlp.0`` / ``mlp.2`` indices), so the map is a mechanical walk.

Bailing MoE V3 — Ling-3.0 (``bailing_hybrid``).

MLA on the last layer of each ``layer_group_size`` block, KDA elsewhere;
dense MLP below ``first_k_dense_replace``, sparse MoE above. Checkpoint
keys are ``model.word_embeddings``, ``attention.*`` and (MLA)
``attention.dense``, so the shared ``_weight_map_helpers`` do not apply.

KDA projections, conv kernels and gate parameters.

Q/K/V bind into the block's fused ``in_proj_qkv`` one ``shard_id`` at
a time; the three conv kernels stay separate parameters and are
concatenated after load.

Bind ``kv_b_proj`` into each MLA layer's attn op.

The op reconstructs V from the cached compressed latent and
cannot do it without the dense weight. KDA layers have no attn
op to bind — the ``isinstance`` check skips them.

Architecture-agnostic :class:`ModelBase` protocol + :class:`ModelDims`.

Every model implementation must:

  - expose ``dims: ModelDims`` (model-global facts: ``num_layers``,
    ``hidden_size``, ``vocab_size``, ``rope_theta``, ``rms_norm_eps``,
    ``intermediate_size``, ``max_position_embeddings``,
    ``tie_word_embeddings``);
  - expose ``layer_specs: list[LayerSpec]`` — the per-layer truth for
    head_dim, num_heads, num_kv_heads, sliding_window, sinks, and the
    MLA / Mamba / GDN / MoE configs;
  - expose ``weight_map() -> dict[str, WeightShardSpec]`` — the
    per-parameter shard plan the loader consumes;
  - implement ``from_safetensors(model_dir, *, device, dtype,
    parallel_cfg)``;
  - implement ``forward(input_ids, batch, cache, attn_ops) ->
    (num_seqs, vocab_size)`` returning per-seq last-token logits.

Cache layout and per-token byte cost belong to the active backend
(:class:`arbi_serve.backends.base.AttentionBackend`), never the model.
The engine threads ``attn_ops`` and the :class:`MultiStatePool` view
into ``forward`` so the model stays backend-agnostic.

:class:`ModelDims` carries only model-global facts; per-layer geometry
(``head_dim``, ``num_heads``, ``num_kv_heads``, ``sliding_window``)
lives on :class:`LayerSpec`.

Per-parameter shard spec the loader consumes.

Maps ``param_name`` (in the model's ``named_parameters()`` namespace)
to a :class:`WeightShardSpec` declaring the safetensors source key
and (optionally) the shard dim / shard id / expert id.

Run prefill or decode and return per-sequence last-token logits.

``input_ids`` is flat ``(N_tokens,)``. The per-sequence
boundaries live in ``batch.cu_seqlens_q``.

``cache`` is the :class:`MultiStatePool` — each layer reads its
per-state-kind view via ``cache.layer_view(i)``.

``attn_ops[i]`` is the per-layer :class:`AttnOp` for whichever
backend the engine has bound to this layer's state kind. The
engine rebuilds this list on every backend swap.

Returns ``(num_seqs, vocab_size)`` float32 logits — the last
token's logits per sequence (sampler input).

When ``return_hidden_state=True`` returns
``(logits, hidden_full)`` where ``hidden_full`` is the FULL
per-token post-final-norm activation just before ``lm_head``
of shape ``(N_tokens, hidden_size)`` in model dtype. The MTP
verify pass uses ``hidden_full`` to compute logits at every
token position (not just the per-seq last); the MTP draft
head consumes the per-seq last slice as its prior-hidden
input. Default ``False`` returns a single tensor; non-MTP
paths are bit-identical.

baidu/Unlimited-OCR text tower — ``UnlimitedOCRForCausalLM``.

MODEL vs ARCHITECTURE. The CHECKPOINT this serves is **baidu's**
(``config.json``: ``model_type: unlimited-ocr``,
``architectures: [UnlimitedOCRForCausalLM]``). It is BUILT ON the
**DeepSeek-OCR architecture** — its language tower declares
``DeepseekOCRForCausalLM`` on a ``DeepseekV2Config`` and the repo ships
``modeling_deepseekv2.py`` / ``deepencoder.py`` — which is why this file
is named ``deepseek_ocr.py`` and why it reuses ``_DeepseekV2MoE`` from
:mod:`arbi_serve.models.deepseek_v3`. Those structural names are correct
and intentional.

What is NOT interchangeable is *behaviour*: the weights are baidu's, so
which instructions the model was trained on and when it emits EOS are
properties of BAIDU's checkpoint. Ground every behavioural claim in
Unlimited-OCR's own artifacts (``modeling_unlimitedocr.py``'s ``infer()``,
``processor_config.json``, baidu's model card) — never in DeepSeek-OCR's
docs. Prompt modes are enumerated in
:data:`arbi_serve.chat_templates.UNLIMITED_OCR_PROMPT_MODES`.

The LM is a DeepSeek-V2 MoE decoder in **plain MHA form**
(``use_mla: false`` — q/k/v/o projections, no latent compression, no
q/k-norm, no bias) with the DeepSeek-MoE FFN split:

  * layer ``idx < first_k_dense_replace`` — dense :class:`GatedSiLUMLP`;
  * layer ``idx >= first_k_dense_replace`` — :class:`_DeepseekV2MoE`
    (router + routed experts + always-on shared experts). Routing is the
    single-group greedy top-k read off ``config.json`` — the reference
    ``MoEGate`` defaults apply: ``scoring_func="softmax"``,
    ``norm_topk_prob=False``, ``routed_scaling_factor=1.0``.

Topology per layer (pre-norm, two-norm — the reference
``DeepseekV2DecoderLayer``)::

  Embedding → [PreNorm → AttentionBlock + residual → PreNorm →
              (GatedSiLUMLP | _DeepseekV2MoE) + residual] × N →
  final RMSNorm → lm_head

Attention is **full causal** on every layer. The checkpoint's
``sliding_window(_size): 128`` is NOT an attention window: the
reference builds its 4-D causal mask WITHOUT a sliding window
(``modeling_deepseekv2.DeepseekV2Model.forward`` →
``_prepare_4d_causal_attention_mask`` with no ``sliding_window``
argument; ``infer()`` even nulls ``config.sliding_window`` before
``generate``). The 128 sizes the reference's decode-time KV RING
(``config._ring_window`` in ``SlidingWindowLlamaAttention``): decode
token ``t`` writes its KV into ring slot ``prefill + (t mod 128)``,
overwriting the oldest retained decode KV, while queries attend the
FULL retained cache (whole prompt + ring). Serving with a true SWA-128
mask blinds the decoder to the image (the prompt's text tokens sit
> 128 positions after most vision tokens) — GPU-observed as
structurally-OCR-shaped hallucination.

The ring is FUNCTIONAL, not a memory bound: it is the model's designed
anti-repetition mechanism. Serving full-attention decode instead
repetition-loops on dense/numeric pages (GPU-verified: the HF
reference with the ring disabled reproduces the exact loops; with the
ring it terminates at EOS — and the verbatim n-gram blocker CANNOT
substitute, because the loops vary numerically). ``decode_kv_ring``
below is therefore consumed at engine build
(``engine/active.py`` → ``FlatPageTable.decode_kv_ring``): decode-row
KV slots wrap at ``prompt + ring`` and the attention KV length caps
there, while RoPE positions keep climbing — exact reference decode
semantics on the paged cache.

The LM-tower geometry lives under ``config.language_config`` (the root
mirrors it).

Vision weights in the same safetensors (``model.sam_model.*``,
``model.vision_model.*``, ``model.projector.*``) are not bound by this
class — the vision tower loads separately via the multimodal seam.

Unlimited-OCR LM-tower shape config.

Lifted from ``config.language_config`` once at boot. ``moe_routing``
carries the gating knobs :class:`_DeepseekV2MoE` consumes verbatim
(scoring_func / norm_topk_prob / routed_scaling_factor + expert
counts); ``None`` means every layer is dense.

One Unlimited-OCR transformer block: pre-norm + MHA + pre-norm + FFN.

The attention block is the plain Llama shape (``use_qk_norm=False``,
bias-free); the FFN is dense :class:`GatedSiLUMLP` when
``layer_spec.moe is None`` and :class:`_DeepseekV2MoE` otherwise.
Implements the uniform :class:`LayerStack` block signature directly.

One full-attention :class:`LayerSpec` per layer.

Every layer is :class:`LayerKind.ATTENTION` with NO sliding
window (the checkpoint's 128 is the reference decode-KV ring,
not a mask — see the module docstring); the dense-vs-MoE FFN
split rides on :attr:`LayerSpec.moe` (``None`` below
``first_k_dense_replace``).

Map each parameter to its safetensors source key + TP shard spec.

Standard DeepSeek HF naming: dense attention under
``self_attn.{q,k,v,o}_proj`` (no q/k-norm), FFN either
``mlp.{gate,up,down}_proj`` (dense layers) or the DeepSeek MoE
subtree (router gate + per-expert + shared experts).

DeepSeek V3 / V2 / V2-Lite — MLA attention.

Topology mirrors official ``DeepseekV3ForCausalLM``:
  Embedding → [PreNorm → MLAAttentionBlock + residual → PreNorm →
              GatedSiLUMLP + residual] × N → final RMSNorm → lm_head.

Attention is :class:`MLAAttentionBlock` (multi-head latent attention
with optional q-LoRA / kv-LoRA per ``MLAConfig``). Layers below
``first_k_dense_replace`` carry a dense gated MLP; the rest carry the
sparse MoE FFN (:class:`_DeepseekV2MoE` — fused stacked experts on a
dense pack, per-expert quant modules on a weight-quantized pack). AWQ
checkpoints (V2-Lite-Chat-AWQ) are supported; EXL3 is also live in
:mod:`arbi_serve.weight_quant`.

``config.json`` ``rope_scaling`` (YaRN on V2-Lite / V3 / R1) drives both
the q_pe / k_pe RoPE table (:class:`YarnRopeScaling`) and the attention
softmax scale (``mscale_all_dim`` squared, via
``LayerSpec.attention_scale``).

V2 / V2-Lite parity: same class topology, differences live on
:class:`MLAConfig` (``q_lora_rank=None`` for V2-Lite, smaller head
counts). Both route through this class unchanged.

Per-layer dispatch routes through :class:`LayerStack` via the
per-arch dispatcher (:func:`_make_deepseek_v3_decoder_layer`). The
inner decoder layers implement the uniform LayerStack block
signature directly, so MLA / dense layer combinations share the
same forward body.

Per-arch construction bundle for :class:`DeepseekV3Model`.

Collects the ctor args that don't fit on :class:`ModelDims` (which
is shared across all archs) but are model-arch-specific knobs
derived from HF ``config.json``. Following the mlc-llm pattern, the
arch's :meth:`from_safetensors` builds one of these and the
constructor consumes it (rather than threading per-layer MLA shape
metadata through the call site).

The MLA shape (``mla_cfg``) is shared across every MLA layer in
all DeepSeek V3 / V2 checkpoints — the config is not
per-layer — so a single :class:`MLAConfig` instance is enough.
:class:`LayerSpec.mla` (per-layer) is the source of truth that
:class:`~arbi_serve.models.mla_block.MLAAttentionBlock` reads
directly.

Args:
    mla_cfg: per-layer MLA shape — same instance shared across
        every MLA layer.
    head_dim: projected per-head Q / K width (``qk_nope_head_dim
        + qk_rope_head_dim``); not the v_head_dim used for V.
    num_attention_heads: per-layer Q-head count.
    num_key_value_heads: per-layer K/V-head count (DeepSeek MLA
        sets this equal to ``num_attention_heads`` because the
        per-head V reconstruction means post-W_UV K and V have
        ``num_attention_heads`` heads).
    yarn: ``rope_scaling`` YaRN schedule for the ``qk_rope_head_dim``
        tail, or ``None`` on a checkpoint that ships none.

Shared body for the three DeepSeek V3 decoder-layer flavors.

All three (MLA + dense MLP, MLA + sparse MoE, dense attention + MLP)
run the identical pre-norm residual block — they differ only in
which token mixer / FFN ``__init__`` wires and the fail-loud
missing-metadata error label. The concrete subclasses set
``_layer_label`` / ``_meta_label`` and wire their submodules; this
base owns the forward + the (unfused) mixer / FFN hooks.

Implements the uniform :class:`LayerStack` block signature directly.

One DeepSeek MLA decoder block: pre-norm + MLA + pre-norm + FFN.

Ships a dense :class:`GatedSiLUMLP` FFN. Layers indexed
``>= first_k_dense_replace`` route through
:class:`_DeepseekV3MoEDecoderLayer` instead, which carries the
sparse MoE FFN.

Route tokens through per-expert MODULES — the quantized-pack path.

Exists only because a per-linear weight-quant backend (AWQ Marlin)
cannot be represented in the stacked :class:`FusedMoE` buffers: the
quant swap converts each expert linear BY PATH into its backend's
GEMM module, so tokens must be gathered and run module-by-module.
The ``unique()``/``nonzero()`` walk host-syncs, so a model that
reaches it is not capturable at any shape. The holder answers
:attr:`~arbi_serve.models.moe_grouped_exl3.GroupedRoutedExperts.capture_safe`
with False and
:func:`~arbi_serve.runtime.capture.preflight._capture_unsafe_component`
turns that into a preflight refusal; the guard below turns any
violation of that contract into a loud error instead of a corrupted
capture.

Expert ids outside ``[0, len(experts))`` contribute zero (EP local
frame); tokens routed to the same expert batch into one call.

DeepSeek-V2 / V3 sparse MoE FFN: router + routed experts + shared experts.

Topology mirrors the official ``DeepseekV2MoE`` /
``DeepseekV3MoE``::

    router_logits = gate(x)              # (N, n_routed_experts)
    weights, ids  = topk(scoring(router_logits), top_k)
    routed_out    = Σ_k weights_k · expert_{ids_k}(x)
    shared_out    = shared_experts(x)    # always-on dense MLP
    y             = routed_out + shared_out

The gating math is DeepSeek-specific and read off ``config.json``
(passed via ``moe_routing``):

  - ``scoring_func`` — ``"softmax"`` (V2 / V2-Lite) or ``"sigmoid"``
    (V3 ``noaux_tc``). V2-Lite uses softmax.
  - ``norm_topk_prob`` — renormalize the top-k weights to sum to 1.
    V2-Lite has this **off** (``norm_topk_prob=false``).
  - ``routed_scaling_factor`` — post-multiply on the routed weights
    (2.5 for V3, 1.0 for V2-Lite).
  - ``n_group`` / ``topk_group`` — group-limited routing. V3 scores
    all 256 experts, keeps the ``topk_group=4`` best of ``n_group=8``
    groups (group score = sum of the group's top-2 members), masks
    the rest to ``-inf`` and takes the top-8 of what survives.
    V2-Lite has ``n_group=topk_group=1``, which is a plain top-k.
  - ``topk_method="noaux_tc"`` (V3) adds a per-expert
    ``e_score_correction_bias`` that steers SELECTION only — the
    returned weights come from the unbiased sigmoid scores.
    V2-Lite uses ``topk_method="greedy"`` and carries no bias.

**Two routed-expert representations, decided by the CHECKPOINT**
(``moe_routing["fused_experts"]``, written by
:func:`_set_expert_dispatch` from the shared
:func:`~arbi_serve.models._moe_dispatch.resolve_expert_dispatch`
rule — every arch that builds this block shares it):

  - **Dense pack, and any quantized pack whose backend answers
    :meth:`~arbi_serve.weight_quant.base.QuantBackend.stacked_moe_spec`**
    (AWQ int4 → ``quant_kind="int4_w4a16"``) → ``experts`` is one
    :class:`FusedMoE`: stacked weights, one binned GEMM per stage,
    sync-free dispatch (the only production dispatch class).
  - **Quantized pack the backend cannot stack** (EXL3 trellis, a
    W4A8 pre-scale) → ``experts`` is a per-expert
    :class:`GatedSiLUMLP` list and the quant swap converts each
    expert linear BY PATH into its backend's GEMM module. Dispatch
    is :func:`_per_expert_sparse_dispatch` — host-syncing, and
    tripwired to raise rather than record a corrupt graph if it is
    ever reached under compile/capture.

The shared-experts MLP has width ``n_shared_experts ·
moe_intermediate_size`` and HF stores it under the single
``mlp.shared_experts.*`` prefix; it stays a per-linear module in
both representations (one MLP, nothing to stack).

One DeepSeek MLA decoder block with a sparse MoE FFN.

Same pre-norm + MLA + pre-norm topology as
:class:`_DeepseekV3MLADecoderLayer`, but the FFN is the
:class:`_DeepseekV2MoE` block (router + routed + shared experts)
instead of a dense gated MLP. Used for every layer indexed
``>= first_k_dense_replace``.

One DeepSeek dense (PAGED_KV) decoder block.

Same pre-norm + attn + pre-norm + MLP topology as the MLA variant,
but the token mixer is the standard :class:`AttentionBlock`. Used
only when a checkpoint mixes MLA and dense layers.

Pick the right DeepSeek V3 decoder layer class for ``layer_spec``.

The single ``if (kind, moe) ==`` switch. A spec with ``moe`` set
(layer index ``>= first_k_dense_replace``) gets the MoE FFN; the
rest get the dense gated MLP.

Holds the ``model.*`` subtree so safetensors keys line up.

The ``layers`` :class:`nn.ModuleList` carries the **inner** decoder
layer instances — one per spec, in spec order. The LayerStack
adapter wrappers (built by :class:`DeepseekV3Model.__init__`)
hold back-references to these same ``inner`` modules rather than
fresh copies, so PyTorch's ``named_parameters`` walk reaches each
parameter via exactly one path (``self.model.layers.{i}.*``) and
safetensors weight-map keys resolve as expected.

DeepSeek V3 / V2 / Kimi K2 causal LM with dense or sparse MoE FFN.

Compile is disabled (``enable_if=False``) because the MLA custom op
(``arbi_serve::mla_attention``) is registered with
``device_types="cuda"`` and cannot be traced by Dynamo. Attempting
torch.compile on this model would hit
"Hit PythonDispatcher dispatch key but PythonDispatcherTLS was not
set" — which corrupts the internal dispatch state and makes every
subsequent torch operation fail.

MLA attention runs RoPE on the q_pe / k_pe tail and reconstructs V
via ``kv_b_proj``. Layers at or above ``first_k_dense_replace``
carry the sparse MoE FFN (:class:`_DeepseekV2MoE`); ``weight_map``
declares its expert / router weights accordingly.

Write the checkpoint's routed-expert dispatch into ``moe_routing``.

Shared by every arch that builds :class:`_DeepseekV2MoE` — the rule
itself lives in
:func:`~arbi_serve.models._moe_dispatch.resolve_expert_dispatch`;
this only supplies the DeepSeek expert-path convention
(``model.layers.{i}.mlp.experts.{e}.{gate,up,down}_proj``, MoE from
``first_k_dense_replace`` up) and the transport keys
:class:`_DeepseekV2MoE` and :func:`_add_deepseek_moe_keys` read.

MoE FFN keys: router gate + routed experts + shared experts.

DeepSeek HF stores:
  - ``mlp.gate.weight`` — the router (dense; never quantized).
  - ``mlp.gate.e_score_correction_bias`` — the ``noaux_tc``
    selection bias (V3 / Kimi K2 only). It lives on the ROUTER in
    the checkpoint but on the module that routes here: the
    ``FusedMoE`` expert stack on a dense pack, the MoE block itself
    on a per-expert (weight-quantized) pack.
  - ``mlp.experts.{e}.{gate,up,down}_proj.weight`` — one gated MLP
    per routed expert.
  - ``mlp.shared_experts.{gate,up,down}_proj.weight`` — the single
    always-on shared MLP (present iff ``n_shared_experts > 0``).

``fused`` selects the destination shape and must match the
constructed block (see :class:`_DeepseekV2MoE`): a dense pack binds
each checkpoint expert into its slot of the stacked FusedMoE buffers
via the arch-agnostic
:func:`~arbi_serve.models._moe_weight_map.add_stacked_moe_expert_keys`;
a weight-quantized pack keeps per-expert module keys (the quant
swap's ``_filter_for_quant`` drops the dense ``.weight`` entries it
replaced; the router gate stays dense).

Construct a DeepSeek V3 / V2 / Kimi K2 model.

Canonical ctor — :class:`DeepseekV3Config` is optional only
because the current safetensors loader already encodes every
shape fact on ``layer_specs``. :meth:`from_safetensors`
always populates ``arch_cfg``.

One-time hook to bind kv_b_proj weights into each per-layer attn op.

Called by the engine after constructing ``attn_ops`` for the MLA
backend. Without this, :class:`MLAAttnOp.forward` cannot
reconstruct V from the cached compressed latent.

Materialize one RoPE cache per layer.

MLA layers RoPE on the q_pe / k_pe tail (``qk_rope_head_dim``),
not the full head dim; dense layers RoPE on the full head dim.
The :class:`RoPECacheRegistry` deduplicates so identical
``(rotary_dim, theta, interleaved, yarn)`` tuples share one cache.

MLA rotates in the INTERLEAVED (GPT-J) convention — channel
``2i`` pairs with ``2i+1``. The reference ``modeling_deepseek.py``
spells this as a de-interleaving ``view(d//2, 2).transpose(-1,-2)``
ahead of a NEOX ``rotate_half``; vLLM spells it
``is_neox_style=False``.

Construct + bind weights from a DeepSeek V3 / V2 / Kimi K2 export.

Reads ``config.json`` and builds one
``LayerSpec(kind=LayerKind.MLA, mla=MLAConfig(...))`` per layer.
``head_dim`` on the LayerSpec is set to ``qk_nope_head_dim +
qk_rope_head_dim`` so RoPECacheRegistry / scale derivation
downstream pick the right size; ``num_kv_heads`` is set to
``num_attention_heads`` (the per-head reconstruction means
post-W_UV K and V have num_attention_heads heads, not 1).

Per-parameter shard spec for the generic loader.

Maps the MLA attention projections and, per layer, either the
dense gated MLP or the sparse-MoE expert / router weights
(:func:`_add_deepseek_moe_keys`), using the standard DeepSeek
HF naming convention.

DeepSeek-V4 (``DeepseekV4ForCausalLM``) — sparse latent attention + FP4 MoE.

Topology::

    embed → expand to hc_mult streams
          → [hc_pre → attn_norm → DSv4LatentAttention → hc_post
             hc_pre → ffn_norm  → DSv4MoE             → hc_post] × N
          → hc_head → norm → head
    (+ speculative stages under mtp.*, whichever shape
     :data:`_MTP_STAGE_KINDS` registers for this checkpoint's ``mtp_kind``)

The checkpoint's names are top-level (``layers.N.attn.*``, ``embed.weight``,
``head.weight``, ``mtp.N.*``) with no ``model.`` prefix, so the module tree
mirrors them exactly and :meth:`DeepseekV4Model.weight_map` needs no
renames apart from the hyper-connection tensors, whose flat
``hc_attn_fn`` / ``hc_attn_base`` / ``hc_attn_scale`` triple binds into one
:class:`DSv4HyperConnection` submodule.

Payloads, and who binds them:

  * routed experts — either E2M1 pairs (``I8`` on disk) with a scale per 32
    reduction elements, which IS the MXFP4 layout :class:`FusedMoE` already
    stacks, or block-wise FP8 E4M3 on the same 128x128 grid as the other
    linears. ``expert_dtype`` says which; both bind through the shared
    stacked-expert helper with no repack.
  * every other linear — FP8 E4M3 with a 128x128 E8M0 scale grid under a
    ``.scale`` suffix, bound by :class:`DeepseekV4FP8Backend`.
  * norms, gates, the compressor projections, the hyper-connection
    tensors, ``attn_sink`` and ``tid2eid`` — dense, bound here. The fp32
    and int64 ones keep their on-disk dtype: rounding a Sinkhorn input or
    an expert-id table to bf16 is not a precision trade, it is a different
    model.

DeepSeek-V4 causal LM.

Compile and cudagraph capture are both structurally refused for this
arch today: the per-query id lists are built inside the forward (the
indexer's top-k is data), and the compressor advances a rolling
accumulator per token. Neither is expressible as a fixed recorded
graph without the id-building and pooling moving into a kernel.

Attention keys for one layer or speculative stage.

The two low-rank stages shard on the axes their shapes imply: ``wq_b``
and ``wo_a`` are column-parallel (whole heads / whole groups per rank),
``wo_b`` row-parallel, and the shared bottlenecks ``wq_a`` / ``wkv``
stay replicated because every rank's heads read the same latent.

Keys a DSpark stage carries beyond the shared block.

The stages form a chain, so these are positional: the first holds the
projection that reads the tapped target layers, the last the output
norm and the two ranking heads. The embedding and LM head are the main
model's, so neither is named here.

Keys a next-token stage carries beyond the shared block.

Every stage of this shape carries the same set, so its position in the
chain selects nothing. The vocab pair is the one variable: it is named
here only for a checkpoint that ships it, because a stage that reads
the trunk's has no such tensors and no modules to bind them into.

One speculative head shape: what is built, and what it binds.

Both halves in one row so they cannot drift — a stage module whose own
keys nobody names leaves its parameters unfilled, and keys naming a
module nobody builds fail the load.

Give every meta BUFFER real storage AT ITS OWN DTYPE.

The shared meta-materialize walk covers parameters and gives them the
engine dtype. This arch keeps its fp32 tensors (hyper-connection
weights, ``ape``, ``attn_sink``, the norms, the fp32 projections) and
its int64 ``tid2eid`` as buffers precisely so that walk cannot widen or
round them: an expert-id table cast to bf16 routes to the wrong
experts above 256, and a Sinkhorn input rounded to bf16 changes every
mixing weight in the model.

Returns the number of buffers materialized.

Meta-build → quant swap → materialize (params, buffers, experts) → load.

The shared :func:`~arbi_serve.models._model_load.build_and_fill` does
not materialize buffers, which for most arches is right — theirs are
computed, not checkpoint-backed. Here they carry loaded weights, so the
dtype-preserving pass runs between the parameter walk and the load.

Construct + bind weights from a DeepSeek-V4 export.

The speculative stage count is cross-checked against the
checkpoint's own keys here rather than trusted from the config,
which can declare fewer stages than the tensors carry.

Run the stack and return per-sequence last-token logits.

``attn_ops`` is unused: this arch's attention reads its own state
view directly rather than through a per-layer op, the same way the
recurrent mixers do.

Per-layer cos / sin tables.

Two distinct tables across the stack: the compressed layers extend
positions with the YaRN schedule on ``compress_rope_theta``, the
window-only layers take plain RoPE on ``rope_theta``. The registry
deduplicates, so this is two allocations however many layers ask.

Checkpoint keys the quant backends consume, which
:meth:`weight_map` therefore does not name.

The FP8 scale grids ride with their weights through
:class:`DeepseekV4FP8Backend`; the routed experts' E8M0 grids ride
through the stacked-expert binding. Naming them here is what lets a
coverage check account for every tensor in the checkpoint index
rather than silently tolerating the ones no entry consumes.

Does this checkpoint key's linear carry an FP8 scale grid?

A compressor's ``wkv`` shares the attention projection's leaf name
and is dense bf16, so the compressors are excluded by path: naming
a grid they do not ship would have a coverage check credit tensors
the checkpoint does not contain.

Gated Delta Net block — Qwen3-Next / Qwen3.5 / Qwen3.6.

The block owns its projections, depthwise causal conv weights, and
delta-rule parameters. The per-request recurrent state lives in
:class:`arbi_serve.cache.recurrent_pool.RecurrentStatePool` — the block is
stateless w.r.t. KV; the engine threads in the per-layer state view on
every step.

Forward kernels
---------------

GPU runtime: ``chunk_gated_delta_rule`` (prefill) and
``fused_recurrent_gated_delta_rule`` (decode) from
``flash-linear-attention``. The block constructs without FLA installed; the
FLA path is taken only when ``hidden_states`` is on CUDA AND FLA is
importable.

CPU smoke / shape tests — pure-PyTorch reference paths cover prefill and
decode (in :mod:`arbi_serve.models._gdn_reference`). They are slow but
numerically faithful, used as the unit-test oracle.

Implementation layout
---------------------

``GDNBlock`` is assembled from mixins (inheritance — zero hot-path
indirection):

  * :class:`arbi_serve.models._gdn_reference._GDNReferenceMixin` — CPU paths.
  * :class:`arbi_serve.models._gdn_fla._GDNFLAMixin` — FLA prefill / verify.
  * :class:`arbi_serve.models._gdn_fla_decode._GDNFLADecodeMixin` — decode.

Optional GPU kernel globals + env-gate readers live in
:mod:`arbi_serve.models._gdn_kernels` and are read via the ``kn`` module
alias so unit tests monkeypatch the single ``_gdn_kernels`` namespace.

One Gated Delta Net block — production form, FLA-backed.

Multi-head shape (Qwen3-Next, Qwen3.5 / Qwen3.6):

  * num_v_heads, num_k_heads — value / key head counts (GVA when
    ``num_v_heads > num_k_heads``, must be divisible).
  * head_v_dim, head_k_dim — per-head widths.
  * key_dim   = num_k_heads * head_k_dim
  * value_dim = num_v_heads * head_v_dim
  * conv_dim  = 2 * key_dim + value_dim   (Q ⊕ K ⊕ V)

Weight layout — separate projections (matches Qwen3.5 / Qwen3.6
HF safetensors export; Qwen3-Next checkpoints use a fused
``in_proj_qkvz`` and bind through this block's ``in_proj_qkv`` /
``in_proj_z`` via the model's ``weight_map``):

  * ``in_proj_qkv``  — hidden → ``[Q, K, V]`` of widths ``[key_dim,
    key_dim, value_dim]`` (fused merged-column linear).
  * ``in_proj_z``    — hidden → ``value_dim`` (the gate stream).
  * ``in_proj_b``    — hidden → ``num_v_heads`` (delta-rule β).
  * ``in_proj_a``    — hidden → ``num_v_heads`` (decay rate α).
  * ``conv1d``       — depthwise causal 1-D conv over the
    concatenated ``[Q, K, V]`` channels of width ``conv_dim``.
  * ``A_log``, ``dt_bias`` — per-head decay parameters; the
    log-space gate is ``g = -exp(A_log) * softplus(a + dt_bias)``.
  * ``norm``         — :class:`_RMSNormGated`, head-wise over the
    value-head output, gated by ``z``.
  * ``out_proj``     — ``value_dim → hidden_size``.

Forward dispatch:

  * Multi-token prefill (per-request token count > 1) — pack to
    a varlen ``(1, T_total, ...)`` and call
    :func:`fla.ops.gated_delta_rule.chunk_gated_delta_rule` with
    ``cu_seqlens``. Final per-request states are written back
    into :attr:`GdnLayerView.recurrent_state`.

  * Per-request-one-token decode (``n_tokens == B``) — stack into
    a ``(B, 1, ...)`` batched call to
    :func:`fla.ops.gated_delta_rule.fused_recurrent_gated_delta_rule`
    seeded from each request's prior recurrent state.

  * CPU smoke / shape tests — when running without CUDA + FLA the
    block falls back to a numerically faithful PyTorch reference
    for both prefill and decode paths so unit tests need no GPU.

The conv state is updated via the standard HF rolling-buffer
semantics — we keep the layout channels-first ``(conv_dim_local,
conv_kernel)`` so live use can switch to ``causal_conv1d_fn`` /
``causal_conv1d_update`` without changing the slab shape.

Return the ``[0, 1, 2, ..., B]`` cu_seqlens tensor for B.

Allocated once per (B, device) pair, cached on the block as
``_cu_seqlens_cache``. Replaces the per-call
``torch.arange(0, B+1, ...)`` in :meth:`_forward_decode_fla` —
~1 launch + Python alloc per GDN layer per decode step
eliminated. Cudagraph-safe because the engine captures decode
at fixed B-buckets ({1, 2}) and the cached tensor's
``data_ptr()`` is stable across replays.

The cache key is ``(B, device)`` — different capture buckets
coexist; rare cross-device calls (test fixtures) get their
own slot. Tensor dtype is int32 to match the kernel's
``cu_seqlens`` argument convention.

Return the fused ``[Q | K | V | Z]`` projection weight.

Shape: ``(2*key_dim_local + 2*value_dim_local, hidden_size)``.
Built once on first call from the four source parameters'
``.weight`` tensors; cached on the module under ``_fused_qkvz``.

Returns ``None`` when fusion is not applicable — currently when
either ``in_proj_qkv`` / ``in_proj_z`` has been swapped to a
quantized linear (the ``is_quantized`` duck-type marker), so its
``.weight`` is not a plain dense tensor. Mirrors
:meth:`_Qwen3_5AttentionBlock._get_fused_qkv_weight`; the caller
falls through to the unfused per-linear path.

Build the fused dense ``[B | A]`` projection weight buffer, ONCE,
in the load path (post weight-load + compaction, pre compile/capture).

``in_proj_b`` / ``in_proj_a`` have ``out=num_v_heads_local``, which
is below Marlin's tile bound (``out % 64 != 0``), so even on an AWQ
checkpoint they are kept DENSE bf16 (public AWQ packs skip sub-tile
projections; cf. ``in_proj_z / in_proj_b / in_proj_a`` are plain
``ColumnParallelLinear``). Either way the forward wants ONE fused
``[B|A]`` GEMM, not two — concatenate the two weights along the
OUTPUT axis. This buffer is read by ``_project_streams_raw_ba``'s
unfused-qkvz branch (the steady AWQ-qkvz config: ``in_proj_qkv`` /
``in_proj_z`` Marlin-quantized → ``_get_fused_qkvz_weight()`` is
None → that branch is taken, and WITHOUT a prebuilt fused [B|A] it
would fall to two separate dense GEMMs per layer instead of one,
adding launch overhead in the MTP verify pass.

Two cases, both producing one contiguous ``(2*num_v_heads_local,
hidden_size)`` buffer:

  * dense bf16 b/a (the 27B-AWQ + Qwen3.5/3.6 case): concatenate the
    two ``.weight`` tensors directly — identical to what
    :meth:`_get_fused_ba_weight` builds lazily, but materialized here
    so the compiled forward reads a stable prebuilt attribute.
  * AWQ-quantized b/a (no public pack does this today, kept for
    completeness): dequantize each into ``dtype`` (the activation
    dtype the unfused fallback casts to) then concatenate.

MATHEMATICALLY exact in both cases — concat-GEMM-then-split computes
the same dot products as two GEMMs over the same weights / input (and
per-output-row dequant+cast then concat == concat then cast). Every
output element is the same sum of the same products.

NOT bit-exact, and nothing here can promise that. The two forms differ
only in the GEMM's N, and BLAS is free to choose a different
microkernel, blocking or vectorisation for a different N — which
reassociates the length-``hidden_size`` reduction and can move the
last place. Whether it does is a property of the CPU/BLAS build, not of
this code: measured over 200k draws it is bit-identical on Zen 4
(oneDNN, 1-8 threads) and it is NOT on the CI runner, which is what
made ``torch.equal`` in the accompanying tests fail there. The
difference is reassociation rounding on a length-K fp32 dot product,
bounded by ~K*eps relative; a genuine fusion fault (wrong half, wrong
split, stale buffer) is O(1) instead, which is the scale the tests
discriminate at.

Stores the result on ``self._fused_ba_dequant`` (a plain attribute,
set ONCE here — never lazily in the forward — so the compiled
fullgraph forward reads it unconditionally without a
``getattr(...None)`` / device / dtype branch that would force a
Dynamo recompile).

Fails loud — never silently degrades — if exactly one of the pair is
quantized, or if a quantized linear is Marlin-active (it has no live
dense buffers to dequantize), since both break the bit-exact fusion.

Return the prebuilt fused dense ``[B | A]`` projection weight.
Built once by :meth:`prepare_fused_ba_dequant` in the load path
(dense-cat or AWQ-dequant-cat); the forward reads it unconditionally.
``None`` only before the load-path prebuild has run.

Return a fused ``[Q|K|V|Z]`` Marlin GEMM descriptor for the AWQ
decode path, or ``None`` when the linears aren't fusable.

``in_proj_qkv`` and ``in_proj_z`` are both column-parallel AWQ
linears reading the same ``hidden``; their already-Marlin-repacked
weights concatenate along the OUTPUT axis bit-exactly (see
:meth:`AWQLinearBase.build_fused_marlin`). Fusing them runs ONE
Marlin GEMM per GDN layer per decode step instead of two, cutting
the kernel-tail + launch count across the model's GDN layers.
Built once and cached; the cached descriptor's tensors have stable
``data_ptr`` (the constituent repacked buffers are load-time and
never re-allocated), so it is cudagraph-capture-safe.

Run the ``[Q|K|V|Z]`` projection + 4-way split — shared by
:meth:`_project_streams` and :meth:`_project_streams_raw_ba`.

Returns ``(q, k, v, z, qkv_view, fused_ba)``:

  * ``q, k`` : ``(N, num_k_heads_local, head_k_dim)``
  * ``v, z`` : ``(N, num_v_heads_local, head_v_dim)``
  * ``qkv_view`` : the ``[Q‖K‖V]`` conv-input stream — a
    ``narrow`` view of the contiguous fused ``qkvz`` (fused /
    Marlin branches) or the raw ``in_proj_qkv`` output
    (four-matmul branch). NOT forced contiguous here; the raw-ba
    caller applies ``.contiguous()`` where its stride contract
    needs it (only the fused-``F.linear`` branch keeps the bare
    view). :meth:`_project_streams` ignores this.
  * ``fused_ba`` : the fused ``[B|A]`` GEMM result
    (``F.linear(hidden, _get_fused_ba_weight())``) when the
    fused branch was taken, else ``None`` — lets each caller
    split ``b/a`` out without recomputing the GEMM.

``gate_fusion_on_cuda`` gates the fused-weight lookups on
``hidden.is_cuda`` (the :meth:`_project_streams` CPU-reference
contract). The GPU-only raw-ba path passes ``False`` so its
fusion selection is unchanged on every device.

Run the four input projections and split QKV.

Returns ``(q, k, v, z, beta, g)`` with shapes:

  * ``q, k`` : ``(N, num_k_heads_local, head_k_dim)``
  * ``v, z`` : ``(N, num_v_heads_local, head_v_dim)``
  * ``beta`` : ``(N, num_v_heads_local)``
  * ``g``    : ``(N, num_v_heads_local)`` — log-space decay
    ``-exp(A_log) * softplus(a + dt_bias)`` (per HF Qwen3.5).

Input-matmul fusion: on GPU we use two fused matmuls
(``[Q|K|V|Z]`` and ``[B|A]``) instead of four. The CPU
reference path still exercises the
four separate modules so the existing unit tests / weight-
loader contract continue to work without a host-side rebuild
of the fused buffer (which is GPU-resident).

The two halves fuse INDEPENDENTLY. On a quantized checkpoint
(AWQ / EXL3 / FP8) ``in_proj_qkv`` / ``in_proj_z`` have no dense
``.weight``, so ``_get_fused_qkvz_weight()`` is ``None`` and
``_project_qkvz`` returns ``fused_ba=None`` — but ``in_proj_b`` /
``in_proj_a`` are sub-tile (``out=num_v_heads_local``) and stay
DENSE on every shipped pack, so the ``[B|A]`` half still fuses off
the load-path-prebuilt ``_fused_ba_dequant`` buffer. Mirrors the
decode twin :meth:`_project_streams_raw_ba`, which has recovered
the same half since the buffer was introduced; without this arm
the GPU prefill / verify path ran FOUR matmuls on exactly the
checkpoints the docstring above promises two for.

Variant of :meth:`_project_streams` that returns ``b`` and ``a``
BEFORE the sigmoid / softplus / log-space exp — i.e. exactly what
the vendored ``fused_sigmoid_gating_delta_rule_update`` kernel
consumes (it does the gating math in-kernel).

Returns ``(q, k, v, z, qkv_concat, b_raw, a_raw)`` with shapes:

  * ``q, k`` : ``(N, num_k_heads_local, head_k_dim)``
  * ``v, z`` : ``(N, num_v_heads_local, head_v_dim)``
  * ``qkv_concat`` : ``(N, conv_dim_local)`` — VIEW of the
    ``[Q‖K‖V]`` slice of ``qkvz``; equivalent to
    ``torch.cat([q,k,v], -1)`` but allocation-free (the source
    tensor is contiguous and the slice preserves stride). Used
    as the input to ``causal_conv1d_update`` on the decode hot
    path.
  * ``b_raw, a_raw`` : ``(N, num_v_heads_local)``

Only used on the GPU fused-kernel path; CPU smoke retains the
post-gating ``_project_streams`` to keep the reference unit
tests (which compare against the four-matmul HF semantics)
unchanged.

Concatenate the per-head Q/K/V projections into the flat
``(n_rows, conv_dim_local)`` ``[Q‖K‖V]`` conv-input stream.

Shared by the prefill / verify / legacy-decode FLA paths, which
all build the same channels-last conv input before the depthwise
causal conv. ``n_rows`` is ``T_total`` (prefill / verify) or
``B`` (decode) — passed explicitly so the symbolic-shape
relation traced under ``@support_torch_compile`` matches each
caller's existing trace exactly.

Split a post-conv ``[Q‖K‖V]`` stream back into its three
components along the last dim (widths ``[key_dim_local,
key_dim_local, value_dim_local]``).

Shared by the prefill / verify / decode FLA paths; each caller
reshapes the returned splits to its own kernel layout.

Per-token causal conv prefill — pure-PyTorch reference.

Returns ``(out, new_conv_state)`` where ``out`` is
``(T, conv_dim_local)`` and ``new_conv_state`` is the rolled
buffer to commit. Used for both CPU smoke tests AND the GPU
path (the reference is portable and small; the conv kernel
itself is not the GDN bottleneck on long contexts — the
delta-rule recurrence dominates).

``snap_buf`` is the optional per-token snapshot slice — when
non-None, the post-token-t buffer is copied into
``snap_buf[t]`` for ``t in range(T)``. The caller passes the
slab-row slice of the layer view's
``snap_conv_state`` tensor (shape ``(T_max, C, K)``); ``T_max
>= T`` is the caller's responsibility (in the uniform-K verify
regime the slate has ``T = step_k + 1 <= T_max = max_k + 1``). Pass
``None`` from non-MTP paths.

Run the GDN mixer for this step.

Routes to the FLA chunk kernel for multi-token prefill,
FLA fused-recurrent for batched per-token decode, and a
PyTorch reference when running without CUDA / FLA (CPU
unit tests).

One token of the gated delta rule. Returns (out_v, new_state).

Math mirrors :func:`torch_recurrent_gated_delta_rule` from
the HF reference: ``state ← state * exp(g) + k ⊗ ((v - k @
state) * beta)``; ``out = q @ state``.

Gated RMSNorm for the Gated Delta Net block.

:class:`_RMSNormGated` is the head-wise gated RMSNorm applied to the
value-head output of :class:`~arbi_serve.models.gdn_block.GDNBlock`,
gated by ``silu(z)``. On CUDA with the vendored Triton kernel available
it routes through :mod:`arbi_serve.kernels.rmsnorm_gated`; otherwise it
falls back to a pure-PyTorch reference.

The kernel-availability globals (``_HAVE_RMSNORM_GATED_TRITON`` /
``rms_norm_gated_silu``) live at module scope here so tests can
monkeypatch them on this module object; :meth:`_RMSNormGated.forward`
reads them as module globals (not ``from``-imported names) so the
patches take effect.

Return the ``native_gated_rmsnorm`` knob.

Imported locally: :mod:`arbi_serve.runtime_flags` pulls config, which
imports model code. ``runtime_flags()`` is a cached snapshot, so this is a
single read and the forward body reads a module-level constant.

RMSNormGated — applied per-head over ``head_v_dim``.

Mirrors the HF ``Qwen3_5RMSNormGated`` semantics: norm in fp32,
multiply by learnable weight in input dtype, then gate by
``silu(z)``. The weight has shape ``(head_v_dim,)`` and is
broadcast across heads / tokens.

On CUDA inputs with the Triton kernel available, the forward
routes through ``arbi_serve.kernels.rmsnorm_gated`` —
ONE Triton launch covering var → rstd → scale → silu-gate,
replacing the multiple PyTorch elementwise launches the reference
path issues (cast / pow / mean / rsqrt / mul / cast / mul / cast /
silu / mul / cast). CPU smoke / unsupported-shape paths fall back
to the pure-PyTorch reference.

Numerical equivalence vs the reference path is asserted in
``tests/test_rmsnorm_gated_triton.py`` (atol=1e-3, rtol=1e-2 on
8 random inputs across B={1,2,4} × num_v_heads={4,8,16} ×
head_v_dim={64,128,256}).

Gemma 4 — ``Gemma4ForConditionalGeneration`` (text tower + vision path).

Loads the text tower under ``model.language_model.*`` (config under
``config.text_config``). The vision tower (``model.vision_tower.*``) and
its ``model.embed_vision.*`` projector are bound only when the checkpoint
carries a ``vision_config`` AND ``ARBI_ENABLE_VISION`` is set (see
:meth:`Gemma4Model.from_safetensors`); text-only boot stays bit-identical
otherwise. Audio weights are ignored.

Per ``layer_types[i] in {"sliding_attention", "full_attention"}``:

  * Sliding layers use ``head_dim`` + ``num_key_value_heads`` + per-
    layer ``sliding_window``; full layers use ``global_head_dim`` +
    ``num_global_key_value_heads`` when set.
  * RoPE θ differs per type via ``rope_parameters[layer_type]``.
  * SWA is transparent to :class:`LayerStack` — every layer is
    :class:`StateKind.PAGED_KV`; the attention backend honours
    ``LayerSpec.sliding_window`` when binding each layer's AttnOp.

Per-arch shape metadata lives on :class:`Gemma4Config` (mlc-llm
pattern). Per-layer dispatch routes through :class:`LayerStack`;
:class:`_Gemma4DecoderLayer` implements the uniform LayerStack block
signature directly. Hetero head_dim, ``attention_k_eq_v``,
``num_kv_shared_layers``, and altup are all wired; ``from_safetensors``
raises loudly only when ``enable_moe_block`` is set (the per-block
top-K MoE FFN is not wired).

Gemma 4 RMSNorm: ``y = norm(x) * weight``.

Gemma 4 dropped the ``(1 + weight)`` parameterization Gemma 2 / 3
used: the safetensors store weights centred around 1 (init=ones)
and the layer math is plain ``* weight`` — see
:class:`transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm`.
Per HF the math also runs in fp32 internally and casts back to the
activation dtype only at exit.

Token embedding × √hidden_size at the model entry.

Gemma-family models scale embeddings by ``sqrt(hidden_size)`` before
the first decoder block; this is folded into the embedding lookup
rather than the first decoder layer to match the HF safetensors
layout (``embed_tokens.weight`` is the unscaled tensor).

One Gemma-4 transformer block (text-only path).

Implements the uniform :class:`LayerStack` block signature directly
— no separate shape-converter wrapper. Sliding-window attention is
invisible at this layer of dispatch: each layer's :class:`AttnOp`
is constructed by the engine with the matching
:attr:`LayerSpec.sliding_window` already baked in.

Layout::

    residual = h
    h = input_layernorm(h)
    h = self_attn(h, ...)
    h = post_attention_layernorm(h)
    h = residual + h

    residual = h
    h = pre_feedforward_layernorm(h)
    h = mlp(h)
    h = post_feedforward_layernorm(h)
    h = residual + h

The four-norm sandwich is the Gemma signature (matches Gemma 2 / 3 /
4 — differs from the Llama / Qwen3 two-norm pre-norm layout).

Holds the ``model.language_model.*`` subtree.

The Gemma-4 multimodal checkpoint nests the text tower under
``model.language_model.*`` (vs Qwen3's ``model.*``). Holding the
parent module here keeps safetensors keys lining up cleanly.

When ``altup_cfg`` is provided the altup model-level adapters are
materialised:
  * ``embed_tokens_per_layer``: a SECOND embedding table of shape
    ``(vocab_size_per_layer_input, num_layers * per_layer_dim)``
    — input-id lookup yields a per-layer per-token embedding.
  * ``per_layer_model_projection``: ``Linear(hidden_size, num_layers
    * per_layer_dim)``, applied to ``embed_tokens(input_ids)`` (the
    primary embedding output).
  * ``per_layer_projection_norm``: a (1+w) RMSNorm on
    ``per_layer_dim``.
The two pieces are summed (after norm + scaling) to form the
per-layer-input tensor each decoder layer consumes.

Generic embedding scaled by a constant — backs ``embed_tokens_per_layer``.

Same pattern as :class:`GemmaScaledEmbedding` but with a
parameterized scale (HF's per-layer embedding uses
``√per_layer_dim`` which differs from the main embedding's
``√hidden_size``).

``quant="int8"`` stores the table as per-row (per-vocab-entry)
symmetric INT8 plus a per-row scale (model dtype) and dequantizes
only the GATHERED rows. The per-layer table is the single biggest
tensor in the Gemma-4-E2B checkpoint (bf16); INT8 halves its
resident footprint, freeing VRAM for KV / concurrency. This is a
gather, so only ``N_tokens`` rows are ever dequantized per step —
a capacity lever, decode-neutral. The payloads are
registered as BUFFERS (no gradient; the quant convention the generic
loader binds through ``named_buffers()`` and the warm-reload path
captures via ``state_dict()``).

Gemma 4 multimodal-arch text-tower causal-LM.

Loads the text path of ``Gemma4ForConditionalGeneration``
checkpoints (E2B + 31B). Vision / audio weights in the same
safetensors are intentionally skipped — this class binds only the
LM tower.

Look up embeddings and scale by ``sqrt(hidden_size)``.

The scale is multiplied as a Python float (a tensor-scalar
multiply), NOT via ``out.new_tensor(scale)``: building a device
tensor from a host scalar inside the forward does a host→device
copy that is illegal under cudagraph stream capture
(``cudaErrorStreamCaptureUnsupported``). The scalar multiply is
bit-identical and capture-safe — required so the whole-forward
decode graph captures (e.g. the ``tkv-bypass`` head_dim>256
split-K decode path).

Bind the host-resident (CPU, bf16) PLE gather table.

``table`` is ``(vocab, dim)`` on the CPU — NOT pinned (it is only ever
RAM). Called once from :meth:`Gemma4Model.from_safetensors` on the
offload path.

Allocate the fixed device + pinned-host staging buffers.

``max_num_tokens`` bounds every forward token count (decode B, verify
N, and any captured prefill bucket ≤ ``max_batched_tokens``); the
buffers are sized once so their ``data_ptr`` is STABLE for the captured
graph to bake. Must be called by the engine BEFORE the cudagraph
capture sweep — a later reallocation would dangle every baked pointer.

Gather this step's PLE rows on the host + async-H2D into ``_dev_buf``.

``device_ids`` is the LIVE device ``input_ids`` (length ``N``) the
forward / captured graph attends with — the authoritative token ids
(``last_sampled_gpu`` on the GPU-resident decode path). We D2H them into
a pinned host buffer, ``index_select`` the host table (CPU), and
non-blocking-H2D the gathered rows into the fixed ``_dev_buf[:N]`` on the
CURRENT stream — issued BEFORE the replay / eager forward so the baked
``_dev_buf`` read sees fresh data, exactly the input-refresh pattern
``input_ids`` / ``positions`` use.

The D2H is a hard host↔device sync (a CPU ``index_select`` can only key
on host ids) — intrinsic to host-offloading a per-token gather on a
GPU-resident-decode engine, since the freshly-sampled token lives only
on the GPU. It is one tiny (N × int64) copy — the unavoidable cost of
the offload.

Pinned-staging reuse guard: the previous step's H2D reads ``_host_stage``
asynchronously, so we ``event.synchronize()`` (host wait) BEFORE
overwriting it — the same discipline as ``PiecewiseBuffers.
wait_staging_free``. Usually already fired; the D2H below dwarfs it.

Look up embeddings and scale by the configured ``embed_scale``.

Scalar-float multiply (not ``out.new_tensor(scale)``): a
host-scalar→device-tensor build inside the forward is illegal
under cudagraph stream capture. Bit-identical, capture-safe.

INT8 path: gather the INT8 rows AND their fp32 scales (two
index_selects — pure device ops, capture-safe), dequant the
gathered rows only (``q * scale``), then apply ``embed_scale``.
Never materializes the full bf16 table.

Offload path: the rows were already gathered on the host and
async-H2D staged into ``_dev_buf`` by :meth:`stage_rows` (issued on
this stream BEFORE the forward / replay). Read the fixed device
buffer and apply ``embed_scale`` — a pure device slice + scalar mul,
capture-safe, and it never touches the host table inside the graph.

Build a per-layer rope-cache list keyed by layer_idx.

Gemma 4 carries different ``(rope_theta, partial_rotary)`` per
layer-type — the registry deduplicates so each unique
``(head_dim, theta, partial)`` triple allocates one shared
cache. Returns a list whose ``i``-th entry is the cache for
layer ``i``; the LayerStack block adapter indexes into it by
``layer_idx``.

Size the PLE host-offload staging buffers (engine boot hook).

No-op unless PLE offload is active. Allocates the fixed device +
pinned-host staging on the model's device, sized to ``max_num_tokens``
(``max_batched_tokens`` — an upper bound on any forward's token count),
so the device buffer's ``data_ptr`` is stable for the cudagraph capture
sweep to bake. MUST run before any capture/warmup forward.

Gather + async-H2D-stage this step's PLE rows before forward/replay.

Called by the runner immediately before the captured-graph replay (or
the eager forward) so the fixed device staging buffer the graph reads is
refreshed with the live step's rows — the same input-refresh contract as
``input_ids`` / ``positions``. Keys on the LIVE device ``input_ids`` (the
authoritative tokens the graph attends with — ``last_sampled_gpu`` on the
GPU-resident decode path) so the PLE contribution and attention never
disagree. The host request record (``host_mirror``) is NOT usable: it
lags the GPU sample by one step until drained. No-op unless PLE offload
is active.

Run the model and return per-seq last-token logits.

``compute_logits=False`` (with ``return_hidden_state=True``) skips
the last-token ``lm_head`` GEMV and returns ``(None, hidden)`` — the
MTP verify pass recomputes ``lm_head`` over all K+1 flat rows itself,
so the last-token head here would be pure waste (mirrors
``Qwen3_5Model.forward``).

Per-layer head_dim / sliding-window / RoPE θ all live on the
:class:`LayerSpec` + the per-model ``_layer_rope_theta`` table;
the engine constructs each layer's :class:`AttnOp` with the
matching sw, so this forward body need only thread the per-
layer ``attn_op`` through. The :class:`LayerStack` iterator
dispatches every layer via the same uniform adapter — SWA is
invisible at this level (it's a backend concern).

When altup is active, computes per-layer-input embeddings at
model entry and threads ``per_layer_inputs`` (shape
``(N_tokens, num_layers, per_layer_dim)``) into the iterator;
each block adapter slices its layer's row before calling the
decoder layer.

Construct + bind weights from a Gemma 4 multimodal export.

Reads ``config.text_config`` for the LM-tower geometry, builds
per-layer :class:`LayerSpec` (per-layer-type head_dim from
``layer_types`` + ``global_head_dim``, per-layer SWA from
``layer_types``), and loads weights from
``model.language_model.*`` keys.

Raises loudly when ``enable_moe_block`` is set — the per-block
top-K MoE FFN is not wired.

Append the Gemma-4 assistant's 4 Q-only draft marker LayerSpecs.

Each draft layer is KV-shared with the verifier's LAST NON-SHARED
layer of the same attention type (sliding→13, full→14 for E2B) via
``kv_source_layer`` (the pool aliases its view; ``skip_kv_proj`` +
the backend's ``skip_kv_write`` make it read-only). ``attention_
scale=1.0`` matches the Gemma-4 softmax scale (validated by the
offline accept harness — ``head_dim**-0.5`` collapses accept).
Rope tables are extended so ``_resolve_per_layer_rope_caches``
covers the markers. Records the assistant path + text_config for
``build_mtp_driver`` (which builds + loads the head).

Per-parameter shard spec for the generic loader.

Gemma 4's multimodal checkpoint nests the text tower under
``model.language_model.*``. Per layer the four-norm sandwich +
per-head q/k norms + (when altup active) the per-layer
adapters + ``layer_scalar`` buffer are saved with HF naming.

Gemma 4 per-arch shape-metadata config (mlc-llm pattern).

Holds :class:`Gemma4Config`, its altup companion :class:`_AltupConfig`,
and the RoPE-parameter resolver helpers. The model module
(:mod:`arbi_serve.models.gemma4`) re-exports :class:`Gemma4Config`,
:class:`_AltupConfig`, and :func:`_resolve_rope_thetas` so their public
import path stays ``arbi_serve.models.gemma4``.

Per-model altup config — Gemma 4 ``hidden_size_per_layer_input`` mechanism.

Captures only the scalar config; the per-layer adapter modules
live on the decoder layer and the model-global projection /
embedding live on :class:`_Gemma4LanguageModel`.

``num_layers`` and ``hidden_size`` are mirrored from
:class:`ModelDims` for convenience (per-layer shapes are derived
from ``per_layer_dim``).

Gemma 4 per-arch shape-metadata dataclass.

Lift of the per-arch fields the model needs to construct itself,
extracted from HF ``config.text_config`` once at boot. Equivalent
to the mlc-llm ``LlamaConfig`` pattern: the model file is
*spec + impl*, not a tangle of HF-AutoConfig probing scattered
through ``__init__``.

:class:`ModelDims` (model-global) and :class:`LayerSpec` (per-
layer) remain the runtime-facing surfaces — :class:`Gemma4Config`
is the deploy-time canonical declaration the boot path can carry
around or serialize into a manifest. :meth:`from_hf_config`
derives one from a raw HF config dict; :meth:`to_layer_specs`
materializes the per-layer :class:`LayerSpec` list (Gemma 4's
homogeneous-PAGED_KV-with-per-layer-variation layout — same
``LayerKind.ATTENTION`` for every layer; sliding-vs-full
attention is a per-spec ``sliding_window`` value, NOT a
per-spec kind).

Args:
    hidden_size: residual width.
    num_layers: dense decoder block count.
    num_attention_heads: per-layer Q-head count.
    num_key_value_heads: per-layer K/V-head count for
        ``sliding_attention`` layers.
    num_global_key_value_heads: per-layer K/V-head count for
        ``full_attention`` layers (``None`` → inherit
        ``num_key_value_heads``).
    head_dim: per-head Q/K/V width on ``sliding_attention`` layers.
    global_head_dim: per-head Q/K/V width on ``full_attention``
        layers (``None`` → inherit ``head_dim``).
    intermediate_size: SwiGLU MLP intermediate width.
    vocab_size: tokenizer vocab size.
    max_position_embeddings: declared max-seq-len (drives the RoPE
        table size — the engine paged-KV cache may exceed this).
    rope_theta_full: RoPE base frequency for full-attention layers.
    rope_theta_sliding: RoPE base frequency for sliding-attention
        layers.
    rope_partial_full: partial-rotary factor for full-attention
        layers (1.0 = full rotation; E2B uses 0.25 here).
    rope_partial_sliding: partial-rotary factor for sliding-
        attention layers.
    rms_norm_eps: per-layernorm epsilon.
    tie_word_embeddings: True iff lm_head shares storage with
        embed_tokens.
    sliding_window: per-layer SWA window for ``sliding_attention``
        layers (``None`` if no sliding layers / SWA disabled).
    layer_types: per-layer type tag — entries in
        ``{"sliding_attention", "full_attention"}``, length
        ``num_layers``.
    attention_k_eq_v: True iff full-attention layers alias V→K
        (Gemma 4 31B). Sliding layers always carry
        ``tie_v_to_k=False``.
    num_kv_shared_layers: count of trailing layers that REUSE the
        K/V cache slot of an earlier same-type layer.
    final_logit_softcapping: tanh-softcapping on lm_head output
        (Gemma 4 caps at 30); ``None`` / 0 = disabled.
    per_layer_dim: ``hidden_size_per_layer_input`` — Gemma 4 altup
        per-layer width. ``0`` if altup is not active.
    vocab_size_per_layer_input: vocab dimension of the second
        (per-layer) embedding table when altup is active. Mirrors
        the main vocab when missing.
    use_double_wide_mlp: doubles the MLP intermediate on
        kv-shared layers (Gemma 4 E2B feature).
    enable_moe_block: per-block top-K MoE FFN toggle; not wired —
        True raises in :meth:`from_safetensors`.

Resolve (full_attention θ, sliding_attention θ) from text_config.

Gemma 4 carries a per-layer-type ``rope_parameters`` mapping. Each
sub-block holds at minimum ``rope_theta``. We pull the
full_attention θ as the model-global default (long-context wall-
time dominator) and the sliding_attention θ for the per-layer
table.

Falls back to a flat ``rope_theta`` field when ``rope_parameters``
is missing (defensive — older fixtures).

Derive a :class:`Gemma4Config` from the raw HF
``config.json`` dict (the *root* — this method picks
``text_config`` for the LM tower).

Applies the same precedence + defaults as
:meth:`Gemma4Model.from_safetensors`, lifted onto a frozen
dataclass. Raises :class:`ValueError` if the dict's
``architectures`` does not include
``Gemma4ForConditionalGeneration`` or if ``text_config`` is
missing.

Project to the runtime-facing :class:`ModelDims` (model-
global subset). Per-layer fields drop here — they live on the
per-layer :class:`LayerSpec` list returned by
:meth:`to_layer_specs`. ``rope_theta`` is the model-global
default; per-layer θ is carried separately on
:attr:`Gemma4Model._layer_rope_theta`.

Per-layer ``kv_source_layer`` table — Gemma 4
``num_kv_shared_layers`` mechanism.

The trailing ``num_kv_shared_layers`` layers each map to the
most recent earlier layer of the same ``layer_type`` within
the non-shared prefix. Layers in the prefix map to ``None``
(own slab).

Materialize the per-layer :class:`LayerSpec` list.

Every layer is :class:`LayerKind.ATTENTION` (so every entry
routes to :class:`StateKind.PAGED_KV` and
:class:`LayerStack` walks them uniformly). The per-layer
variation Gemma 4 carries — sliding-window size, per-type
head_dim, per-type num_kv_heads, partial-rotary factor,
attention_k_eq_v, kv_source_layer, double-wide MLP — all
lives on :class:`LayerSpec` fields the attention backend
and pool already consume.

gpt-oss (OpenAI ``GptOssForCausalLM``) — 20B / 120B.

Topology::

    Embedding
    → [PreNorm → AttentionBlock (+per-head sink, alternating SWA)
       → PreNorm → router + clipped-SwiGLU MoE] × N
    → RMSNorm → lm_head

Arch facts driven from ``config.json``:

  * ``layer_types`` — per-layer ``"sliding_attention"`` /
    ``"full_attention"``; a sliding layer's window is ``sliding_window``.
  * ``attention_bias`` — q / k / v / o all carry a bias.
  * per-head learned attention sinks (``self_attn.sinks``), surfaced via
    :attr:`LayerSpec.attention_sinks` and bound onto the per-layer attn
    op by :func:`arbi_serve.engine.attention_sinks.bind_attention_sinks`.
  * ``rope_scaling`` — YaRN (:class:`YarnRopeScaling`).
  * ``swiglu_limit`` — the expert activation clamp. The matching
    ``alpha`` is NOT a config key in any released gpt-oss checkpoint nor
    in ``transformers``' ``GptOssConfig``; it is read from
    ``swiglu_alpha`` when a checkpoint supplies one and otherwise takes
    :data:`_REFERENCE_SWIGLU_ALPHA`, the value hardcoded in
    ``transformers.models.gpt_oss.modeling_gpt_oss.GptOssExperts``.

Expert-stack layout: :class:`~arbi_serve.models.moe.FusedMoE` holds
``w13_weight`` as ``(E, 2N, H)`` with gate rows ``[0, N)`` and up rows
``[N, 2N)``, and ``w2_weight`` as ``(E, H, N)`` plus the per-expert
biases ``w13_bias`` / ``w2_bias``. Released gpt-oss checkpoints store
the experts MXFP4-packed as one 4-D tensor per stack per layer with
gate/up INTERLEAVED along the output dim, so
:meth:`GptOssForCausalLM.weight_map` covers the dense tensors only and
:func:`~arbi_serve.models._moe_stacked_loader.bind_stacked_expert_stacks`
binds the stacks.

gpt-oss MoE FFN — biased router + biased clipped-SwiGLU experts.

The router is a replicated ``hidden → n_experts`` linear WITH bias
(``mlp.router`` in the checkpoint). Its top-k values are renormalised
by :class:`FusedMoE`'s ``norm_topk_prob``, which makes
``softmax(all) → top_k → renormalise`` identical to the reference's
``top_k → softmax(top_k)``.

Checkpoint keys holding this model's expert stacks.

:meth:`weight_map` excludes them — one 4-D tensor per stack per
layer fills every expert slot, which no per-parameter
:class:`WeightShardSpec` expresses. :meth:`_bind_expert_stacks`
binds them; this names them so a coverage check accounts for
every key in the checkpoint index.

IBM Granite 4.1 dense.

Topology mirrors HF ``GraniteForCausalLM``:
  Embedding × embedding_multiplier
  → [PreNorm → Attn + residual_multiplier·sublayer → PreNorm
     → SwiGLU MLP + residual_multiplier·sublayer] × N
  → final RMSNorm → tied lm_head → logits / logits_scaling.

Granite differs from a stock LLaMA-shaped dense in four scalars
exposed in ``config.json``:

* ``embedding_multiplier`` — multiplies the embedding output.
* ``residual_multiplier`` — scales every sublayer (attn / MLP) output
  before the residual add.
* ``logits_scaling`` — *divides* logits before sampling (entropy
  expansion).
* ``attention_multiplier`` — replaces the default ``1/sqrt(head_dim)``
  softmax scale. Plumbed through ``LayerSpec.attention_scale`` so it
  flows to ``backend.make_attn_op(scale=...)`` without any kernel
  change.

No Q/K-norm (Granite ≠ Qwen3), no sliding window, no softcap, no MoE,
no Mamba — pure dense GQA + RoPE + SwiGLU + RMSNorm. It therefore rides
the same construction / forward stem as the other dense arches: the
shared :class:`AttentionBlock` (``use_qk_norm=False``),
:class:`PreNormDecoderLayer` residual skeleton (the ``residual_multiplier``
folds into the ``_mixer`` / ``_ffn`` hooks), and
:class:`LayerStackModelMixin` (RoPE registry, positional LayerStack
dispatch, quant-aware meta-materialize). Only the embedding scale and
the ``logits_scaling`` divide stay Granite-specific in ``forward``.

One Granite transformer block: pre-norm + attn + pre-norm + MLP.

Same dense pre-norm topology as LLaMA / Qwen3 with no per-head Q/K
norms, with the sublayer (attention / MLP) output scaled by
``residual_multiplier`` before each residual add. The scale folds
into the ``_mixer`` / ``_ffn`` hooks so the shared residual skeleton
in :class:`PreNormDecoderLayer` applies unchanged. Granite does not
fuse the post-attention norm with the residual add.

Run attention then MLP, each as a pre-norm residual scaled by the
Granite ``residual_multiplier``.

Uniform :class:`LayerStack` block signature
(``hidden, positions, meta, view, *extras``). ``attn_op`` arrives
PRE-RESOLVED on the hot path (the model dispatcher reads
``attn_ops[layer_idx]`` outside the trace via
``_arbi_attn_op_extras_idx``); the ``isinstance(..., list)`` branch
only fires on the eager/test path that passes the full per-layer
list positionally.

Kimi Delta Attention block — Ling-3.0 (``bailing_hybrid``).

Per-channel forget gate, so ``g`` is ``(T, HV, K)``. State and conv slabs
are GDN-shaped, hence ``LayerKind.KDA`` -> ``StateKind.GDN``. The kernels
take ``transpose_state_layout=True`` and read the pool's ``(N, HV, V, K)``
rows directly — no transpose at the boundary.

One KDA block.

Checkpoint keys per layer ``N`` (see the model's ``weight_map``):
``attention.{q,k,v}_proj`` fuse into :attr:`in_proj_qkv`;
``attention.{q,k,v}_conv1d`` fuse into the conv weight;
``f_proj`` decay, ``b_proj`` beta, ``g_proj`` output gate,
``A_log`` / ``dt_bias`` gate params, ``o_norm``, ``o_proj``.

Pure-PyTorch KDA — slow, portable, the unit-test oracle.

Per-channel decay: ``state ← state * exp(g)[:, :, None] + k ⊗
((v - k @ state) * beta)``, ``out = q @ state``, with q/k
L2-normalized to match ``use_qk_l2norm_in_kernel=True``.

Per-layer architecture spec — drives backend, state kind, cache shape.

Every model exposes ``layer_specs: list[LayerSpec]`` whose entries
declare per-layer geometry (head_dim, num_heads, num_kv_heads,
sliding_window) and per-layer architectural variants (MLA / Mamba / GDN
/ ShortConv / MoE / sinks). Homogeneous models like Qwen3 dense have N
identical :class:`LayerSpec` entries modulo ``layer_idx``; hetero models
like NemotronH / Qwen3-Next / DeepSeek-V3 / LFM2 mix kinds.

The :class:`StateKind` enum is the routing key the engine and pool use
to decide which backend's :class:`AttnOp` and which slab a given layer
talks to. Adding a kind = adding a backend that exposes
``state_kind() == NEW_KIND`` plus a per-state-kind pool view.

What kind of per-layer state lives in the cache.

Used by the engine to route the right per-layer view to the right
backend's :class:`AttnOp`. Adding a kind = adding a backend that
exposes it via ``AttentionBackend.state_kind(...)`` plus a slab
shape and a per-step metadata builder.

Per-layer Mamba-2 / SSM shape.

Drives :class:`~arbi_serve.models.mamba2_block.Mamba2Block` — the
NemotronH chunk_scan_combined + n_groups + headdim path.

Args:
    state_dim: SSM hidden state dimension (``ssm_state_size`` in HF).
    conv_kernel: depthwise causal-conv kernel length.
    intermediate_size: pre-out-proj width; equals
        ``num_heads * head_dim``.
    num_heads: SSM head count (``mamba_num_heads`` on NemotronH).
    head_dim: per-head V-state width (``mamba_head_dim`` on
        NemotronH).
    n_groups: number of B/C state groups (``n_groups`` on NemotronH).
    chunk_size: chunk-scan chunk size (``chunk_size`` on NemotronH).

Per-layer Gated Delta Net shape (Qwen3-Next, Qwen3.5 / Qwen3.6).

Multi-head fields mirror upstream HF config keys:

  * ``num_v_heads``  — ``linear_num_value_heads``  (e.g. 48 in 27B).
  * ``num_k_heads``  — ``linear_num_key_heads``    (e.g. 16 in 27B).
  * ``head_v_dim``   — ``linear_value_head_dim``   (e.g. 128).
  * ``head_k_dim``   — ``linear_key_head_dim``     (e.g. 128).

GVA (Grouped Value Attention) is implied when ``num_v_heads >
num_k_heads`` (must be divisible) — FLA's
:func:`chunk_gated_delta_rule` and
:func:`fused_recurrent_gated_delta_rule` apply it automatically.

``state_dim`` equals ``num_v_heads * head_v_dim`` (the value
stream width through the out-projection); the per-request
recurrent state matrix is ``(num_v_heads, head_k_dim,
head_v_dim)`` consumed by FLA.

Per-layer KDA shape (Ling-3.0). Extends :class:`GDNConfig`: the
state and conv slabs are identical, only the block math differs.

Args:
    lower_bound: clamp floor on the log-space decay; ``None`` = unclamped.
    safe_gate: guarded gate accumulation.
    fused_gate_proj: ``f_proj``/``g_proj`` are single linears, not the
        ``*_a_proj``/``*_b_proj`` low-rank pair.

Per-layer ShortConv shape (LFM2 / LFM2-MoE).

LFM2's mixer is ``in_proj -> (split B,C,x) -> 1-D causal conv on
(B*x) -> gate by C -> out_proj``. The conv operates per-channel along
the token axis; the per-request state buffer holds the last
``conv_kernel - 1`` activations needed to continue the conv across
autoregressive steps.

Args:
    conv_dim: channel count fed to the conv (LFM2: ``hidden_size``).
    conv_kernel: 1-D causal conv kernel length. LFM2-8B-A1B uses 3.
    bias: whether the conv has a bias (LFM2-8B-A1B: ``False``).

Per-layer DeepSeek-V4 sparse latent-attention shape.

One 512-d latent per token serves as BOTH K and V for all
``num_heads`` query heads (``head_dim`` on the :class:`LayerSpec`);
its trailing ``rope_head_dim`` channels carry RoPE and the leading
``head_dim - rope_head_dim`` are NoPE.

``compress_ratio`` selects the layer family:

  * ``0`` — sliding window only (the first two layers and every
    DSpark stage). No compressor, no indexer, plain RoPE.
  * ``4`` — window + a compressed stream pooled from every 4 tokens
    with OVERLAPPING windows, and an :class:`Indexer` that scores its
    own 4-token-pooled stream to pick ``index_topk`` of them.
  * ``128`` — window + a non-overlapping 128-token-pooled stream,
    every entry of which is attended (no indexer).

``kv_quant_group`` / ``fp4_quant_group`` are the QAT round-trip group
widths the checkpoint was trained with: the latent's NoPE channels are
quantize-dequantized through FP8 per ``kv_quant_group``, and the
indexer's Hadamard-rotated query / stream through FP4 per
``fp4_quant_group``. They are part of the forward numerics, not a
storage choice — dropping them changes the served distribution.

Per-layer architecture spec.

Drives backend, state kind, cache shape, and forward-kernel
parameters. For homogeneous models (all-Qwen3-dense), every entry
in the model's ``layer_specs`` list is equal modulo ``layer_idx``.

LFM2 hybrids: alternating ``ATTENTION`` / ``SHORT_CONV`` (and per-
layer dense-vs-MoE FFN selection in :class:`LFM2MoeForCausalLM`).

``backend_hint``: per-layer backend name override. Reserved for
per-layer mixed-bit-width work; not yet wired, so the runtime
asserts it is ``None`` at engine boot.

Standard transformer building blocks for arbi-serve.

Holds RMSNorm, Embedding, the GatedSiLUMLP MLP wrapper, and the
RoPECache + RoPECacheRegistry. The :class:`Linear` symbol is an alias
for :class:`ReplicatedLinear` so existing call sites still resolve;
new code should use :class:`arbi_serve.models.linear.ReplicatedLinear`
or one of the parallel variants directly.

Tensor shapes follow flat-token-batching conventions: every per-token
tensor is ``(N_tokens, ...)`` where N_tokens spans the whole batch.
Per-sequence reshapes happen at the attention boundary only.

Pre-norm RMSNorm with learnable scale.

``y = x * rsqrt(mean(x^2) + eps) * weight``

F.rms_norm exists in PT 2.4+ but is silently fp32-promoted on
bf16 inputs in some builds; this fp32-explicit form is portable
and matches the reference Qwen3 numerics bit-for-bit.

Two call shapes:

* ``norm(x)`` — plain RMSNorm of ``x`` via the fp32-explicit Python
  path.
* ``norm(x, residual)`` — fused (residual + RMSNorm) on CUDA. The
  ``residual`` argument is updated IN PLACE to ``residual + x``;
  the return is RMSNorm of the new residual. Replaces the manual
  ``x = residual + h; residual = x; x = pre_norm(x)`` chain in pre-
  norm transformer blocks. Falls back to the unfused Python path
  (eager add + plain ``forward``) on CPU or when the Triton kernel
  module isn't importable — same numerics either way.

Read the ``native_rmsnorm`` knob once, at import.

Function-local import: ``arbi_serve.runtime_flags`` pulls config which
imports model code, so a module-level import here would cycle.
``runtime_flags()`` is a cached snapshot, so this stays a single read and
the forward body keeps reading a plain Python constant — no per-call
branch for Dynamo to specialize on or graph-break over.

RMSNorm with the ``y = norm(x) * (1 + weight)`` convention.

Weights are stored zero-centred (init at zero, learn the offset);
the forward adds 1 before scaling. Used for the input / post-attention
/ final norms and the per-head q/k norms on Qwen 3.5 / 3.6.

Two call shapes (matching :class:`RMSNorm`):

* ``norm(x)`` — plain RMSNorm of ``x`` with the ``(1 + weight)``
  scale. fp32-explicit Python path.
* ``norm(x, residual)`` — fused (residual + RMSNorm) on CUDA. The
  ``residual`` argument is updated IN PLACE to ``residual + x``;
  the return is RMSNorm-of-new-residual scaled by ``(1 + weight)``.
  ``1.0 + self.weight.float()`` is cached (see :meth:`_folded`) and
  handed to the kernel, whose fp32 weight-mul is bit-equivalent to
  the eager ``(x32 * (1 + w)).to(in_dtype)`` ordering within bf16 ULP.
  Falls back to the unfused Python path on CPU or when the Triton
  kernel module isn't importable.

The fold is precomputed into the ``folded_weight`` buffer because
``self.weight`` is a live parameter the surrounding custom op can't
fold across — recomputing it per forward lowered as a standalone
kernel outside the op boundary (one launch per norm per step). The
buffer is (re)built by :meth:`prepare_for_compile` in the load path
(post-compaction, pre-capture), never lazily in the forward, so the
compiled fullgraph forward reads it unconditionally — no Python
branch, no ``.data_ptr()`` / ``._version`` guard that would force a
Dynamo recompile.

Token embedding lookup. Plain ``F.embedding`` underneath.

Held as a Parameter so the loader can copy weights into it
uniformly with the linears, and so tied lm_head re-uses the
same tensor without an alias hop.

Token embedding sharded along the vocab dimension across TP ranks.

On a TP=2 deployment of a vocab=152K, hidden=5120 model the dense
vocab axis: each rank owns ``vocab_size_padded / tp_size`` rows.

Forward contract: each rank looks up only its local rows; tokens
outside the local vocab range get a zero embedding, then a single
all-reduce sums the partial results so every rank ends up with the
same final embedding tensor (each token has exactly one nonzero
contributor by construction). When TP=1 the implementation
short-circuits to plain ``F.embedding`` — bit-equivalent to
:class:`Embedding`, no collective.

Vocab padding: the sharded slab has ``ceil(vocab_size / tp_size) ×
tp_size`` rows so every rank has the same shape (CUDAGraph capture
requires uniform shapes across ranks). Padding rows are zero —
they never get hit because token ids in ``[vocab_size, padded)``
are not tokenizer-reachable.

Tied lm_head: when the model ties the lm_head to the embedding
tensor (Qwen 0.8B / 9B), the lm_head matmul reads
``self.weight`` directly. Because the weight is sharded along the
vocab dim, the matmul produces a per-rank ``(N, vocab_per_shard)``
logit slice; the lm_head call site is responsible for the all-
gather to assemble the full vocab.

Callable lm_head that reads weights from a tied
:class:`VocabParallelEmbedding`.

The embedding's ``weight`` is sharded along the vocab dim so the
matmul produces a per-rank ``(N, vocab_per_shard)`` slice; an
all-gather assembles the full vocab logits for sampling. Holds
the embedding in a 1-element list so ``nn.Module.__setattr__``
doesn't double-register the parameter (the main model already
owns the embedding tensor).

Gated MLP: ``down(act(gate(x)) * up(x))`` (Llama/Qwen GeGLU family).

``activation`` selects the gate nonlinearity — ``"silu"`` (Llama/Qwen,
the default) or ``"gelu_tanh"`` (Gemma's ``gelu_pytorch_tanh``).

THREE gate/up strategies (mutually exclusive):

  * default — separate ``gate_proj`` / ``up_proj`` linears, one matmul each.
  * ``use_merged_gate_up=True`` — one :class:`MergedColumnParallelLinear`
    GEMM (gate‖up stacked). Fewer matmuls, but the fused dense param is NOT
    quant-swappable (EXL3/FP8 ``quant_class_for`` raises) — use only when
    the model is never served quantized.
  * ``lazy_fuse_forward=True`` — keeps SEPARATE ``gate_proj`` / ``up_proj``
    params (quant-swappable + per-projection LoRA) but issues ONE
    ``F.linear`` over a lazily-built ``[gate|up]`` weight on the CUDA
    non-LoRA hot path. Same numerics; quant-safe (falls back to the
    two-matmul path when a projection is quantized / on meta / LoRA-active
    / CPU). Costs one duplicated ``[gate|up]`` tensor of VRAM per layer, so
    prefer it for smaller intermediates. This is the mode Qwen3.5/3.6 use.

``tp_sharded=False`` builds all three projections replicated. A routed
MoE expert under full expert parallelism is the case: the rank owns
WHOLE experts, so the intermediate dim is not split and no projection
reduces.

YaRN NTK-by-parts RoPE scaling (``rope_scaling.rope_type == "yarn"``).

Fields mirror the HF ``config.json`` keys. ``attention_factor``
defaults to the paper's ``0.1 * ln(factor) + 1`` when the checkpoint
does not pin one.

``mscale`` / ``mscale_all_dim`` are the DeepSeek V2 / V3 extension:
the cos/sin table carries the *ratio* of the two magnitude terms and
the attention softmax scale carries ``mscale_all_dim``'s term
squared (:attr:`softmax_scale_factor`).

Pre-computed cos / sin table for NEOX- or interleaved-style RoPE.

``interleaved=False`` (NEOX, the default) pairs channel ``i`` with
channel ``i + head_dim/2``; ``interleaved=True`` (GPT-J) pairs
channel ``2i`` with channel ``2i+1``. DeepSeek V2/V3 MLA rotates its
``qk_rope_head_dim`` tail in the interleaved convention.
Cached up to ``max_position_embeddings`` at init; positions beyond
that raise (the engine clamps requests to ``max_context``).

``yarn`` replaces the plain ``theta ** (-2i/d)`` inverse frequencies
with the YaRN NTK-by-parts schedule and post-scales the whole table
by its ``mscale`` (gpt-oss ships ``factor=32`` over a 4096-token
pretrain window).

One :class:`RoPECache` per ``(head_dim, rope_theta)`` pair.

Hetero-head-dim models (vision encoders, MLA layers) need
multiple cos/sin tables; the registry keys a shared table per
geometry so layers with matching ``(head_dim, theta)`` reuse one
on-device tensor pair.

For Qwen3 dense, every layer's spec has the same ``(head_dim,
theta)``; the registry returns the same instance every call.

Fused residual-add + RMSNorm path.

On CUDA + when the vendored Triton kernel module imports
cleanly, dispatches through ``torch.ops.arbi_serve.fused_add_rms_norm``
— one launch instead of nine (1 add + 8 unfused norm). The
custom op declares ``residual`` as an in-place mutation so
Inductor / cudagraph capture sees a stable boundary.

Otherwise (CPU smoke, or the kernel module fails to import on
an unusual triton build): fall back to the plain unfused path
— bit-equivalent numerics, just no fusion.

Rebuild the folded ``1.0 + weight.float()`` buffer from the
live weight. Called in the load path after weights land on-device
and after flat-slab compaction, before compile/capture — and again
on any weight reload. Keeps the forward's read unconditional and
Dynamo-safe (no per-call cache branch).

Fused residual-add + RMSNorm with the ``(1 + weight)`` scale.

Hands the cached ``1.0 + self.weight`` fold to
``torch.ops.arbi_serve.fused_add_rms_norm``. Falls back to plain
eager (mutating ``residual`` in place to match the kernel
contract) on CPU or when the kernel module isn't importable.

RTN-quantize this rank's vocab shard to fp8-e4m3 with per-row scales.

Per vocab row ``r``: ``scale[r] = max(|W[r]|) / 448`` (the e4m3 finite
max) and ``q[r] = clamp(W[r]/scale[r], -448, 448)`` stored fp8. The bf16
``weight`` Parameter is then DROPPED so only the fp8 payload (half the
bytes) plus a per-row bf16 scale stay resident; :meth:`_lookup`
dequantizes just the gathered rows on the fly. Data-free (rows quantize
independently), so the result is bit-identical to slicing a globally
quantized table under any TP. Returns the resident quantized bytes.

``persistent`` controls whether the fp8 buffers join ``state_dict``. The
KV win depends on this: run PRE-compaction with ``persistent=True`` so
the streaming compactor packs the fp8 payload into its slab and the bf16
table never contributes to the (shared, un-shrinkable) bf16 weight slab
— the freed ~half table is what the KV pool grows into. Post-compaction
the bf16 table is already a view into that shared slab, so freeing it
reclaims nothing; the pre-compaction seam is the one that pays off.

Pre-bind fp8 slab views so the next :meth:`weight_loader` RTN-quantizes
the checkpoint table STRAIGHT into them (cold EXL3 direct-slab path).

The counterpart of :meth:`quantize_to_fp8` for the exl3-direct binder,
which sizes the embed slot as fp8 ``[vocab, hidden]`` + a per-row bf16
scale at slab-allocation time (``cfg.embed_quant == "fp8"``). The bf16
``weight`` Parameter — a meta placeholder here — is dropped and never
materialized in the weights slab: the loader quantizes into these views
instead of copying a bf16 table into a bf16 slot. Same fp8-e4m3 per-row
payload / scale / multiplier convention as :meth:`quantize_to_fp8`, so
:meth:`_lookup` serves it identically.

RTN-quantize ``src_rows`` (this shard's real vocab rows) into the
pre-bound ``weight_fp8`` / ``weight_scale`` slab views; zero any pad
rows. Runs once, from :meth:`weight_loader` under the fp8 slab-fill arm.

Drops the (never-materialized) ``weight`` placeholder at the end — its
only job was to resolve this dispatch.

Reshape this bf16 embed into empty quantized placeholders for the
WARM flat-dump reload (``mode`` == "fp8").

The warm path builds the model under ``skip_weight_load`` — the embed
arrives as a bf16 ``weight`` meta param. The cold dump this boot reloads
was captured with the embed ALREADY quantized (``weight_fp8`` +
``weight_scale`` persistent buffers, no bf16
``weight``), so the warm graph must match: drop the bf16 ``weight`` and
register EMPTY quant buffers, which :func:`_presize_placeholder_buffers`
then grows to the dump's shape and the DMA fills. The bf16 table is
never allocated, so the KV pool grows on the warm path exactly as it
does cold. Idempotent-safe: a no-op if already quantized.

:meth:`forward` writing the logits into ``out``; returns ``out``.

Fuses only at TP=1, where the local matmul IS the full-vocab
answer. At TP>1 the all-gather owns its destination and the trim
is a slice of it, so there is nothing to write into — fall back to
forward + copy (correct, just unfused). The tied head takes no
adapter state; ``lora_state`` exists for signature parity with
:meth:`arbi_serve.models.linear.LinearBase.forward_into`.

Bind per-linear LoRA target keys so adapters resolve 1:1 with the
HF / PEFT export module names. Without this every ``lora_target_key``
stays ``None`` and LoRA silently no-ops. Merged fuses gate+up into one
adapter target; the separate / lazy-fuse forms keep them distinct.

Lazily build (+ cache) the ``[gate|up]`` weight, or ``None`` when
fusion does not apply (a projection quantized or still on meta). The
per-call check is a single ``is not None`` test so Dynamo emits no
guard on the cached tensor.

Positions this schedule spans, given the config's declared window.

HF's convention is that ``max_position_embeddings`` is ALREADY
post-extension and ``original_max_position_embeddings`` records
the pre-extension value — ``declared`` unchanged for a well-formed checkpoint.

It is load-bearing only when a schedule is OVERRIDDEN onto a
config whose ``max_position_embeddings`` was not bumped with it
(Ling-3.0's 256K recipe injects ``rope_scaling`` alone). Declaring
``factor`` against ``original`` states the span unambiguously, so
the arithmetic is the operator's own — not an inference from a
capacity knob.

Build from a full HF ``config.json`` mapping; ``None`` when the
checkpoint declares no RoPE scaling.

Raises ``NotImplementedError`` for any ``rope_type`` other than
``yarn`` — a silently-ignored scaling schedule is wrong output,
not a degraded one.

Apply RoPE in-place semantics (returns new tensors).

``q``/``k`` shape: ``(N_tokens, num_heads, head_dim)``.
``positions`` shape: ``(N_tokens,)`` int — absolute position
in the sequence. The engine builds positions per-step from
per-sequence start + per-token offset; flat 1-D simplifies the
gather here.

Look up a cached cos/sin pair from the engine's BootManifest
warm-restart cache, if one is attached.

Walks ``self._engine_ref._cached_rope_tables`` (set by
:meth:`arbi_serve.runtime.boot_artifacts.RopeCacheArtifact.hydrate`).
Returns the ``{"cos", "sin"}`` dict or ``None`` if no engine ref is
attached or no entry matches the key.

LFM2 — Liquid AI hybrid attention + ShortConv.

Per ``config.layer_types[i]``: ``"full_attention"`` (GQA, Q/K
LayerNorm, RoPE, paged-KV) or ``"conv"`` (:class:`ShortConvBlock` —
1-D causal conv + gate, per-request buffer in
:class:`ShortConvStatePool`). Dense MLPs throughout — the MoE variant
is :mod:`arbi_serve.models.lfm2_moe`.

Per-layer dispatch routes through :class:`LayerStack`; the inner
decoder layers implement the uniform LayerStack block signature
directly. The per-arch dispatcher
(:func:`_make_lfm2_decoder_layer`) picks the right class. Weight
keys mirror the HF ``model.layers.{i}.<name>`` convention (see
:meth:`weight_map`).

Per-arch construction bundle for :class:`LFM2ForCausalLM`.

Collects the LFM2-specific construction knobs that don't fit on
the shared :class:`ModelDims` (which is generic across all archs).
Following the mlc-llm pattern, the arch's :meth:`from_safetensors`
builds one of these and the constructor consumes it.

Args:
    short_conv_cfg: per-layer ShortConv config — same instance
        shared across every conv layer (LFM2 uses one consistent
        ShortConv shape across the whole stack). ``None`` for the
        theoretical pure-attention configuration.

Holds ``model.*`` subtree so safetensors keys line up.

Builds one inner decoder-layer instance per :class:`LayerSpec`
via :func:`_make_lfm2_decoder_layer`. The :class:`LayerStack`
adapter wrappers (built by :class:`LFM2ForCausalLM.__init__`)
hold back-references to these same ``inner`` modules rather than
fresh copies, so PyTorch's ``named_parameters`` walk reaches each
parameter via exactly one path (``self.model.layers.{i}.*``) and
safetensors weight-map keys resolve as expected.

LFM2 hybrid (attention + ShortConv) causal LM.

Layer interleave is read from ``config.layer_types[i] in
{"full_attention", "conv"}`` and baked into ``layer_specs[i].kind``.
Both attention and conv layers carry a dense ``GatedSiLUMLP`` FFN
(no MoE — the MoE variant is :class:`LFM2MoeForCausalLM` in
:mod:`arbi_serve.models.lfm2_moe`).

Produce ``(ModelDims, layer_specs)`` from the LFM2 / LFM2-MoE config dict.

Shared between :class:`LFM2ForCausalLM` and the MoE variant — the
geometry math (head_dim, GQA divisibility, layer_types interleave,
ShortConvConfig per conv layer) is identical; only the per-layer
FFN choice differs and is handled in the model classes.

Forward over the hybrid stack.

Per-layer dispatch lives on the per-kind decoder layer
(:class:`_LFM2AttentionLayer` / :class:`_LFM2ShortConvLayer`),
which implements the uniform LayerStack block signature
directly. The iterator threads per-:class:`StateKind` metadata
+ state-view from the dicts below; each layer pulls its
kind-specific extras out of the ``**extra`` bundle.

  - ATTENTION → consumes ``state_views[PAGED_KV]``
    (:class:`MultiStatePool`) and ``batch_meta[PAGED_KV] =
    batch.attn_meta``; per-layer ``attn_ops[i]`` threaded via
    ``**extra``.
  - SHORT_CONV → consumes ``state_views[SHORT_CONV]`` and
    ``batch_meta[SHORT_CONV] = batch.short_conv_meta``; no
    ``attn_op`` (no AttnOp seam for ShortConv).

LFM2-MoE — Liquid AI hybrid attention + ShortConv + sparse MoE.

Same hybrid backbone as :class:`LFM2ForCausalLM` (attention + ShortConv
mixers per ``layer_types[i]``) plus per-layer FFN choice driven by
``num_dense_layers``: layers ``[0, num_dense_layers)`` use dense
``GatedSiLUMLP``; later layers use the fused sparse FFN
(:class:`~arbi_serve.models.moe.FusedSparseMoEBlock` over
:class:`~arbi_serve.models.moe.FusedMoE` — stacked expert weights, one
binned GEMM per stage, sync-free dispatch). Routing follows
LFM2-8B-A1B's ``Lfm2MoeSparseMoeBlock``: sigmoid scoring with optional
learned bias, top-k = ``num_experts_per_tok``, optional
``norm_topk_prob`` / ``routed_scaling_factor``.

EP-aware: at ``--ep-size N`` each rank binds a contiguous slice of
``num_experts // ep_size`` experts; the router runs replicated and one
trailing all-reduce over the EP group sums each token's mixture.

Per-layer dispatch routes through :class:`LayerStack`; the inner
decoder layers implement the uniform LayerStack block signature
directly. The per-arch dispatcher
(:func:`_make_lfm2_moe_decoder_layer`) picks one of three flavors
keyed by ``(spec.state_kind(), spec.is_moe)``: PAGED_KV+dense /
SHORT_CONV+dense / either+MoE.

Per-arch construction bundle for :class:`LFM2MoeForCausalLM`.

Collects the ctor args that don't fit on :class:`ModelDims` (which
is shared across all archs) but are model-arch-specific knobs
derived from ``config.json``. Following the mlc-llm pattern, the
arch's :meth:`from_safetensors` builds one of these and the
constructor consumes it (rather than threading 4 positional args
through the call site).

Args:
    moe_routing: kwargs forwarded to
        :class:`~arbi_serve.models.moe.FusedSparseMoEBlock` for every
        MoE-bearing layer. Includes ``scoring_func``,
        ``norm_topk_prob``, ``routed_scaling_factor``,
        ``use_expert_bias``. Same instance shared across every MoE
        layer (LFM2-MoE uses one routing config across the whole
        stack — the only per-layer MoE knob is the
        :class:`MoEConfig` on the spec itself).

Attention mixer + dense MLP. Used when ``spec.moe is None`` and
``spec.kind == LayerKind.ATTENTION`` (i.e. ``layer_idx <
num_dense_layers`` and the attn variant of the layer interleave).

Implements the uniform :class:`LayerStack` block signature directly.

Either mixer + fused sparse MoE FFN. Used for every spec with
``spec.moe is not None`` (i.e. ``layer_idx >= num_dense_layers``).

Implements the uniform :class:`LayerStack` block signature directly;
branches on ``layer_spec.kind`` at forward time to pick the mixer
argv. The FFN half is identical across both mixer kinds.

Pick the right LFM2-MoE decoder layer class for ``layer_spec``.

The single ``if (kind, moe) ==`` switch — three distinct
classes by ``(spec.state_kind(), spec.is_moe)``:

  * ``(PAGED_KV, False)``  → :class:`_LFM2MoEAttnDecoderLayer`
  * ``(SHORT_CONV, False)`` → :class:`_LFM2MoEShortConvDecoderLayer`
  * ``(_, True)``           → :class:`_LFM2MoEMoEDecoderLayer`

Holds ``model.*`` subtree so safetensors keys line up.

The ``layers`` :class:`nn.ModuleList` carries the **inner** decoder
layer instances — one per spec, in spec order. The LayerStack
adapter wrappers (built by :class:`LFM2MoeForCausalLM.__init__`)
hold back-references to these same ``inner`` modules rather than
fresh copies, so PyTorch's ``named_parameters`` walk reaches each
parameter via exactly one path (``self.model.layers.{i}.*``) and
safetensors weight-map keys resolve as expected. The adapters
live outside ``self.model`` in a plain Python list to avoid
double-registration.

LFM2 hybrid + sparse MoE FFN for layers >= num_dense_layers.

The MoE FFN is the fused sparse block
(:class:`~arbi_serve.models.moe.FusedSparseMoEBlock`): sync-free,
fixed-shape dispatch, so decode / piecewise capture and EP compose
like any other arch.

MoE FFN: router gate + (optional) bias + stacked routed experts.

The routed experts bind through the arch-agnostic
:func:`~arbi_serve.models._moe_weight_map.add_stacked_moe_expert_keys`
straight into the fused block's ``w13_weight`` / ``w2_weight``
stacks. LFM2-shaped here: the per-expert linears are named
``w1`` (gate) / ``w3`` (up) / ``w2`` (down), the router key is
``feed_forward.gate``, and HF stores the score bias as
``feed_forward.expert_bias`` (no ``.gate.`` qualifier) while the
fused block holds it as ``experts.e_score_correction_bias``.

Produce ``(ModelDims, layer_specs)`` for LFM2-MoE.

Reuses :func:`build_lfm2_layer_specs` for the attention / conv
interleave and for ``ModelDims``, then attaches a
:class:`MoEConfig` to layers ``[num_dense_layers,
num_hidden_layers)`` so the FFN factory builds the fused MoE block
instead of a dense MLP.

Linear-layer factory: replicated + tensor-parallel variants.

A model declares its projections via the four classes here:
:class:`ReplicatedLinear` (no comm), :class:`ColumnParallelLinear`
(output dim sharded across the TP group; forward returns the local
shard), :class:`RowParallelLinear` (input dim sharded; forward
all-reduces the partial output), and :class:`MergedColumnParallelLinear`
(fuses MLP gate+up via a per-loader ``shard_id`` selector).

Every parallel linear short-circuits to :class:`ReplicatedLinear`'s
math when its shard axis has size 1: no NCCL handshake, no extra view,
no all-reduce.

Each parallel linear is cut by, and reduces / gathers over, ONE named
axis — ``shard_axis``, defaulting to ``"tp"`` (the whole TP group).
``"attn_tp"`` selects the attention-TP set, which under
``attn_dp_size > 1`` is a strict subset of the TP group: an attention
projection's partial must be reduced over exactly the ranks that share
its tokens, because the remaining TP ranks are computing different
requests entirely. The axis fixes BOTH halves together — the weight
slice and the collective's group — so the two can never disagree.

Adapter dispatch. Each parallel linear accepts an optional
``lora_state`` kwarg. At construction the model binds a
``lora_target_key``; the post-matmul forward calls
:func:`arbi_serve.adapters.adapter.apply_adapter` which short-
circuits when there is no adapter, no key, or no per-layer weight,
and otherwise dispatches to the adapter state's
:meth:`apply_to_output`. LoRA's
:class:`arbi_serve.adapters.lora.state.LoraBatchState` is the first
:class:`Adapter` impl; future siblings (IA3, prefix-tuning, BitFit,
vector-DPO) plug into the same dispatcher without per-linear edits.

For a LoRA :class:`AdapterState`, :meth:`apply_to_output` runs::

    y = x W^T + Σ_k 1[lora_id == k] * (α_k / r_k) * (x A_k^T) B_k^T

where ``A_k``, ``B_k`` are the adapter's per-target weights. The
BGMV kernel is launched once per linear forward, even when the batch
carries mixed adapters (Punica / BGMV semantics).

Common surface for the linear-layer family.

Subclasses are :class:`nn.Module` with a ``weight: nn.Parameter`` of
a *local* shape (already sharded by ``tp_rank`` for parallel
variants) and a ``weight_loader`` hook the loader calls with
the full safetensors tensor + an optional ``shard_id`` qualifier
for fused linears.

Adapter target keying. Each parallel linear holds an optional
``lora_target_key`` set by the model — typically
``f"layers.{i}.{module_name}"`` (e.g. ``"layers.5.q_proj"``). When
``None`` the layer is adapter-opaque and ignores any passed
``lora_state`` / ``adapter_state``. Models opt in per-linear via
:meth:`set_lora_target_name`. Semantically this key is the generic
"adapter layer key" the dispatcher routes against.

Linear, no parallelism. Bit-equivalent to ``nn.Linear``.

Used for the input-and-output-replicated path: lm_head / embed (when
not row-sharded), router weights inside MoE, sink biases. Optional
``bias`` rides inside the single ``F.linear`` call so the numerics
match the ``nn.Linear`` reference exactly.

Output dim sharded across TP group.

Forward returns the local (out_features // tp_size) shard. Used for
Q / K / V projections feeding into a row-parallel ``o_proj`` that
all-reduces at the end. ``gather_output=True`` issues an all-gather
after forward — used when the consumer expects the full tensor
(e.g. lm_head when sharded by output dim).

Optional ``bias`` is sharded along the output dim exactly like the
weight (each rank owns its shard's bias slice), applied inside the
local ``F.linear`` — before the optional all-gather, so the gathered
output composes to the full biased projection.

TP=1 fast path: ``self.weight`` is the full (out_f, in_f) tensor and
forward is identical to :class:`ReplicatedLinear`.

Untied lm_head: vocab (output) dim sharded across the attention-TP set.

Like :class:`ColumnParallelLinear` with ``gather_output=True`` but
vocab-aware: the output dim is padded up to a multiple of tp_size so
every rank holds the same shard shape, the padded logit columns are
trimmed back to the real (org) vocab after the all-gather so they
never leak into sampling, and the weight loader places the
checkpoint's ``org_vocab`` rows into the first rows of the padded
per-shard tensor (zeroing the tail of a short last shard). For a
tp-divisible vocab padded == org and this is bit-identical to a plain
``ColumnParallelLinear(gather_output=True)``.

The vocab trim lives here (not in the generic
:class:`ColumnParallelLinear`, which is used for many non-vocab
projections that must NOT be trimmed).

Input dim sharded across TP group.

Forward runs the local matmul on the sharded input then all-reduces
the partial output. Used as the "exit" linear of attention / MLP /
MoE: the upstream column-parallel linears feed each TP rank a slice
of the inner dimension; this layer reduces them back to one.

``input_is_parallel=True`` (default) means the caller is already
handing in a per-rank shard. ``input_is_parallel=False`` would
scatter, which we don't currently need.

Optional ``bias`` is replicated (never sharded — it applies to the
full output) and added exactly ONCE: after the all-reduce at TP>1
(each rank's matmul is a partial sum over the sharded input dim, so
adding it per-rank would scale it by tp_size), fused into the single
``F.linear`` on the TP=1 fast path.

TP=1 fast path: identical to :class:`ReplicatedLinear`.

Fuses multiple column-parallel projections that share the input.

Used for the MLP ``gate_up_proj`` (concatenation of ``gate_proj`` +
``up_proj``). Halves the per-step matmul launch overhead AND halves
the load-time copy count by routing both safetensors tensors into
the same parameter via ``shard_id`` ∈ ``{"gate", "up"}``.

Layout: weight is ``(sum(out_features_per_shard) // tp_size,
in_features)`` with shards stacked along dim 0 in declared order.

Run ``module(x)``, landing the result in ``out`` when one is given.

The lm_head epilogue's single dispatch point. ``out is None`` is the
ordinary path (live decode, prefill, every eager forward): allocate and
return, exactly as before. ``out`` given is the CAPTURE path: the
caller owns a persistent full-vocab buffer and wants the head's GEMM to
write into it, because an in-graph ``(B, vocab)`` allocation is
retained for the lifetime of every captured shape and full-vocab rows
are the largest single term in the cudagraph capture pool.

Dispatch is duck-typed on ``forward_into`` so it also covers heads that
are not :class:`LinearBase` (:class:`~arbi_serve.models.layers.TiedLMHead`)
and degrades to a plain copy for heads that are neither (a quant-backend
swap installs its own module) — those keep the intermediate, but still
land in the caller's buffer, so the shared-buffer half of the win holds
for every arch and every quant backend.

Apply a head's declared logit epilogue to a raw ``lm_head`` projection.

The three terms trained heads compose after the vocab projection, in the
order they compose them: ``scale`` (granite ``logits_scaling``) divides,
``multiplier`` (``output_multiplier``) multiplies, and ``softcap``
(``final_logit_softcapping``) applies ``tanh(x / cap) * cap``.

ONE spelling of the arithmetic, so a caller that runs the vocab projection
ITSELF reaches the same result as one that goes through
:meth:`~arbi_serve.models._layer_stack_model_mixin.LayerStackModelMixin._lm_head`.
Two callers run it themselves: the sharded verify's rank-local shard GEMM,
which has no ``_lm_head`` to route through and must still produce the
values the gathered path would have, and the DFlash drafter's vocab
projection, which calls the bare head module directly.

Callers on a hot path resolve "declares any of them" once off the config
and skip this call entirely; the guards here make the epilogue safe to
call unconditionally, not free to reach.

Every term is skipped when it is identity, so a head declaring none
returns ``logits`` unchanged and issues no kernel.

``in_place`` mutates ``logits`` rather than allocating, for the caller that
owns a destination buffer (the capture-time ``batch.logits_out``) and must
not re-introduce the vocab-sized allocation that buffer exists to avoid.
The arithmetic — and hence the result — is the same either way.

Bind an adapter-target key — the engine threads it through
the per-step :class:`AdapterState` (today: :class:`LoraBatchState`)
to look up per-layer adapter weights for this linear. Pass
``None`` to mark the layer as adapter-opaque.

Semantically the key is the generic "adapter layer key" the
dispatcher routes by.

Generic adapter dispatcher. Replaces the per-linear LoRA hook.

Short-circuits to ``y`` when there's no adapter on the batch
OR the linear is adapter-opaque OR no per-layer weight exists
for this linear's key. Otherwise dispatches to the adapter
state's :meth:`AdapterState.apply_to_output`. The kwarg name
``lora_state`` is preserved for back-compat — it accepts any
:class:`arbi_serve.adapters.adapter.AdapterState` (LoRA today,
IA3 / prefix-tuning / BitFit / vector-DPO tomorrow).

:meth:`forward` writing its result into ``out``; returns ``out``.

Base implementation is the honest fallback — run ``forward`` and
copy. Subclasses whose matmul can name its destination override
this to skip the intermediate entirely (see
:meth:`ReplicatedLinear.forward_into`). Only the lm_head epilogue
calls this today, to keep the captured decode graph's full-vocab
logits out of the cudagraph capture pool; see
:func:`linear_into`.

Replicated matmul written straight into ``out`` — no intermediate.

Bit-identical to :meth:`forward`: ``invariant_linear(out=...)``
dispatches to the same ``mm``/``addmm`` ``F.linear`` does. The LoRA
correction may return a fresh tensor (the adapter kernels are not
destination-passing); copy it back so the contract "``out`` holds
the result" always holds.

Local output-shard matmul written straight into ``out``.

Only fuses when this rank's local matmul IS the whole answer:
``out``'s width must equal the local shard width, which rules out
the ``gather_output`` TP>1 case (the all-gather allocates its own
destination, so there is nothing to fuse into) and the
vocab-padded head whose caller-visible width is the trimmed org
vocab. Both fall back to the base copy — correct, just unfused.
At TP=1 ``self.weight`` is the full matrix and this is the
:class:`ReplicatedLinear` path.

Generic Llama dense model — ``LlamaForCausalLM``.

Covers the vanilla Llama topology used by Llama 2/3, TinyLlama, and the
many ``model_type: "llama"`` derivatives (e.g. MiniCPM5-1B, which is a
plain Llama checkpoint with no muP scalars). It is **architecture-
generic**: any checkpoint whose ``config.json`` declares
``architectures: ["LlamaForCausalLM"]`` and ``model_type: "llama"``
loads here, no per-derivative module.

Topology mirrors the official ``LlamaForCausalLM``:
  Embedding → [PreNorm → AttentionBlock + residual → PreNorm →
              GatedSiLUMLP + residual] × N → final RMSNorm → lm_head.

This is exactly the Qwen3 dense path **minus** the per-head Q/K RMSNorm
(Llama has none) — so it reuses the shared :class:`AttentionBlock`
(``use_qk_norm=False``), :class:`GatedSiLUMLP`, :class:`RMSNorm`, and the
:class:`LayerStackModelMixin` construction/forward boilerplate. No
attention bias (Llama is bias-free, like Qwen3 / Granite).

Deliberately fail-loud on features this generic path has not been
validated against rather than silently mis-serving:
  * ``rope_scaling`` (linear / dynamic / llama3 / yarn) — raises until a
    scaling variant is wired + tested. MiniCPM5-1B has ``rope_scaling:
    null`` so it is unaffected.
  * attention bias (some community Llama forks) — raises.

Generic Llama dense shape config.

All shape fields + the HF parse / ``ModelDims`` / ``LayerSpec``
projection live in :class:`DenseDecoderConfig`. Llama only declares
the ``model_type`` it accepts and fails loud on the two knobs the
generic path hasn't been validated against (rope scaling, attention
bias) so a mis-served model surfaces at boot, not as quiet drift.
Uniform full-context attention — no sliding window.

One Llama transformer block: pre-norm + attn + pre-norm + MLP.

Identical to :class:`Qwen3DecoderLayer` except the attention block
carries no per-head Q/K RMSNorm (``use_qk_norm=False`` — Llama has
none). Implements the uniform :class:`LayerStack` block signature
directly and delegates the residual skeleton to
:class:`PreNormDecoderLayer`.

Mamba-2 mixer — NemotronH (and future Mamba-2 hybrids).

Three axes define the shape:

  * **Multi-head SSM state.** ``num_heads`` heads of width
    ``head_dim``, giving a state of shape ``(num_heads, head_dim,
    state_dim)`` per request.
  * **Grouped B / C.** ``n_groups`` independent (B, C) state channels —
    each group is shared across ``num_heads // n_groups`` heads, much
    like GQA on attention.
  * **Chunk-scan kernel.** Prefill uses
    :func:`mamba_chunk_scan_combined` (varlen via ``cu_seqlens``).
    Decode uses :func:`selective_state_update` with the multi-head
    state shape.

Weight layout (mirrors HF NemotronHMamba2Mixer):

  * ``in_proj``     — ``hidden → [gate, hidden_states_B_C, dt]`` of
    widths ``[intermediate, conv_dim, num_heads]``, where
    ``conv_dim = intermediate + 2 * n_groups * state_dim``.
  * ``conv1d``      — depthwise causal 1-D conv on
    ``hidden_states_B_C`` with kernel ``conv_kernel``; weight shape
    ``(conv_dim, 1, conv_kernel)``, bias ``(conv_dim,)``. SiLU
    activation baked in via the kernel's ``activation="silu"`` arg.
  * ``A_log``       — ``(num_heads,)`` fp32 SSM decay parameter;
    ``A = -exp(A_log)``.
  * ``D``           — ``(num_heads,)`` fp32 input-skip term.
  * ``dt_bias``     — ``(num_heads,)`` softplus pre-bias for dt.
  * ``norm``        — Zamba2-style gated RMSNorm with group_size
    ``intermediate // n_groups``; weight shape
    ``(intermediate,)``.
  * ``out_proj``    — ``intermediate → hidden_size``.

The CPU pure-PyTorch reference path is shape-faithful (per-token
selective-state update for decode) so unit tests stay GPU-free.
The GPU prefill path packs varlen into ``(1, T_total, ...)`` and
calls ``mamba_chunk_scan_combined``.

A prefill row arriving with ``past > 0`` — a second-or-later
``chunk_prefill`` chunk of one prompt, or the duplex lane's ``1 + N``
context-injection row — seeds the conv window and the SSM state from
that row's own recurrent slab entry and scatters the post-chunk state
back, exactly as :class:`~arbi_serve.models.short_conv_block.ShortConvBlock`
and :mod:`~arbi_serve.models._gdn_fla` already did for their own
recurrences. Both upstream kernels take the seed as a first-class
argument (``causal_conv1d``'s ``initial_states``,
``mamba_chunk_scan_combined``'s ``initial_states``), so nothing here is
a workaround: it is the seam the scheduler's chunked-recurrent-prefill
contract (``Scheduler.schedule``'s "the state threads chunk-to-chunk
through the EXISTING GDN/Mamba slab seam") assumes.

RMSNormGated with per-group variance (Zamba2 / NemotronH form).

Matches HF :class:`Zamba2RMSNormGated`: ``y = silu(gate) * x``
grouped into ``hidden_size // group_size`` groups, then variance
is computed per-group; final scale by ``self.weight``.

The weight has shape ``(hidden_size,)`` — full intermediate-size
width, not per-group.

One Mamba-2 SSM block.

Forward signature matches the other recurrent blocks so the
model's per-layer dispatch can call any of them uniformly:

  ``forward(hidden, state_view, meta) -> (N_tokens, hidden_size)``

Per-request state lives in
:class:`RecurrentStatePool.Mamba2LayerView`. The MTP verify forward
and both partial-accept rollback modes come from
:class:`~arbi_serve.models._mamba2_verify_rollback._Mamba2VerifyRollbackMixin`.

Width of ``dt`` in this rank's in_proj output.

A backend that replicates a sub-block instead of sharding it reports
the sub-block's full width here; the even-split default applies to
every backend that does not.

Split in_proj output into ``(gate, hidden_B_C, dt)``.

Sizes ``[intermediate_local, conv_dim_local, dt_local]``, where
``dt_local`` is normally ``num_heads_local`` but is the FULL
``num_heads`` when the backend had to replicate the ``dt`` sub-block
rather than shard it (``dt`` is one element per head, too narrow for
EXL3's 128-wide Hadamard block). This rank's heads are then sliced
out of the replicated stripe.

Run the Mamba-2 mixer over one step (prefill or decode).

Dispatch:
  * MTP verify (``meta.verify_pass``): the K+1 tokens per row go
    through the PER-TOKEN decode kernels — see
    :meth:`_forward_verify`. The chunk scan writes only the
    post-all-K+1 state, so partial-accept rollback would have
    nothing to return to.
  * GPU + kernels: per-request-one-token decode →
    ``causal_conv1d_update`` + ``selective_state_update``.
    Multi-token prefill → ``causal_conv1d_fn`` +
    ``mamba_chunk_scan_combined``.
  * CPU / kernels missing: pure-PyTorch reference for both
    decode and prefill (slow but numerically faithful).

One per-token Mamba-2 SSM step, batched over ``P`` rows.

Returns ``(y, new_ssm_state)``. Math (matches HF torch_forward
use_precomputed_state branch):

  dt_softplus = softplus(dt + dt_bias)              # (P, H)
  dA = exp(dt_b * A)
  B_per_head, C_per_head = repeat n_groups → num_heads
  dB = dt_b * B_per_head[..., None, :]
  dBx = dB * x[..., :, None]
  ssm_new = ssm_state * dA + dBx
  y = (ssm_new * C_per_head[..., None, :]).sum(-1) + D[:, None] * x

Multi-token prefill — pure-PyTorch reference.

Walks each request as a per-token loop over the SSM. Slow but
numerically faithful — used by CPU shape tests.

**Continued prefill works here.** The loop seeds ``conv_state`` /
``ssm_state`` from the row's own slab entry before its first token
and writes the post-chunk state back after its last, which IS the
chunk-to-chunk continuation. Do NOT add a ``torch._check(past <= 0)``
here to mirror a kernel-path restriction: this path does not have
that restriction, and on a box without the ``mamba-ssm`` /
``causal-conv1d`` wheels (every bare-metal dev box; the Docker image
builds them) this is the path a chunked seed prompt runs on. A raise
inside the step reaches ``run_forever``'s error branch, which sets
``finish_reason="error"`` for EVERY row in the slate — a duplex
connection goes silently dark with no error on the wire (design doc
§7.43.1).

A row's seed is masked by ``meta.has_initial_state`` — the
engine's precise ``prompt_consumed > 0`` mask on persistent
prefill steps, ``seq_lens > 1`` otherwise — matching
:class:`ShortConvBlock` / GDN. It is belt-and-braces rather than
load-bearing: a first chunk's slab row is zero-cleared by
``alloc_for_request``, so seeding from it unmasked is already
the fresh-row identity.

Bucketing: ``state_indices`` is per-SEQ (shape (B,)); per-token
bucketing comes from ``cu_seqlens_q`` (shape (B+1,)). When
``cu_seqlens_q`` is missing (e.g. early CPU smoke tests with one
request), we synthesize ``[0, n_tokens]``.

Compile status: the per-row ``T == 0`` skip below is Dynamo-guard
safe (``guard_or_false`` — see its own comment). The per-TOKEN
loop further down (``for t in range(T)``) is NOT: ``T``
is an unbacked SymInt under ``torch.compile`` (every operand
traces back to a ``.item()`` read), and ``range()`` needs a
concrete trip count Dynamo cannot extract from it — verified live
(``torch._dynamo.exc.UserError: Could not extract specialized
integer from data-dependent expression``, thrown from this exact
``for`` line) for EVERY multi-token call through this path, not
only ``T == 0`` ones. This is why boxes that fall back to this
reference path for real traffic (GPU kernels unavailable — see
``_HAVE_CAUSAL_CONV1D`` / ``_HAVE_MAMBA_SSM2`` above) still need
``ARBI_COMPILE_OFF=1`` / ``--no-cuda-graphs`` for anything that
prefills through it, even after the ``T == 0`` fix. Closing this
the rest of the way needs the per-token recurrence rewritten
around a data-dependent-trip-count-safe construct (e.g.
``torch._higher_order_ops.while_loop``/``scan``) rather than a
plain Python ``for`` — out of scope here; this reference path's
real job is CPU-shape-test correctness, not a compiled hot path
(the compiled hot path is the kernel decode/prefill methods
below, which have no such loop).

Decode-step kernel path (one token per request).

Uses ``causal_conv1d_update`` (SiLU-baked rolling conv) and
``selective_state_update`` (multi-head SSM step, takes
``(B, num_heads, head_dim, state_dim)`` state). Both kernels
accept a ``conv_state_indices`` / ``state_batch_indices``
argument that gathers/scatters directly into the underlying
slab — no host-side stack/unstack round-trip.

``(A, D, dt_bias)`` broadcast to :func:`selective_state_update`'s
multi-head signature: ``A (H, head_dim, dstate)``, ``D (H, head_dim)``,
``dt_bias (H, head_dim)``.

ONE source for both per-token ``selective_state_update`` call sites —
the decode / verify forward (:meth:`_decode_step_kernels`) and the
partial-accept replay (:meth:`_replay_masked_steps_kernel`) — so the
replay cannot drift from the forward through a differently-built
operand.

One decode step through the per-token kernels.

Advances BOTH slabs in place on the rows ``state_indices``;
returns the post-conv stream ``(B, conv_dim_local)`` and ``y``
``(B, num_heads_local, head_dim)``. Shared by the decode forward
and the MTP verify forward's per-token loop, so the two are the
same kernel pair by construction.

Multi-token prefill kernel path.

Uses ``causal_conv1d_fn`` (per-request) for the conv stream
and ``mamba_chunk_scan_combined`` for the SSM. Each row is run
as its own batch-1 call, seeded with that row's own carried
state, so continued prefill and fresh prefill can share a step.
(``mamba_chunk_scan_combined``'s VARLEN path does not take an
initial state, but this loop never uses it: it calls the kernel
once per row with ``cu_seqlens=None``, and THAT signature takes
``initial_states``.)

Bucketing: ``state_indices`` is per-SEQ (shape (B,));
per-token bucketing comes from ``cu_seqlens_q``.

Capture-safe at B=1 (the piecewise prefill capture bucket).
Per-row ``int(state_indices[i].item())`` host syncs and
``int(cu[i].item())`` boundary reads are replaced with
``index_select`` / ``index_copy_`` on the persistent
``meta.state_indices`` buffer the engine pre-binds via
:class:`PiecewiseBuffers.recurrent_state_indices`. B>1 paths
miss the captured pool and fall through to the eager varlen
loop unchanged.

The prefill fold: conv, ONE scan launch, and the savepoint emit.

The Mamba-2 twin of
:meth:`~arbi_serve.models.gdn_block.GDNBlock._fold_emit_delta_rule`,
wrapping :meth:`_mamba2_kernels` so a savepoint is staged by the one
code path both prefill shapes take.

WHAT THE PLAN IS AND WHERE IT COMES FROM
========================================
Nothing here is read from the metadata or from ``self``: the step's
plan is published host-side before the forward runs
(:mod:`arbi_serve.cache._fold_emit_staging`) and looked up by the
identity of ``cu`` — the boundary tensor this launch was handed. No
plan, or another tensor's plan: the plain launch.

WHAT "THE LAUNCH ROW" MEANS HERE
================================
GDN issues one varlen launch for the whole batch, so its plan's row
indexes ``cu_seqlens`` and the kernel finds it by ``i_n``. The SSD
scan takes no ``cu_seqlens``, so this block issues one batch-1 launch
PER ROW (:meth:`_forward_kernel_prefill`) and the row is which
iteration is running — a host integer, free to compare. The plan's
row and this launch's must be equal, and this launch's token count
must be that row's segment of the plan's boundaries; GDN's whole-
launch token count would pass on the wrong row here. The emit table
the kernel reads is correspondingly this row's single entry of the
step's table, sliced, never a value read out of it.

WHY THE SAME BITS
=================
The scan carries each sequence's state in fp32 across its chunks and
stores it in fp32 at the end; the emit is that same register at an
earlier chunk, in the slab's own layout. Nothing about the scan's
blocking, order or rounding changes, so the output and the final
state are the plain call's and the emitted state is what a scan cut
at that chunk would have ended with. ``offset`` is a multiple of the
fold grid by construction — ``savepoint_fold_split`` refuses rather
than rounds — so the chunk index it names is exact. Measured in
``tools/mamba_fold_emit/check_emit_state.py``.

BEFORE IT EMITS, THE LAUNCH PROVES IT IS THE PLAN'S
===================================================
The slab row it will scatter into must be the row the snapshot's host
buffers were laid out for, compared by device address, no sync; where
the row is already a host integer (the B>1 loop resolves it) it is
compared directly too. Every mismatch refuses by NAME on
``gdn_fold_split_staged`` and stages nothing; the commit then refuses
the snapshot whole. A partial snapshot is never stored.

The conv slab row a prefill cut at the boundary would have written.

Mamba's rolling conv buffer is ``(D, K)`` with slot 0 OUT of the
carry: a chunk's post-state is the ``K-1`` pre-conv inputs before the
boundary in slots ``1..``, and slot 0 is zero — the next token's
``causal_conv1d_update`` shifts it out unread, so a prefill leaves it
zero and a snapshot that filled it would not be the row a cut
produced. So the gather is the staging's index table minus its oldest
entry, and the oldest column is zeroed rather than filled.

Index-driven and unconditional, like GDN's: the same kernels whether
or not a plan is armed, so the compiled and the captured graph are
the same either way.

Run conv + chunk-scan over one varlen chunk.

``conv_init`` / ``ssm_init`` are this row's carried recurrent
state — ``None`` for a fresh first chunk, the previous chunk's
post-state for a continued one.

The SSM carry is the kernel's own ``initial_states`` argument.
The CONV carry is not, and cannot be here: this block feeds
``causal_conv1d`` a ROW-MAJOR ``x`` on purpose (its channel-last
path hits a misaligned-address error at large ``D`` — 9728 on
NemotronH), and the extension refuses that combination outright
(``RuntimeError: initial_states is only supported for channel
last layout`` — measured, not assumed). So the carry is applied
by PREPENDING the K-1 remembered inputs to the chunk and
discarding their outputs. For a causal depthwise conv of width K
that is not an approximation of seeding, it IS seeding: every
output the chunk keeps sees exactly the K-1 real predecessors the
kernel's own ``initial_states`` would have supplied. It costs
K-1 extra columns (3 on NemotronH) and keeps the layout that
dodges the alignment bug.

Returns ``(y_flat, final_ssm, final_conv)``:
  * ``y_flat``: ``(T, intermediate_local)``.
  * ``final_ssm``: ``(1, num_heads_local, head_dim, state_dim)``.
  * ``final_conv``: ``(1, conv_dim_local, conv_kernel - 1)``.

B=1 prefill — capture-safe vectorized path.

Whole hidden_states is the single request's chunk. Slab gather
via ``index_select`` / scatter via ``index_copy_`` on the
persistent ``meta.state_indices`` buffer.

The gather is what makes a CONTINUED chunk correct here. Passing no
seed and then overwriting the slab with the chunk's own final state
is not a loud failure on this path: at B=1 (the duplex lane's every
step, and the piecewise prefill capture bucket) a second chunk would
restart the recurrence from zero, throw away everything the prompt
said before it, and return plausible garbage.

MLA (Multi-head Latent Attention) block — DeepSeek V2 / V3, Kimi K2.

DeepSeek's MLA layer is a low-rank-factored attention with a separate
RoPE tail:

    Q     = q_a_proj(hidden_size → q_lora_rank)
    Q     = q_a_layernorm(Q)
    Q     = q_b_proj(q_lora_rank → num_heads × (qk_nope + qk_rope))

    KV    = kv_a_proj_with_mqa(hidden_size → kv_lora_rank + qk_rope)
    KV_c  = kv_a_layernorm(KV[..., :kv_lora_rank])
    KV    = kv_b_proj(KV_c → num_heads × (qk_nope + v_head))   ← used for K_nope, V

    K_pe  = RoPE(KV[..., kv_lora_rank:]) — single shared head, broadcast
    Q_pe  = RoPE(Q[..., qk_nope:])

The cache stores one (kv_lora_rank + qk_rope_head_dim) vector per token
per layer (the "shared" KV); on attend, V is reconstructed by re-applying
``W_UV`` (the V-half of ``kv_b_proj``) to the cached compressed vector.
Decode ABSORBS the factors — ``W_UK`` folds into the query and ``W_UV``
unfolds the latent result — so nothing materialises a bf16 mirror of the
cache.

When ``q_lora_rank is None`` (DeepSeek V2-Lite has it null), the Q
low-rank factor is skipped and ``q_proj`` is a single
``hidden_size → num_heads × qk_head_dim`` linear.

One MLA attention block.

Constructed from a :class:`LayerSpec` whose ``mla`` field carries
``MLAConfig(kv_lora_rank, qk_rope_head_dim, qk_nope_head_dim,
v_head_dim, q_lora_rank)``. All projections are
:class:`LinearBase` subtypes cut by the ``"attn_tp"`` shard axis; an
attention-TP set of size 1 short-circuits to plain matmul.

The forward signature matches Qwen3's ``AttentionBlock`` for
drop-in compatibility within ``DeepseekV3Model.forward``: it consumes
a per-layer state view (the MLA paged slab) and the kind-specific
``MLAMeta`` populated by the metadata builder.

Hand the per-layer attn op the dense ``kv_b_proj`` weight for V re-projection.

The MLA backend reconstructs V via ``F.linear(kv_c, W_UV)`` and
therefore needs the dense ``(out, in)`` weight matrix — not a
callable linear. On a dense checkpoint that is just
``kv_b_proj.weight``. On a quantized checkpoint (AWQ V2-Lite),
``kv_b_proj`` was swapped to a quant Linear with no ``.weight``
Parameter; because the layer marked it ``retain_dense_for_readout``
at construction, the quant load captured the dense weight from
the raw int4 buffers (exact torch dequant), exposed here via
:meth:`readout_dense_weight`. (An identity probe through Marlin
is not a valid substitute — Marlin's GEMM does not column-isolate
an identity input, so the probe returns ~0.)

MoE / Expert Parallelism — the fused expert path.

:class:`FusedMoE` is the only production expert dispatch: stacked
``(E, 2N, H)`` / ``(E, H, N)`` expert weights driven by the Triton
fused-MoE kernel (one binned GEMM per stage), sync-free and fixed-shape
— which is what makes an MoE decode step cudagraph-capturable — and
TP/EP-aware. :class:`FusedSparseMoEBlock` wraps it with a replicated
router for arches without extra FFN structure. The per-expert-module
walk exists only as the parity oracle under ``tests/_moe_oracle.py`` —
it must never be importable from production (a reference dispatch one
arch away from being served is a silent several-fold slowdown), and the
one legitimate per-expert-module production case (a weight-quant backend
that cannot hand over a stacked expert payload — see
:meth:`~arbi_serve.weight_quant.base.QuantBackend.stacked_moe_spec`)
lives arch-local and capture-tripwired in
:mod:`arbi_serve.models.deepseek_v3`.

Routing (:func:`_route`)::

    router_logits = gate(hidden)  # (N, n_experts)
    scores = sigmoid_or_softmax(router_logits.float())
    choice = scores + e_score_correction_bias  # selection only
    choice = group_limited_mask(choice)  # n_group > 1
    indices = topk(choice, k=n_active).indices
    weights = scores.gather(-1, indices)  # UNBIASED scores
    if norm_topk_prob:
        weights = weights / (weights.sum(-1, keepdim=True) + 1e-20)
    weights = weights * routed_scaling_factor

``e_score_correction_bias`` biases SELECTION only — the returned weights
are gathered from the unbiased scores.

Routing is per-token, so non-MoE rows in the same engine step pay zero
cost.

Mask ``choice`` down to the ``topk_group`` highest-scoring groups.

Group score is the sum of the top-2 members when an
``e_score_correction_bias`` is in play (DeepSeek-V3 ``noaux_tc``),
the group max otherwise (DeepSeek-V2 ``group_limited_greedy``).
Non-selected groups become ``-inf``, so they can never win the
trailing top-k whatever sign the scores carry.

Compute (router_weights, router_indices) from raw router logits.

``expert_bias`` (``e_score_correction_bias``) is added to the SCORES
— after the sigmoid/softmax, never to the logits — and steers
selection only: the returned weights are gathered from the unbiased
scores. ``n_group > 1`` restricts selection to the ``topk_group``
best expert groups first (DeepSeek-V3).

``forced_indices`` ``(N, top_k)`` REPLACES the top-k selection while
leaving everything downstream intact — the weights are still gathered
from this router's scores at those ids, still normalized, still scaled.
A layer that routes by token id rather than by hidden state
(DeepSeek-V4's hash-routed layers) supplies them; the scores are what
the checkpoint trained the weights against, so gathering them is the
routing, not an approximation of it.

Sync-free and fixed-shape end to end, so an MoE decode step stays
cudagraph-capturable.

Expert activation whose pre-activations are clamped before the gate.

Replaces the default ``silu(gate) * up`` with::

    g = clamp(gate, max=limit)
    u = clamp(up, -limit, limit)
    out = (u + up_shift) * g * sigmoid(alpha * g)

The clamp is part of the TRAINED forward for the architectures that
declare a limit, not a numerical guard: the weights were fit against
a saturating activation, so dropping it changes the served
distribution wherever a pre-activation reaches the limit.

``alpha`` and ``limit`` come from the arch config; nothing here
defaults them. ``up_shift`` is the constant the clamped up-projection
is offset by — ``1.0`` for gpt-oss, ``0.0`` for a plain clamped
SwiGLU (DeepSeek-V4).

True iff ``device`` has native fp8 tensor cores — Ada ``sm_89`` or
Hopper/Blackwell ``sm_90+``.

The fused-MoE Triton kernel's ``use_fp8_w8a8`` path emits a ``tl.dot`` in
``float8e4nv`` (native e4m3), which Ampere (``sm_80`` / ``sm_86``) does not
implement — Triton there rejects the dtype outright (``type fp8e4nv not
supported in this architecture``). Off native-fp8 cards the experts must be
dequantized to bf16 first. Same arch gate as
:func:`arbi_serve.weight_quant.fp8.triton_gemm.block_fp8.is_eligible`.

Triton-fused MoE FFN block — stacked experts, TP- and EP-aware.

API surface::

    moe = FusedMoE(
        hidden_size=H,
        intermediate_size=N,  # per-expert, not 2*N
        num_experts=E,  # global expert count (router width)
        top_k=K,
        scoring_func="softmax",  # or "sigmoid"
        norm_topk_prob=True,
    )
    out = moe(hidden_states, router_logits)  # (T, H) -> (T, H)

Weight tensors held by this module (``E_local = E //
expert_shard_count``, ``N_local = N // moe_tp_size``):

- ``w13_weight``: ``(E_local, 2*N_local, H)`` — gate+up, stacked
  along N. Per-expert convention is ``gate_proj`` first (``[:N]``),
  ``up_proj`` second (``[N:2N]``); the SiLU+mul activation expects
  this order. Loaders must respect it.
- ``w2_weight``:  ``(E_local, H, N_local)`` — down projection.

**Expert parallelism.** The router is replicated (it scores all ``E``
experts on every rank), so each rank sees the same top-k. This rank
holds the contiguous slice ``[expert_shard_rank * E_local, +E_local)``;
ids are shifted into the local frame and ids outside it become ``-1``,
for which the kernel writes zeros.

Experts are sharded along both available axes, and the module reads
the single composed ``expert_shard_rank`` rather than either one:

  * across TP groups (``--ep-size N``, the 2-D world topology);
  * within a TP group (``--enable-expert-parallel`` /
    ``--moe-ep-size N``): attention shards TP-wise over N ranks and
    the experts shard EP-wise over those same N ranks.

**Tensor parallelism.** The intermediate dim is sharded across the
``moe_tp_size`` ranks that share this rank's expert shard (gate/up
rows, down columns — the fused twin of ``ColumnParallelLinear`` × 2 +
``RowParallelLinear``). At ``moe_ep_size == 1`` that is the whole TP
group; at ``moe_ep_size == tp_size`` it is this rank alone and the
expert stack is not cut at all.

**Combine.** One all-reduce over the TP group and one over the EP
group. Those two groups tile the world, and every summand — expert
subset × intermediate slice — lands on exactly one rank of that tiling
at every ``moe_ep_size``, so the collective sequence is invariant to
the expert-sharding mode. See :meth:`_fused_experts`.

**Quantization.** ``quant_kind="fp8_w8a8"`` carries the fp8_e4m3
payload in the same ``(E, *, *)`` tensors. ``weight_block_size``
selects the scale layout:

- ``None`` — one fp32 scale per expert per stage (``w13_scale`` /
  ``w2_scale`` shaped ``(E_local,)``), activations quantized with a
  single tensor-wide amax.
- ``(bn, bk)`` — DeepSeek-V3 / Qwen3.6 native block-wise fp8: one
  scale per ``bn x bk`` weight tile, so ``w13_scale`` is
  ``(E_local, 2*N_local/bn, H/bk)`` and ``w2_scale`` is
  ``(E_local, H/bn, N_local/bk)``. The stored value is the
  checkpoint's ``weight_scale_inv``, applied by multiplication —
  the same convention as
  :func:`arbi_serve.weight_quant.fp8.loader.dequantize_fp8` and the
  dense ``arbi_fp8::block_fp8_mm`` GEMM. Activations are quantized
  per ``(token, bk)`` group.

``quant_kind="mxfp4"`` carries OCP MXFP4 expert stacks: ``w13_weight``
is ``(E_local, 2*N_local, H // 2)`` uint8 (two E2M1 nibbles per byte,
LOW nibble = even K) and ``w13_scale`` is ``(E_local, 2*N_local,
H // 32)`` uint8 (E8M0 biased exponents); ``w2_weight`` /``w2_scale``
are ``(E_local, H, N_local // 2)`` / ``(E_local, H, N_local // 32)``.
Activations stay bf16. The TP cut of the intermediate dim must land on
whole 32-element scale blocks; the expert (EP) axis is unconstrained.

``quant_kind="nvfp4"`` carries modelopt NVFP4 expert stacks: the same
nibble-packed ``w13_weight`` / ``w2_weight`` the MXFP4 kind uses, with
a 16-element scale block whose factor is fp8_e4m3 (``w13_scale``
``(E_local, gu*N_local, H // 16)``) and a per-tensor fp32
``weight_scale_2`` per source linear (``w13_gscale`` ``(E_local, gu)``,
``w2_gscale`` ``(E_local, 1)``). Activations stay bf16. The TP cut of
the intermediate dim must land on whole 16-element scale blocks.

``quant_kind="int4_w4a16"`` carries an AWQ/GPTQ int4 expert stack in
the layout ``fused_moe_kernel_gptq_awq`` reads, with
``quant_group_size`` (``gs``) the checkpoint's group along the
reduction axis. ``w13_weight`` is ``(E_local, 2*N_local, H // 2)``
uint8 (two int4 per byte along K, LOW nibble = even K),
``w13_scale`` is ``(E_local, 2*N_local, H // gs)`` and
``w13_qzeros`` is ``(E_local, N_local, H // gs)`` uint8 (two
zero-points per byte along N, LOW nibble = even N); ``w2_weight`` /
``w2_scale`` / ``w2_qzeros`` are ``(E_local, H, N_local // 2)`` /
``(E_local, H, N_local // gs)`` / ``(E_local, H // 2, N_local // gs)``.
Activations stay at the engine dtype. Both reduction axes (``H`` and
``N_local``) must be multiples of ``gs``.

**Per-expert FFN bias.** ``use_ffn_bias`` adds ``w13_bias``
``(E_local, 2*N_local)`` and ``w2_bias`` ``(E_local, H)``, which the
kernel adds after dequantization and before the router-weight multiply.
``w13_bias`` rides the moe_tp row cut of ``w13_weight``; ``w2_bias``
applies to the FULL output, so at ``moe_tp_size > 1`` only
``moe_tp_rank == 0`` contributes it — the trailing all-reduce would
otherwise sum it once per rank. Distinct from ``use_expert_bias``,
which is the ROUTER correction added to the logits before top-k.

``activation`` replaces the default ``silu(gate) * up`` when set.

Router + :class:`FusedMoE` routed experts — the generic sparse FFN.

The block a new MoE arch drops in as its FFN module: it supplies a
:class:`~arbi_serve.models.layer_spec.MoEConfig` plus routing knobs
and gets the fused sync-free dispatch — and with it cudagraph
capture, EP, and TP composition — for free. Holds:

  - ``gate``: replicated router linear (hidden → n_experts; scores
    all experts on every rank so EP ranks agree on the top-k);
  - ``experts``: one :class:`FusedMoE` carrying this rank's stacked
    local expert slice (bind via
    :func:`~arbi_serve.models._moe_weight_map.add_stacked_moe_expert_keys`)
    and, when ``use_expert_bias``, the ``e_score_correction_bias``
    parameter.

Arch families with extra FFN structure (Qwen-MoE's sigmoid-gated
shared expert) wrap :class:`FusedMoE` themselves; this block refuses
a ``shared_experts`` count rather than silently dropping the term.

Give the expert stacks real storage on ``device``.

Arch ``from_safetensors`` calls this after the quant swap and the
meta-materialize walk — that walk only retypes PARAMETERS, and
these are buffers (see :meth:`__init__`). Returns bytes allocated.

Three states arrive here and each is left correct:

  * meta (the normal cold build) → allocate the full stack;
  * zero-element placeholder (``skip_weight_load``: the warm
    flat-dump fill or a donor share supplies both the size and the
    bytes) → left alone;
  * already real (a donor share aliased the donor's storage) →
    left alone, so the alias is never broken.

Run gate -> top_k -> fused gate/up GEMM -> SiLU+mul ->
fused down GEMM -> weighted sum over top_k.

``hidden_states``: ``(T, H)`` (caller flattens batch+seq).
``router_logits``: ``(T, E)`` raw scores from the router head.

``router_indices`` ``(T, top_k)`` bypasses the top-k selection with
ids the caller already knows (a hash-routed layer reads them off the
token id). The weights still come from THIS router's scores at those
ids, so the selection changes and nothing else does.

``overlap`` is an optional zero-arg callable producing a ``(T, H)``
term to add to the routed mixture — an always-on shared expert is
the canonical one. It must depend only on this block's input, never
on the routed output, because it is invoked while the EP all-reduce
of the routed output is in flight (see
:meth:`GroupCoordinator.all_reduce_overlapped`). Passing it is
arithmetically identical to adding the term after ``forward``
returns; the only difference is that its compute hides the
collective's latency. Omitted ⇒ nothing changes.

Materialize the fp8 expert stacks as dense bf16, caching them.

The Ampere (``sm<89``) fallback for the fp8 GEMM — mirrors
:meth:`arbi_serve.weight_quant.fp8.linear.FP8Linear._ensure_dequantized`
for the 3-D expert stacks. :func:`dequantize_fp8` is per-matrix, so
each local expert's ``(rows, cols)`` slice is dequantized with its own
scale (per-tensor ``(n_local,)`` or block-wise ``(n_local, .., ..)``).

One-shot at the first forward; the cached bf16 buffers have stable
``data_ptr``s, so a later cudagraph capture bakes in valid addresses
(the loader binds the fp8 weights before any forward, exactly like the
fp8 linear's lazy dequant). Costs a full bf16 mirror of the expert set
— the memory price of running an fp8 checkpoint on a card that cannot
compute in fp8; native-fp8 cards never allocate it.

Dynamically quantize ``x`` to fp8 for the fused GEMM.

Returns ``(x_fp8, x_scale)`` in the layout the kernel's
``use_fp8_w8a8`` branch reads: a 0-d tensor-wide scale when the
weights carry per-tensor scales, or a ``(M, ceil(K / bk))`` fp32
grid when they carry block-wise scales — matching the dense
block-fp8 GEMM's per-``(token, K-group)`` quantization, which
never lets one large-magnitude row flatten its siblings.

Routed experts as per-expert modules, driven by a grouped GEMM.

The peer of :class:`~arbi_serve.models.moe.FusedMoE`: that one serves packs
whose experts stack into one tensor, this one serves packs whose payload the
fused kernels cannot read — an EXL3 trellis, an AWQ W4A8 pre-scale — and
which therefore have to stay per-linear.

Both legs live here on purpose. Every arch that grew its own copy of this
decision ended up with a hole: refusing every non-EXL3 pack, or reaching a
per-expert dispatch that raises under cudagraph capture. Which leg is live is
only knowable AFTER the quant swap binds each linear's ``_inner``, so the
component — not the arch — is the only thing that can answer
:attr:`capture_safe`.

Widest step the grouped leg can take, given the gate and the budget.

The grouped leg runs whenever :func:`expert_major_fires` is false, so the
answer is ``threshold - 1`` once the gate is armed and ``max_step_rows``
when it is off.

Per-expert modules keyed by GLOBAL expert id, dispatched as one group.

A :class:`~torch.nn.ModuleDict` SUBCLASS rather than a module holding one:
the checkpoint path is ``<mixer>.experts.{gid}.<proj>``, and an extra
attribute level would insert a segment the quant swap's prefix remap
cannot resolve.

Keys are GLOBAL expert ids. Under expert parallelism a LOCAL index would
collide with another shard's expert 0 through that remap's identity
fallback, and the linear would be swapped twice.

``forward`` takes routing already computed by the arch — archs disagree
about routing (noaux_tc vs plain softmax) and agree about dispatch.

Whether this rank's routed dispatch may be cudagraph-captured.

True only once every projection carries a bound EXL3 inner: the
grouped leg is fixed-shape and sync-free, the per-expert fallback
host-syncs on the routed ids and raises under capture. Read by
:func:`~arbi_serve.runtime.capture.preflight._capture_unsafe_component`,
which is the whole-forward capture gate for decode, prefill and
MTP verify alike.

Every routed projection carries a bound EXL3 inner.

``_inner`` is set only by :mod:`arbi_serve.weight_quant.exl3.linear`,
so this is exactly "the ``exl3_mgemm`` legs are reachable"; without it
the forward takes the per-expert walk and no mgemm call is made.
Side-effect free — unlike :meth:`_exl3_groups` it caches no pointer.

Rebase routing to this shard, neutralising off-shard slots.

For the legs whose kernel takes only a LOCAL pointer table and no
expert range: ``exl3_moe`` (expert-major) and the per-expert
dispatch. Off-shard slots keep their position with a zeroed weight
and an in-range index, so the dispatch keeps a fixed shape and never
syncs on the routing.

Whether this row count clears the expert-major threshold.

Read per forward rather than cached: the threshold is a boot flag but
``rows`` is the captured bucket, so each captured graph bakes in its
own answer. The threshold itself is resolved from THIS layer's shard
count — the crossover differs between a resident expert set and a
sharded one, and the boot-time slot-cap validator resolves it the same
way, so the two can never disagree about which leg a step takes.

Slot-major dispatch: one ``exl3_mgemm`` slot per (token, expert).

``indices`` are GLOBAL expert ids. Under expert sharding the kernel
takes this rank's ``[min_index, max_index)`` and rebases them itself:
it skips every off-shard slot and leaves that slot's output row at
the zero the op pre-fills, so the trailing routed combine multiplies
a zero row by its unmasked weight. Slot positions are untouched, the
launch shape is the full slot set, and nothing reads the routing on
the host.

Allocate this layer's arena-backed attention output buffer.

Called eagerly from the model's per-layer dispatch seam; the
result replaces the arena in ``extras`` so the compiled block
forward receives a tensor. See
:meth:`~arbi_serve.models._decoder_layer.PreNormDecoderLayer.
alloc_arena_buffers` — this block does not share that skeleton
(four norms, not two) but the contract is the same.

NemotronH — NVIDIA's hybrid Mamba-2 + attention + MLP-only stack.

Per ``config.hybrid_override_pattern`` (mapped to
``layers_block_type``): each layer carries exactly one of Mamba-2,
attention (NoPE GQA, no qk_norm), or a standalone non-gated relu² MLP
(unlike LFM2 where every block has both mixer AND MLP).

Weights live under ``backbone.*`` (not ``model.*``); ``lm_head`` is
untied + bias-free. Per-mixer key layouts:
  Mamba-2:  in_proj / conv1d.{weight,bias} / A_log / D / dt_bias /
            norm.weight / out_proj.weight
  Attention: q_proj / k_proj / v_proj / o_proj  (no q/k_norm)
  MLP-only:  up_proj / down_proj  (no gate; relu² activation)

Per-layer dispatch routes through :class:`LayerStack`; the inner
decoder layers implement the uniform LayerStack block signature
directly. The per-arch dispatcher
(:func:`_make_nemotron_h_decoder_layer`) picks the right class.
MLP-only layers carry :attr:`StateKind.NONE` — the iterator threads
``None`` for both ``batch_meta`` and ``state_view`` and the MLP
layer discards them. MoE-only blocks (Nemotron-3.5-Lightning) take
the same NONE state with a routed FFN in place of the dense MLP.

Decode NemotronH ``hybrid_override_pattern`` into block-type list.

Mapping (per the model's bundled ``configuration_nemotron_h.py``):
  ``M`` → ``"mamba"`` (Mamba-2 mixer)
  ``*`` → ``"attention"`` (NoPE GQA)
  ``-`` → ``"mlp"`` (standalone MLP-only block)

  ``E`` → ``"moe"`` (routed MoE FFN block)

NemotronH-3-Nano-4B has 21 M / 17 - / 4 * across 42 layers;
Nemotron-3.5-Lightning-30B-A3B has 23 M / 23 E / 6 * across 52.

Per-arch construction bundle for :class:`NemotronHModel`.

Collects the ctor args that don't fit on :class:`ModelDims` (which
is shared across all archs) but are model-arch-specific knobs
derived from the HF NemotronH ``config.json``. Following the
mlc-llm pattern, the arch's :meth:`from_safetensors` builds one of
these and the constructor consumes it (rather than threading
several positional / keyword args through the call site).

Args:
    mamba_cfg: per-layer Mamba-2 config — same instance shared
        across every Mamba layer (NemotronH uses one consistent
        Mamba-2 shape across the whole stack).
    head_dim: per-attention-head dimension (NemotronH-3-Nano-4B:
        128). Derived from ``attention_head_dim`` / ``head_dim`` /
        ``hidden_size // num_attention_heads`` in the HF config.
    num_q_heads: pre-TP attention head count
        (``num_attention_heads``).
    num_kv_heads: pre-TP KV head count (``num_key_value_heads``;
        equal to ``num_q_heads`` when GQA is disabled).
    layers_block_type: per-layer block type — one of
        ``"mamba"`` / ``"attention"`` / ``"mlp"``. Decoded from
        ``hybrid_override_pattern`` (or read directly from the HF
        ``layers_block_type`` field when present).

NemotronH MLP — non-gated, relu² activation.

Shape:
  ``up_proj``   ``hidden_size → intermediate_size``
  activation    ``act²`` (relu² per ``mlp_hidden_act='relu2'``)
  ``down_proj`` ``intermediate_size → hidden_size``

No gate (unlike Llama / Qwen / LFM2). Activation is squared ReLU
(``F.relu(x).pow(2)``) — same as the upstream HF NemotronHMLP.

MLP-only block — PreNorm + non-gated relu² MLP + residual.

Implements the uniform :class:`LayerStack` block signature directly.
Unique to NemotronH: a standalone FFN-only block (no token mixer,
no per-layer cache state). State_kind is NONE — both ``batch_meta``
and ``state_view`` arrive as ``None`` and the layer discards them.

Routed MoE mixer — DeepSeek-V3 noaux_tc router over relu² experts.

Each expert is the same non-gated :class:`_NemotronHMLP` the dense
MLP-only block uses. Experts stay per-module so a per-linear quant
backend can bind each one by path.

Resolve ``(fused_experts, quant_kind, quant_group_size)`` for the
routed experts from the pack's CAPABILITY.

A dense pack, or a quantized pack whose backend can serve the experts
stacked, takes the fused path; anything else keeps per-expert modules.
Decided from the checkpoint because the layer class is fixed at
construction, before the quant swap runs. Same rule as
:class:`_DeepseekV2MoE`.

Holds the ``backbone.*`` subtree so safetensors keys line up.

The ``layers`` :class:`nn.ModuleList` carries the **inner** decoder
layer instances — one per spec, in spec order. The LayerStack
adapter wrappers (built by :class:`NemotronHModel.__init__`) hold
back-references to these same ``inner`` modules rather than fresh
copies, so PyTorch's ``named_parameters`` walk reaches each
parameter via exactly one path
(``self.backbone.layers.{i}.*``) and safetensors weight-map keys
resolve unchanged from the pre-migration layout. The adapters
live outside ``self.backbone`` in a plain Python list to avoid
double-registration.

NemotronH causal LM — Mamba-2 + attention + MLP-only hybrid.

The model is text-only (no vision tower); the HF arch class string
is ``NemotronHForCausalLM``. Weights live under the ``backbone.*``
prefix (not ``model.*``) — the loader rebases automatically via
``weight_map``. The lm_head is untied (``tie_word_embeddings=False``).

Bind the Mamba-2 mixer params for layer ``base``.

``base`` is the per-layer suffix (e.g. ``"layers.5"``, no top-level
checkpoint namespace). ``prefix`` supplies that namespace —
``"backbone"`` (default) reproduces NemotronH's existing
``backbone.layers.{i}.*`` keys unchanged; the NemotronLabs-
VoiceChat STT backbone (same Mamba-2 mixer shape, different
checkpoint) passes ``prefix="stt_model.llm"``. Module path (the
``wm`` dict key) and safetensors source key share the same
prefixed base — every NemotronH-family checkpoint keeps dst == src
for this subtree.

TP sharding (supported only when ``n_groups % tp_size == 0`` — see
the caller's guard in ``from_safetensors``): every rank owns a
contiguous, whole-number-of-groups slice of heads, so B and C's
group-major layout slices cleanly along with everything else.

  * ``in_proj``: a SINGLE fused tensor whose output rows stack
    ``[gate(I) | x(I) | B(G*S) | C(G*S) | dt(H)]`` in that exact
    order (``hidden_states_B_C`` = ``[x | B | C]`` per
    :meth:`Mamba2Block._split_xBC`, itself the middle third of
    in_proj's ``[gate | hidden_states_B_C | dt]`` split). A plain
    contiguous TP cut over the WHOLE tensor would land mid-block
    (e.g. rank 0 getting all of gate plus half of x); ``fused_split``
    slices each sub-block independently by its own even TP cut and
    re-concatenates in the same order — bit-identical to
    ``Mamba2Block``'s own ``in_proj``/``_split_projection``/
    ``_split_xBC`` local-shape expectations.
  * ``conv1d.{weight,bias}``: depthwise, channel dim mirrors
    ``hidden_states_B_C`` = ``[x(I) | B(G*S) | C(G*S)]`` exactly
    (same convention as GDN's ``conv1d_weight`` in
    ``_qwen3_5_weights.py`` — see that file for the precedent this
    mirrors). Must use the SAME per-sub-block split as in_proj's
    middle third, not an independent even split over the whole
    conv_dim.
  * ``A_log`` / ``D`` / ``dt_bias``: per-head ``(num_heads,)``
    vectors, not part of in_proj — plain ``shard_dim=0`` (each
    rank's contiguous head range).
  * ``norm.weight``: Zamba2 gated-RMSNorm weight, full
    ``intermediate_size`` width (unlike GDN's per-head-width norm,
    which is TP-invariant by construction) — plain ``shard_dim=0``,
    rank-local BEFORE ``out_proj`` (no cross-rank reduction needed;
    matches ``Mamba2Block.norm`` operating on the local
    ``intermediate_local`` slice).
  * ``out_proj``: standard row-parallel — ``shard_dim=1`` slices the
    input (intermediate) dim along the same head-contiguous ranges;
    ``RowParallelLinear.forward`` all-reduces the partial output.

At TP=1 every cut above is a no-op identity slice, so the bind is
unchanged from the pre-TP replicated form.

Bind the routed-MoE mixer params for layer ``base``.

Each expert is a non-gated NemotronHMLP (``up_proj`` + ``down_proj``,
the same pair :func:`_add_mlp_keys` binds), kept per-expert rather
than stacked so a per-linear quant backend can bind each expert by
path. ``gate`` carries the DeepSeek-V3 ``noaux_tc`` selection bias.

Forward over the hybrid stack.

Per-layer dispatch is data-driven via :class:`LayerStack`:
each adapter pulls its kind-specific metadata + state-view out
of the per-:class:`StateKind` lookup tables built below, plus
any per-layer extras (``rope_cache``, ``attn_ops[i]``) out of
``**extra``.

The lookup tables are populated only for state kinds the model
actually uses — :attr:`StateKind.NONE` (MLP-only layers) is
intentionally absent so the dict ``.get`` returns ``None`` and
the MLP adapter discards the slot.

NemotronLabs-VoiceChat-11B — STT backbone + the composite served model.

The real checkpoint (``NVIDIA-NemotronLabs-VoiceChat-11B``) is a
composite speech-to-speech model: a perception (audio) tower, this
text backbone, an RNNT decoder, and a separate TTS stack, all bundled
under one ``model.safetensors``. This file has two top-level classes:

  * :class:`NemotronVoiceChatBackboneModel` — ONLY the backbone: the
    hybrid Mamba-2 / attention / MLP-only decoder plus its two output
    heads (``lm_head`` for text, ``function_head`` for function-call
    tokens). It does not import or reference the perception tower, the
    TTS stack, or the audio codec — usable stand-alone for
    text-in/text+function-out.
  * :class:`NemotronVoiceChatModel` — the composite served model:
    backbone + perception tower (audio-in) + TTS backbone/MoG head +
    RVQ-VAE codec (audio-out), wired for TURN-BASED serving. See its
    own docstring for the full public-interface contract. The RNNT
    decoder/joint (``stt_model.rnnt_decoder.*`` / ``rnnt_joint.*``, 15
    tensors) is bound as :attr:`NemotronVoiceChatModel.rnnt_head`
    (:class:`~arbi_serve.models.audio.nemotron_voicechat_rnnt
    .NemotronVoiceChatRnntHead`) — the checkpoint's own streaming
    endpoint detector, decoded per 80 ms frame by the duplex path's
    EOU-gated BOS admission signal (design doc §7.32); turn-based ASR
    still uses the backbone's own ``lm_head``, never this head.

Architecturally the backbone is a NemotronH hybrid stack (see
``nemotron_h.py``): per-layer Mamba-2 / NoPE-GQA-attention / standalone
relu²-MLP, selected by a ``hybrid_override_pattern`` string, same as
NemotronH-3-Nano. This file reuses NemotronH's layer-block classes,
per-arch dispatcher, and weight-map helpers directly (imported, not
copy-pasted) — only the checkpoint's top-level naming and the
untied-dual-head epilogue differ.

Checkpoint naming (verified against the real safetensors header at
``/mnt/k8scache/models/NVIDIA-NemotronLabs-VoiceChat-11B/model.safetensors``,
1632 tensors total — this file binds only the ``stt_model.{embed_tokens,
llm.*,lm_head,function_head}`` subset, 342 tensors):

  ``stt_model.embed_tokens.weight``        [131072, 4480]
  ``stt_model.llm.layers.{i}.norm.weight`` [4480]            — pre-mixer norm, every layer
  ``stt_model.llm.layers.{i}.mixer.*``     — Mamba-2 / attention / MLP-only, per hybrid_override_pattern
  ``stt_model.llm.norm_f.weight``          [4480]            — final norm
  ``stt_model.lm_head.weight``             [131072, 4480]    — untied
  ``stt_model.function_head.weight``       [131072, 4480]    — untied, second head over the SAME final hidden state

``tie_word_embeddings=False``: ``embed_tokens``, ``lm_head``, and
``function_head`` are three separate tensors in the checkpoint (byte-
compared, not aliased). NoPE attention (no rotary embeddings) — see
:class:`NemotronVoiceChatBackboneConfig` for why ``rope_theta`` is
carried but unused.

``config.json`` at the checkpoint root is a NeMo Hydra TRAINING config
(``data`` / ``exp_manager`` / ``model`` / ``trainer`` / ``_rnnt_merge_info``
top-level keys) — there is no HF-style ``architectures[]`` /
``hidden_size`` / ``hybrid_override_pattern`` section to read shapes
from. The one relevant field it DOES carry —
``model.stt.model.pretrained_llm == "nvidia/NVIDIA-Nemotron-Nano-9B-v2"``
— names the real source of shape truth: this backbone is initialized
from NVIDIA's Nemotron-Nano-9B-v2 (a NemotronH hybrid). Since that
model's HF ``config.json`` is not bundled here, :data:`_BACKBONE_SHAPE`
hardcodes the values directly (cross-checked against this checkpoint's
own tensor shapes via the safetensors header, not re-derived at load
time) rather than parsing a config section that doesn't exist.

Shape truth for the NemotronLabs-VoiceChat-11B STT backbone.

Frozen dataclass mirroring the HF NemotronH config fields this
backbone needs (see module docstring for why there is no real
``config.json`` section to parse them from). :meth:`from_hf_config`
still takes the checkpoint's raw Hydra config dict — not to read
shapes out of it (there are none), but to assert it names the
expected base LLM before handing back the hardcoded shape bundle,
so a wrong/renamed checkpoint directory fails loud instead of
silently booting with the wrong shapes.

Holds the ``stt_model.llm.*`` subtree — decoder layers + final norm.

No embedding here: unlike NemotronH (whose ``backbone`` owns the
embedding table too), this checkpoint nests ``embed_tokens`` one
level up, directly under ``stt_model`` — see
:class:`_NemotronVoiceChatInner`.

Holds the ``stt_model.*`` subtree: embed_tokens, llm, lm_head,
function_head — the four tensors/subtrees this backbone binds.

``lm_head`` and ``function_head`` live here (not on the top-level
model) because the checkpoint nests both directly under
``stt_model.*``, not under ``stt_model.llm.*``.

NemotronLabs-VoiceChat-11B STT backbone — Mamba-2 + attention +
MLP-only hybrid decoder with two untied output heads.

This is ONLY the backbone half of the composite VoiceChat
checkpoint (see module docstring) — no perception tower, no RNNT
decoder, no TTS stack. A later integration pass composes this with
those pieces; until then this class is usable stand-alone for
text-in/text+function-out.

Deviates from the plain :class:`arbi_serve.models.base.ModelBase`
forward contract (single ``(num_seqs, vocab_size)`` logits tensor)
because there are two heads sharing one hidden state — see
:meth:`forward`.

``LOGITS_HEAD_NAMES`` is the capability marker the GENERIC engine
dispatch (:func:`arbi_serve.runtime.forward_exec.collapse_dual_head_logits`,
the cudagraph capture builders in ``runtime/capture/{decode,prefill,
mixed}.py``) reads to recognize this "returns a tuple of per-head
logits, not one tensor" shape — checked by model ATTRIBUTE, not by
hardcoded ``isinstance`` on this class, mirroring the
``output_modalities``/``mm_bindings`` capability-check convention
(:func:`arbi_serve.multimodal.output.resolve_output_spec` /
:func:`arbi_serve.multimodal.registry.resolve_mm_bindings`). The
generic dispatch samples ONLY the first-named ("text") head — the
function-call head is intentionally reachable only through the
purpose-built realtime turn loop (``runtime/nemotron_voicechat_turn.py``
→ ``EngineSttStep``), which drives :meth:`forward` directly and
samples BOTH heads itself, bypassing this generic seam entirely. A
plain ``/v1/completions``-style caller therefore gets ordinary text
generation (including any text-parsed tool-call syntax the served
text stream itself contains); this model's dedicated function-call
head is not surfaced there. See :meth:`forward`.

The reference's ``AddFusion.forward`` (``nemo.collections.speechlm2
.parts.fusion``), specialized to this checkpoint's real per-modality
weights and its ``user_text_embeds is None`` shape (see
:data:`DUPLEX_TEXT_CHANNEL_WEIGHT` / :meth:`NemotronVoiceChatBackboneModel
.fuse_stt_step_embeds` for why). A free function (not a method) so it
is directly unit-testable against hand-computed expected values
without constructing a model.

Replicated (never TP-sharded), 1:1 weight map for a submodule whose
attribute names already mirror the checkpoint's own key suffix (after
stripping ``src_prefix``) — a mechanical walk over both
``named_parameters()`` and ``named_buffers()`` (real, checkpoint-backed
buffers like the TTS char-encoder's continuation/BOS-flag lookup
tables and its frozen ``audio_prompt_projection_W``), no remapping.

Mirrors :meth:`NemotronVoiceChatPerceptionTower.weight_map`'s own
convention (see that method's docstring for the two-piece dst/src
split this generalizes) — written once here for the pieces (TTS
backbone, audio codec) that don't carry their own version of it,
rather than duplicating the walk per piece.

The composite NemotronLabs-VoiceChat-11B: STT backbone + audio-in
perception tower + TTS backbone/MoG head + RVQ-VAE audio codec, wired
for TURN-BASED serving (full duplex is explicitly out of scope — see
the approved project plan, ``docs/omni_serving.md``'s duplex-roadmap
section).

=== Public interface (the contract a turn-orchestration loop drives) ===

This class does NOT implement ``output_modalities()`` /
:class:`arbi_serve.multimodal.output.OmniOutputSpec` — unlike
Step-Audio-2's single interleaved discrete-token stream, this model
has TWO text-side logit heads plus a wholly separate, non-discrete
TTS generation pathway (MoG head → continuous RVQ codes → codec
decode); forcing that shape through the id-range-classified
``OmniOutputSpec`` seam would be the wrong kind of reuse. Instead a
bespoke turn-orchestration loop (``arbi_serve/runtime/
nemotron_voicechat_turn.py`` — NOT written by this pass; a
different, not-yet-written task owns it) drives the pieces below
directly:

  * **Text/function decode** (STT half): :meth:`forward` — identical
    signature/contract to :meth:`NemotronVoiceChatBackboneModel.forward`
    (delegates to it entirely; see below), including
    ``return_hidden_state=True`` for whatever downstream use the
    loop needs of the full hidden state (NOT TTS conditioning — see
    the note on ``context_hidden_state`` below). The turn-scoped
    driver (:class:`~arbi_serve.runtime.nemotron_voicechat_stt_step
    .EngineSttStep`) calls this directly, one decode step at a time,
    NOT through the ordinary engine scheduler — see that class's own
    docstring for why (two untied heads don't fit the generic
    one-head dispatch contract).

  * **Audio-in frame fusion**: :meth:`fuse_stt_step_embeds` — the
    per-step embedding formula a genuinely frame-lockstep STT decode
    loop needs (one backbone step per audio frame, fusing that
    frame's continuous embedding with the previous step's own
    emitted text-/function-channel tokens — see that method's own
    docstring for the reference citation). Feeds :meth:`forward`'s
    ``inputs_embeds`` parameter directly.

  * **Audio input**: ``self.mm_bindings["audio"]`` — the standard
    :class:`~arbi_serve.multimodal.registry.MediaBinding` seam,
    wired exactly as
    :class:`~arbi_serve.models.step_audio2.StepAudio2Model` wires
    its own audio tower. Populated only when ``ARBI_ENABLE_AUDIO``
    is set; a text-only boot leaves it empty and every multimodal
    branch in :meth:`forward` (inherited from the backbone) no-ops.

  * **TTS turn start**: :meth:`prime_tts_turn` — warms the TTS
    Gemma3 backbone's KV cache with a speaker-conditioning prefix
    (``audio_prompt_latent``, e.g. ``self.tts_default_voice_prompt_latent``
    for the bundled "Aria" voice, or a
    precomputed
    ``.pt`` bundle's ``audio_prompt_latent``). Call once per turn,
    BEFORE the first :meth:`generate_tts_frame` call.

  * **TTS per-frame generation**: :meth:`generate_tts_frame` — one
    autoregressive frame: TTS backbone step → CFG+MoG code sampling
    → codec decode → waveform chunk. Call repeatedly, threading
    ``prev_code``/``past_key_values`` from one call's return values
    into the next call's arguments (see its own docstring for the
    exact state-threading contract).

  * **Weight loading**: :meth:`from_safetensors` / :meth:`weight_map`
    pieces (see "Weight loading" below).

  * **Boot-time VRAM profiling**: :meth:`profile_worst_case_tts_frame`
    — a duck-typed hook the generic activation-peak profiler
    (:func:`arbi_serve.runtime.profile_peak.profile_tts_codec_peak`)
    looks up by name, not part of the turn-orchestration contract
    above. Runs one representative :meth:`generate_tts_frame` call so
    the codec-decode/backbone/MoG-sampling transient — otherwise
    invisible to the generic ``ModelRunner.forward`` synthetic probe
    — is folded into the KV-pool sizing budget.

  * **Boot-time TTS paged-KV reserve**:
    :meth:`tts_kv_growth_reserve_bytes` — a separate duck-typed hook
    :func:`~arbi_serve.engine.inprocess_capture.serving_floor_for_grow`
    calls directly. Its closed-form capacity covers the speaker prefix,
    page rounding, both CFG rows, every admitted connection, and the
    configured session-duration ceiling. The single-frame activation
    profile and this persistent page slab are independent budget terms.

=== Why ``generate_tts_frame`` does NOT take a raw STT hidden-state
tensor (deviating from this class's original one-line scoping note) ===

The reference NeMo source (``RVQEARTTSModel.forward``'s
``context_hidden_state`` parameter, fed through a dedicated
``self.embed_context = nn.Linear(context_hidden_size, hidden_size)``)
DOES support exactly that kind of external-LM-hidden-state
conditioning in general — but this specific checkpoint's
``config.json`` sets ``context_hidden_size`` to ``None``, so
``embed_context`` is never constructed for VoiceChat-11B (confirmed
two ways: the real checkpoint's ``tts_model.tts_model.*`` header has
no ``embed_context.*`` key among its 418 tensors, and
:meth:`NemotronEarTTSBackbone._prepare_conditioning`'s port already
omits the ``context_hidden_state`` branch for exactly this reason —
see that method's docstring). The reference's own real per-frame
call (``DuplexEARTTS.infer_codes_one_step``) confirms the STT/TTS
text coupling for THIS checkpoint is instead entirely at the
DISCRETE TEXT level: the previous TTS-side subword token id is
looked up through the TTS backbone's OWN embedding table
(``self.embed_tokens(prev_subword_id)``, itself gated behind the
same always-``None`` ``context_hidden_size`` check and therefore
ALSO inert for this checkpoint) — never a continuous tensor handed
across from the STT backbone's 4480-dim hidden space into the TTS
backbone's unrelated 1152-dim one. So the turn-orchestration loop's
real STT→TTS coupling is: sample text from the STT backbone's
``text_logits``, decode it to a string, and re-tokenize that string
into the TTS side's own subword vocabulary as this method's
``subword_ids``/``subword_mask`` — the mechanism
:class:`~arbi_serve.models.audio.nemotron_ear_tts.CharAwareSubwordEncoder`
(``self.tts_backbone.embed_subword``) actually implements.

=== Weight loading ===

:meth:`from_safetensors` loads all four pieces from the same
checkpoint in one pass, reusing each piece's own already-correct
construction/weight-map logic rather than re-deriving any tensor
naming:

  * the STT backbone is fully self-contained — its own
    ``_apply_quant_and_materialize`` + :func:`load_model_weights`
    (the SAME generic, quant-/TP-aware loader path
    ``NemotronVoiceChatBackboneModel.from_safetensors`` uses standalone);
  * the perception tower reuses its own
    :meth:`NemotronVoiceChatPerceptionTower.weight_map` (dst-prefixed,
    src-key rebased onto ``stt_model.perception.*`` here — the
    tower's own method deliberately returns tower-LOCAL src keys, see
    its docstring, so this integration pass is where the checkpoint
    prefix gets re-attached);
  * the TTS backbone and the audio codec have no ``weight_map()`` of
    their own (neither piece's own agent needed one, since neither
    is independently served) — :func:`_mechanical_weight_map` above
    does the same mechanical, no-remapping walk
    :meth:`NemotronVoiceChatPerceptionTower.weight_map` already
    established as this codebase's convention for "attribute names
    already mirror the checkpoint 1:1" pieces;
  * four small buffers this composite owns directly (see "Tensors
    this composite owns" below) that don't belong to any one piece.

None of the non-backbone pieces are ever TP-sharded (matching their
own weight-map conventions) — :meth:`_load_replicated_weights`
binds them with a plain, unsharded ``.copy_()`` walk instead of
routing through the generic (TP-/quant-aware, and therefore
heavier) :func:`load_model_weights` a second time. ``TP>1`` is not
implemented for the same reason the backbone alone isn't (Mamba-2's
fused ``in_proj`` splitter isn't built) — see :meth:`from_safetensors`.

=== Tensors this composite owns directly (not part of any one piece) ===

Cross-referenced against the real checkpoint's safetensors header
(see this PR's coverage self-check): four small tensors live under
``tts_model.*``/``tts_model.tts_model.*`` that neither
:class:`~arbi_serve.models.audio.nemotron_ear_tts.NemotronEarTTSBackbone`
nor :class:`~arbi_serve.audio.nemotron_audio_codec.NemotronAudioCodecDecoder`
claims (both pieces' own docstrings say as much — the TTS backbone's
says "every top-level entry EXCEPT ``rvq_embs``, which the codec
module owns", but the codec's OWN docstring says it carries no such
buffer either; this composite is where that orphaned tensor finally
gets a home):

  * ``tts_rvq_embs`` (``tts_model.tts_model.rvq_embs``,
    ``[31, 1024, 512]``) — the per-quantizer-depth embedding table
    :func:`~arbi_serve.models.audio.mog_head.sample_frame_codes` and
    :class:`~arbi_serve.models.audio.nemotron_ear_tts.NemotronEarTTSBackbone`
    need. Stored at its REAL checkpoint shape (no sentinel row) so
    the weight-map coverage check stays a plain shape-preserving
    1:1 map like every other entry; :meth:`generate_tts_frame` pads
    it with the extra zero sentinel row those two functions expect
    (``depthsum_embedding``'s documented ``codebook_size + 1``-row
    interface) at call time via ``F.pad``.
  * ``tts_control_codes`` (``tts_model._control_codes``, ``[3]``
    int64) and ``tts_codec_silence_tokens``
    (``tts_model.codec_silence_tokens``, ``[31]`` int64) — precomputed
    control/silence code vectors the reference's
    ``DuplexEARTTS.decode_one_audio_step``/``infer_codes_one_step``
    use to sanitize control codes before codec decode and to seed
    silence frames respectively (``nim_extract/.../duplex_ear_tts.py``,
    ``replace_control_speech_codes`` / ``inference_force_speech_silence_on_eos``).
    Exposed here for the turn-orchestration loop to use the same
    way; this pass does not itself implement that sanitization logic
    (see the TODOs below) — ``tts_control_codes``' exact 3-entry
    ordering (hypothesized to be ``[speech_pad_id, speech_bos_id,
    speech_eos_id] == codebook_size + {0, 1, 2}``, matching the
    reference's ``speech_pad_id``/``speech_bos_id``/``speech_eos_id``
    properties) is NOT independently confirmed against real values —
  * ``tts_default_voice_prompt_latent`` (``tts_model.audio_prompt_latents.Aria``,
    ``[1, 37, 1152]``) — the bundled default ("Aria") voice's
    precomputed speaker-conditioning latent (see
    the reference preparation module's
    docstring for the exact derivation and why ``T=37`` for the
    checkpoint's own 3.0s default prompt duration), ready to pass
    straight into :meth:`prime_tts_turn` with no prep script run
    needed.

``weight_map()`` covers exactly 1630 of the real checkpoint's 1632
tensors (verified against the real safetensors header). The
2-tensor gap is
``stt_model.perception.preprocessor.featurizer.{fb,window}`` (the
mel filterbank + Hann window) — deterministic, recomputed in closed
form rather than loaded; this is
:class:`NemotronVoiceChatPerceptionTower`'s OWN pre-existing,
already-tested convention (see its ``weight_map`` docstring and
``test_nemotron_voicechat_perception.py::test_weight_map_covers_real_checkpoint_header``).
The streaming RNNT decoder/joint (``stt_model.rnnt_decoder.*`` (9) +
``stt_model.rnnt_joint.*`` (6)) loads into :attr:`rnnt_head` — the
duplex EOU admission signal's endpoint detector (design doc §7.32).

=== Documented TODOs for the turn-orchestration-loop task (genuinely
belongs there, not here) ===

  1. ``self.tts_backbone.embed_subword.set_char_map(subword_id_to_char_ids)``
     MUST be called once (using arbi-serve's own tokenizer — this
     module has none) before any ``subword_ids``-conditioned
     :meth:`generate_tts_frame`/:meth:`prime_tts_turn` call, or the
     char-aware subword encoder silently produces degenerate
     (empty-char-sequence, near-zero) text-conditioning embeddings.
     See :class:`~arbi_serve.models.audio.nemotron_ear_tts.CharAwareSubwordEncoder`'s
     own docstring — this is the one piece of TTS conditioning state
     that is genuinely NOT a checkpoint tensor (unlike the
     BOS/continuation flag buffers, which ARE real checkpoint data
     and load automatically via :meth:`weight_map`).
  2. Seeding ``prev_code`` for the very FIRST :meth:`generate_tts_frame`
     call of a turn: this pass exposes ``self.tts_codec_silence_tokens``
     as a reasonable default seed (broadcast to ``[B, 1, num_quantizers]``),
     but the reference's exact ``offline_inference`` protocol instead
     seeds with the audio PROMPT's own last code frame
     (``prompt_codes[:, -1:]``) when available — a
     ``prepare_nemotron_voicechat_voice.py``-baked custom-voice bundle
     DOES save ``prompt_codes`` for this; the bundled default
     "Aria" latent does not (only the latent itself is in the
     checkpoint). Closer reference fidelity for a custom voice is
     the orchestration loop's call.
  3. ``guidance_enabled`` MUST stay constant for an entire turn across
     :meth:`prime_tts_turn` and every subsequent
     :meth:`generate_tts_frame` call — both methods independently
     batch-double their inputs/outputs (``past_key_values`` included)
     when CFG is on, mirroring the reference's own per-call doubling
     (see :meth:`NemotronEarTTSBackbone.forward`'s ``guidance_enabled``
     branch); mixing an unguided ``prime_tts_turn`` with a guided
     :meth:`generate_tts_frame` (or vice versa) desyncs the KV cache
     batch dimension. This class does not enforce the invariant
     itself (it is stateless — see the state-threading contract on
     :meth:`generate_tts_frame`); the orchestration loop owns turn
     lifecycle and must not vary it mid-turn.
  4. Control-code sanitization before codec decode
     (``replace_control_speech_codes`` in the reference — swapping
     any stray BOS/EOS/PAD control code that leaks into a generated
     frame for a real silence code) IS implemented in
     :meth:`generate_tts_frame`, on its ``sanitize_control_codes``
     default, using this class's own ``tts_control_codes``/
     ``tts_codec_silence_tokens`` buffers — at the decode call site
     and on the decoded copy only, exactly where the reference puts
     it. It cannot live in the orchestration loop any more: the
     streaming vocoder cache (design doc §7.20) makes decode a
     stateful, exactly-once-per-frame call, so a caller cannot
     re-decode a frame after the fact to repair it.
  5. Chat-template / prompt-markup boundary tokens around an
     audio-input placeholder run (Step-Audio-2's
     ``<audio_start>``/``<audio_end>``-equivalent) are intentionally
     NOT invented here (``extra={}`` on the ``"audio"``
     :class:`~arbi_serve.multimodal.registry.MediaBinding`) — the
     approved project plan scopes the chat-template work as a
     separate piece; a real convention should come from there, not
     be guessed here.
  6. ``AUDIO_PATCH_ID``'s exact collision-safety (see the constant's
     own comment above) is unconfirmed without VoiceChat-11B's own
     real tokenizer artifact; a live parity run should double-check
     before this ships to production traffic.

The head KV/quality calibration hooks for reference logits.

This model has no top-level ``lm_head`` attribute (dual-head, see
class docstring) — generic single-head-assuming code (the TKV
calibration capture/drift drivers,
:func:`arbi_serve.calibration.engine_capture.resolve_lm_head_module`)
reads this instead. Points at ``LOGITS_HEAD_NAMES[0]`` ("text"),
the SAME head the generic engine dispatch samples — calibration
stays consistent with what a plain ``/v1/completions`` caller
actually sees.

``lm_head`` over an ARBITRARY set of post-final-norm rows.

Overridden rather than inherited from
:meth:`LayerStackModelMixin.logits_from_hidden` — the mixin's
default routes through ``self._lm_head`` → ``self.lm_head``, which
this dual-head model doesn't have at the top level (see
:attr:`calibration_lm_head`'s docstring). This is the seam
:mod:`arbi_serve.engine.logprobs` needs to score prompt positions
other than the last (``--logprobs``/``--prompt-logprobs``); routes
through the SAME "text" head :attr:`calibration_lm_head` and the
generic single-tensor engine dispatch both use. No
softcap/scale epilogue: neither is a NemotronH-family concept
(that's Gemma/Granite-only), so the plain head call is exact.

Forward over the hybrid stack, returning BOTH heads' logits.

Returns ``(text_logits, function_logits)`` — both
``(num_seqs, vocab_size)`` — computed from the SAME per-seq
last-token hidden state, via ``lm_head`` and ``function_head``
respectively (untied, independent weights).

``return_hidden_state=True`` additionally returns the full
per-token post-final-norm hidden state (``(N_tokens,
hidden_size)``, same contract as
:meth:`arbi_serve.models.base.ModelBase.forward`'s single-head
variant) as a third tuple element, for a later TTS-coupling
milestone: ``(text_logits, function_logits, hidden_full)``.

Unlike single-head archs this does NOT route through
:meth:`LayerStackModelMixin._lm_head` / ``batch.logits_out`` —
that helper assumes one head at a fixed ``self.lm_head``
attribute path and one destination buffer; with two heads
sharing one hidden state, the cudagraph-capture destination-
buffer optimization it provides is left as a follow-up once
this backbone is wired into the full composite model.

``inputs_embeds``: ``(num_tokens, hidden_size)`` precomputed
embeddings that BYPASS ``embed_tokens``/the multimodal merge
below entirely — the frame-lockstep STT turn loop
(``arbi_serve/runtime/nemotron_voicechat_stt_step.py``'s
``EngineSttStep``) uses this to feed one audio frame's FUSED
embedding (:meth:`fuse_stt_step_embeds`) directly, since that
embedding is not any single token's ``embed_tokens`` row.
``input_ids`` must still be a real, same-length tensor even when
``inputs_embeds`` is given (its VALUES are unread in that case —
only its shape/dtype/device matter to the surrounding
``ScheduledBatch``/attention-metadata plumbing). ``batch.mm`` is
ignored when ``inputs_embeds`` is given (the caller has already
done all embedding-level fusion itself).

Multimodal merge (audio-in, ``inputs_embeds is None`` only): if
``batch.mm`` is populated (the composite
:class:`NemotronVoiceChatModel` below sets it exactly as
:class:`~arbi_serve.models.step_audio2.StepAudio2Model` does for
its own audio tower), the perception tower's pre-encoded
embeddings overwrite the placeholder-token rows of the
embedding stream before the layer stack runs. This backbone
still does not import or construct any tower itself (see module
docstring) — it only conditionally reads a generic, engine-side
``batch.mm`` seam, so plain text-in callers (``batch.mm is
None``) are completely unaffected.

Plain token embedding for rows that carry no embed override.

The seam ``forward_exec.resolve_embed_override`` uses to serve a
MIXED batch — some rows carrying a fused
``pending_embed_override``, some not. The duplex lane emits
exactly that shape: its seed-prefill and context-injection rows
carry prompt / caller-supplied tokens rather than one fused audio
frame, so no override is parked on them, while steady-state frame
rows in the SAME slate do carry one (design doc §7.9.14).

Deliberately mirrors :meth:`forward`'s own ``inputs_embeds is
None`` branch, MINUS the ``batch.mm`` multimodal merge — an
un-overridden row's audio, when it has any, arrives through the
override rather than through ``batch.mm``, and a duplex request
never carries ``mm`` feats at all. A caller that needs the merge
must not route through here.

One frame-lockstep STT decode step's FUSED input embedding.

Direct port of the reference's default fusion strategy —
``AddFusion`` (``nemo.collections.speechlm2.parts.fusion``,
``fuse_method`` defaults to ``"add"`` and this checkpoint's
config.json never overrides it) — applied exactly as
``DuplexSTTModel._step_inference``/``_step_zero``
(``nemo.collections.speechlm2.models.duplex_stt_model``) call it:
a per-modality weighted sum of the previous step's own text-
channel token embedding, this step's continuous audio-frame
embedding, and the previous step's function-channel token
embedding. ``user_text_embeds`` (the reference's optional ASR
channel) is omitted — this checkpoint trains with
``predict_user_text=False`` (config.json), so the reference's
own ``AddFusion.forward`` never receives that argument for
VoiceChat-11B either.

The three weights (:data:`DUPLEX_TEXT_CHANNEL_WEIGHT`,
:data:`DUPLEX_USER_CHANNEL_WEIGHT`,
:data:`DUPLEX_FUNCTION_CHANNEL_WEIGHT`) are this checkpoint's
real ``model.stt.model.duplex_{text,user,function}_channel_weight``
config.json values, not a guess.

Args:
    audio_frame_embeds: ``(n, hidden_size)`` — this step's
        continuous frame embedding(s), straight from
        :meth:`~arbi_serve.models.audio.nemotron_voicechat_perception
        .NemotronVoiceChatPerceptionTower.encode`'s per-frame rows
        (same flattened multi-sequence convention as
        ``embed_tokens(input_ids)`` elsewhere in this class — NOT
        a ``(B, T, H)`` batch tensor).
    prev_text_token_ids/prev_function_token_ids: ``(n,)`` long —
        the previous step's own sampled text-/function-channel
        token id (:data:`TEXT_PAD_TOKEN_STR`'s id, the "blank",
        before either channel has emitted a real token — see
        that constant's docstring).

Returns ``(n, hidden_size)``, ready to pass as this backbone's
own :meth:`forward`'s ``inputs_embeds``.

Per-parameter shard spec for the generic loader.

Every dst key already carries the full ``stt_model.*`` prefix
(matching :class:`_NemotronVoiceChatInner`'s attribute
naming), so every src key is identical to its dst key — no
rebasing needed for the embed/norm/head entries. The per-layer
mixer entries reuse NemotronH's ``_add_*_keys`` helpers with
``prefix="stt_model.llm"`` (see ``nemotron_h.py``) instead of
duplicating the ~150 lines of key-mapping logic.

See :meth:`NemotronVoiceChatBackboneModel.logits_from_hidden`
— this composite is what ``eng.model`` actually holds for a
normal boot (see :attr:`LOGITS_HEAD_NAMES`'s comment above), so
the ``arbi_serve.engine.logprobs`` seam needs this delegation on
THIS class, not just the backbone.

Full delegation to :meth:`NemotronVoiceChatBackboneModel.forward`
— same signature, same ``(text_logits, function_logits[, hidden_full])``
contract. The backbone itself reads ``batch.mm`` (populated by the
engine from this class's ``mm_bindings["audio"]``) to merge in the
perception tower's embeddings, so there is no multimodal-specific
code to duplicate here — see
:meth:`NemotronVoiceChatBackboneModel.forward`'s own docstring.

Full delegation to
:meth:`NemotronVoiceChatBackboneModel.embed_input_ids` — see that
method's own docstring. Needed on THIS class for the same reason
:attr:`SUPPORTS_EMBED_OVERRIDE` is: the composite, not the
backbone, is what ``eng.model`` holds for a normal boot, and
``resolve_embed_override`` looks the hook up on ``eng.model``.

Warm the TTS decoder's KV state with the speaker-conditioning
prefix for the first :meth:`generate_tts_frame` call of a turn.

Args:
    audio_prompt_latent: ``[B, T_prompt, 1152]`` (or ``[2B, ...]``
        if ``guidance_enabled`` — see below), e.g.
        ``self.tts_default_voice_prompt_latent`` (the bundled
        "Aria" voice) or a precomputed voice asset
        ``.pt`` bundle's ``audio_prompt_latent``.
    subword_ids/subword_mask: optional system-prompt text
        conditioning during warmup (rare — the bundled "Aria"
        voice's own baking used none, per
        ``prepare_nemotron_voicechat_voice.py``'s docstring).
    guidance_enabled: MUST match every subsequent
        :meth:`generate_tts_frame` call for this turn — see the
        class docstring's

Mirrors the reference's turn-start warmup call
(``DuplexEARTTS.offline_inference``'s ``outputs =
self.tts_model(**init_inputs)`` line,
``nim_extract/.../duplex_ear_tts.py``): ``audio_prompt_latent``
stands in for the code-embedding of the prompt segment (it IS
already backbone-hidden-space, per how
``prepare_nemotron_voicechat_voice.py`` derives it — depthsum
embedding → ``embed_code`` → the frozen
``audio_prompt_projection_W``), fused with any text conditioning
via the SAME ``gated_fusion_audio_text`` module
:meth:`NemotronEarTTSBackbone.forward` uses, then run once
through the Gemma3 stack with ``use_cache=True``. Reuses
:meth:`NemotronEarTTSBackbone._prepare_conditioning` directly
(private-by-convention but same-package; avoids re-deriving its
null-embedding/CFG-flag logic) rather than duplicating it.

Returns connection-owned paged state on a serving engine. The
module-only reference path returns a ``past_key_values`` list.

One autoregressive TTS frame: backbone step → CFG+MoG code
sampling → codec decode → waveform chunk.

Composes exactly the three pieces named in the class docstring's
TTS-frame contract — :meth:`NemotronEarTTSBackbone.forward` (one
Gemma3 step), :func:`~arbi_serve.models.audio.mog_head.sample_frame_codes`
(CFG-guided iterative RVQ unmasking), and
:meth:`NemotronAudioCodecDecoder.decode` (RVQ dequant + vocoder)
— with no new math of its own.

State-threading contract (this method is otherwise STATELESS —
the orchestration loop owns turn lifecycle):

Args:
    prev_code: ``[B, 1, num_quantizers]`` long — the PREVIOUS
        frame's revealed codes (this step's autoregressive
        "input token"; see class docstring TTS KV state, from
        :meth:`prime_tts_turn` (first call) or the PREVIOUS
        :meth:`generate_tts_frame` call's returned cache.
    subword_ids/subword_mask: the text CURRENTLY being spoken,
        tokenized into the TTS side's own subword vocabulary
        (NOT the STT backbone's token ids — see the class
        docstring's "Why generate_tts_frame does NOT take a raw
        STT hidden-state tensor" section). Requires
        ``self.tts_backbone.embed_subword.set_char_map(...)`` to
        have been called once (class docstring MoG+CFG
        sampling knobs — defaults are this checkpoint's own real
        ``inference_guidance_scale``/``inference_top_p_or_k``/
        ``inference_noise_scale`` config values (``config.json``,
        cross-checked against
        :mod:`arbi_serve.models.audio.mog_head`'s module
        docstring); ``num_iter=8`` matches the reference's own
        hardcoded default (``DuplexEARTTS._get_generation_config``).
        A truthy ``guidance_scale`` enables CFG batch-doubling —
        see class docstring ``None`` decodes this
        frame in isolation — every conv left-pads with zeros and
        the overlap-add restarts, which puts an audible seam at
        every 80 ms boundary and a waveform that barely resembles
        the same codes decoded as one sequence (design doc §7.20
        has the measurements). Caller owns the object's lifetime,
        exactly like ``past_key_values``.
    sanitize_control_codes: replace BOS/EOS/PAD sentinel codes with
        ``tts_codec_silence_tokens`` in the codes handed to the
        vocoder (the returned ``new_code`` stays RAW for the
        autoregressive history — the reference sanitizes only at
        its decode call site). Applied unconditionally rather than
        behind a ``.any()`` probe: the probe is a host sync, and a
        second decode to repair an already-decoded frame would
        advance ``codec_cache`` twice.
    force_silence_code: DISCARD this frame's sampled codes and use
        ``tts_codec_silence_tokens`` instead — for the vocoder AND
        for the returned ``new_code``, which is the caller's next
        ``prev_code``. This is the reference's PAD/idle
        substitution (``model.py:658-693``), which replaces
        ``step_acoustic_tokens`` itself and then feeds that same
        substituted value back into its TTS stream, so a muted
        stretch re-anchors the autoregressive audio history to
        silence rather than accumulating PAD-conditioned
        hallucination inside it (design doc §7.56).
        Deliberately a parameter here rather than a fixup at the
        call site: the substitution has to land BETWEEN sampling
        and the vocoder decode to keep ``codec_cache``'s conv
        left-context on silence too, and a caller cannot reach
        that point. The MoG sampling loop is SKIPPED entirely when
        this is set — its output is discarded either way, and it
        is the expensive half of the frame.
        ``sanitize_control_codes`` is unrelated and orthogonal: it
        repairs sentinel codes for the vocoder only and keeps the
        RAW code in the history; this replaces both.

Returns:
    ``(waveform_chunk, new_code, present_key_values)``:

      - ``waveform_chunk``: ``[B, wav_to_token_ratio]`` float —
        this frame's audio samples.
      - ``new_code``: ``[B, 1, num_quantizers]`` long — this
        frame's revealed codes; becomes ``prev_code`` for the
        NEXT call.
      - ``present_key_values``: this frame's updated KV state;
        becomes ``past_key_values`` for the NEXT call. Native
        paging updates and returns the same connection-owned state.

Encode a worst-case dummy utterance for the boot activation profiler.

The perception tower runs eagerly, outside every pool and every captured
graph, so its transient needs a home in the serving floor the same way a
vision tower's does. Called by
:func:`arbi_serve.runtime.profile_peak.profile_perception_peak` (a
duck-typed hook) inside the same reset-peak window as the text / vision /
audio / TTS probes.

Sized at the perception capture ladder's widest bucket — the widest mel
input the captured encode path serves. A no-op when the audio path is
off.

Run one representative :meth:`generate_tts_frame` call for the
boot-time activation-peak profiler.

The frame-lockstep TTS chain (TTS decoder step → CFG+MoG
iterative code sampling → :class:`~arbi_serve.audio.nemotron_audio_codec.NemotronAudioCodecDecoder`
decode) is profiled before the serving paged-KV runtime is attached.
The generic
``ModelRunner.forward`` synthetic prefill the profiler otherwise
runs never exercises it. Called by
:func:`arbi_serve.runtime.profile_peak.profile_tts_codec_peak`
(a duck-typed hook: the generic profiler looks this method up by
name and no-ops when absent), inside the SAME reset-peak window
as the text/vision/audio probes, so its transient folds into the
measured ``max(...)`` peak the same way theirs do.

Primes the turn with :attr:`tts_default_voice_prompt_latent`
(the bundled "Aria" voice, already real weights by the time boot
profiling runs) and drives one silence-seeded frame with
:meth:`generate_tts_frame`'s own defaults — ``guidance_scale=0.2``
(this checkpoint's real inference config, so the CFG batch-doubling
that is on the hot path in production is exercised here too) and
``num_iter=8``, whose iterative unmasking loop's true peak lands
wherever it lands INSIDE this one call, which is what matters
since the caller resets peak stats before and reads
``max_memory_allocated`` after the whole call, not per iteration.
``subword_ids`` is left unset (skips the tiny character-encoder
module) — the profiler measures memory, not numbers, and that
module's activation footprint is negligible next to the 28-layer
backbone and the codec's upsampling stages. A no-op when the
audio path is off (``self.tts_backbone``/``self.audio_codec``
unset) — text-only boots are unaffected.

Physical capacity of the connection-bounded TTS paged-KV slab.

The closed-form size includes the speaker prefix, whole-page
rounding, page zero, every admitted connection, and both CFG rows.
It therefore reserves exactly the static allocation constructed by
the attached ``tkv-bypass`` runtime. Returns zero with audio disabled.

``(padded_table, per-depth squared norms)`` derived from
``tts_rvq_embs``, built once and reused.

Both are pure functions of a frozen codebook, so both are built on
first use and held for the process's life rather than rebuilt on
every 80 ms tick — the padded table is a ``[31, 1025, 512]``
allocation and the norms a 31-way reduction. The padded shape is
what :func:`~arbi_serve.models.audio.mog_head.depthsum_embedding`
and :func:`~arbi_serve.models.audio.mog_head.sample_frame_codes`
take; the norms are
:func:`~arbi_serve.models.audio.mog_head.depthsum_encoding_step`'s
``embs_sq``. Invalidated by :meth:`_load_replicated_weights`, the
only writer of ``tts_rvq_embs``.

Fill every entry of ``weight_map`` via a plain, unsharded copy.

Every entry here (perception tower, TTS backbone, audio codec,
and this composite's own directly-owned buffers) is replicated
across TP ranks and carries no quantization — unlike the STT
backbone (loaded separately, above, via the generic quant-/
TP-aware :func:`load_model_weights`), none of these pieces need
the ``shard_dim``/``shard_id``/``fused_split`` machinery, so a
direct ``.copy_()`` walk is the correct, simplest loader here —
mirrors :meth:`NemotronAudioCodecDecoder.load_from_full_state_dict`'s
plain ``load_state_dict``, generalized to a name-mapped source
key instead of an identical-name state dict (the whole reason
this composite needs a per-entry map rather than that method's
own prefix filter: attribute names get an extra ``<piece>.``
hop the checkpoint doesn't have).

Per-parameter shard spec for the generic loader (and this
composite's own coverage self-check) — see the class docstring's
"Weight loading" and "Tensors this composite owns directly"
sections for the per-piece strategy this method implements.

Resolve a ``session.voice`` name to a TTS speaker-conditioning
latent, the ``audio_prompt_latent`` :meth:`generate_tts_frame` /
:meth:`prime_tts_turn` need.

Falls back to :attr:`tts_default_voice_prompt_latent` (the bundled
"Aria" voice) for ``None``, ``""``, ``"default"``, and every name in
:data:`_NON_LOOKUP_VOICE_NAMES` (the realtime layer's own
cross-model sentinel/preset names — none of them name a real
``<model_dir>/voices/<name>.pt`` entry, so treating them as one
would either loop back to this exact default or 404 against a name
this checkpoint never had). Any other name loads
``<self.model_dir>/voices/<voice>.pt`` — the convention
the voice preparation step writes to (a
dict with an ``"audio_prompt_latent"`` key) — and caches the
resolved tensor in-process so a repeat lookup skips the disk read.

A voice name that fails to resolve (no ``model_dir``, missing file,
unreadable/malformed checkpoint) falls back to the default WITH a
logged warning, never a raised exception — a live duplex or turn-
based session must not die over an operator typo in a voice name.

Qwen3 dense model — standalone implementation.

Topology mirrors the official ``Qwen3ForCausalLM``:
  Embedding → [PreNorm → AttentionBlock + residual → PreNorm →
              GatedSiLUMLP + residual] × N → final RMSNorm → tied lm_head.

Per-head Q / K RMSNorm (Qwen3-specific) lives inside
:class:`AttentionBlock`. Per-arch shape metadata lives on
:class:`Qwen3Config` (mlc-llm pattern). :class:`Qwen3DecoderLayer`
implements the uniform :class:`LayerStack` block signature directly
— no separate shape-converter class.

Qwen3 dense shape config.

Shared shape lift / HF parse / ``ModelDims`` + ``LayerSpec``
projection live in :class:`DenseDecoderConfig`. Qwen3 adds only the
optional sliding-window field (populated from the HF
``use_sliding_window`` + ``sliding_window`` pair). The Qwen3-specific
per-head q/k-norm is a DecoderLayer behaviour, not a shape field.

One Qwen3 transformer block: pre-norm + attn + pre-norm + MLP.

Implements the uniform :class:`LayerStack` block signature directly.
The ``attn_op`` is engine-provided per-step — the layer holds no
backend state of its own. The post-attention norm fuses the residual
add (``_fused_post_norm``), so the residual skeleton in
:class:`PreNormDecoderLayer` takes the fused branch.

Embed tokens, run the decoder layer stack, and return last-token
logits (optionally with the full post-final-norm hidden state).

``compute_logits=False`` (with ``return_hidden_state=True``) skips
the last-token ``lm_head`` GEMV + TP all-gather and returns
``(None, hidden)`` — the MTP verify pass recomputes ``lm_head`` over
all K+1 rows itself, so the last-token head is per-step waste.

Qwen 3.5 / Qwen 3.6 — text-only causal-LM (`Qwen3_5ForConditionalGeneration`).

Hybrid attention + GDN per ``text_config.layer_types``:

  * ``"linear_attention"`` → :class:`GDNBlock` (FLA gated delta net).
  * ``"full_attention"``   → the shared
    :class:`arbi_serve.models.attn.AttentionBlock` with
    ``attn_output_gate=True`` (``q_proj`` outputs ``(num_heads,
    head_dim * 2)``; the second half is a sigmoid gate multiplied into
    the attention output before ``o_proj``).

Per-layer dispatch routes through :class:`LayerStack`; the inner
decoder layers implement the uniform LayerStack block signature
directly. The per-arch dispatcher
(:func:`_make_qwen3_5_decoder_layer`) picks the right class.
MTP-marker specs are filtered out of the iterator: the bundled MTP
head dispatches its own decoder block separately, after the main
loop produces post-final-norm hidden.

Weight-key namespace: the HF `Qwen3_5ForConditionalGeneration`
checkpoint nests text-tower tensors under ``model.language_model.*``;
``model.visual.*`` and external ``mtp.*`` keys are ignored
(text-only). Partial-rotary RoPE: ``partial_rotary_factor = 0.25`` —
only the first ``head_dim // 4`` channels carry rotation.

Holds the ``model.language_model.*`` subtree.

MTP-marker specs (``is_mtp_layer=True``) are skipped here — those
exist only to drive per-layer slab + attn-op allocation for the
bundled MTP head's own attention block; the main-model layer loop
must not instantiate a duplicate decoder layer for them. The head
holds the actual MTP decoder layer modules (see
:class:`Qwen3_5MtpHead`).

The ``layers`` :class:`nn.ModuleList` carries the **inner** decoder
layer instances — one per non-MTP spec, in spec order. The
LayerStack adapter wrappers (built by :class:`Qwen3_5Model.__init__`)
hold back-references to these same ``inner`` modules rather than
fresh copies, so PyTorch's ``named_parameters`` walk reaches each
parameter via exactly one path (``self.model.layers.{i}.*``) and
safetensors weight-map keys resolve unchanged from the
pre-migration layout. The adapters live outside ``self.model``
in a plain Python list to avoid double-registration.

Qwen 3.5 / Qwen 3.6 text-only causal LM.

Construction is parameterized by the ``text_config`` sub-dict of
`Qwen3_5ForConditionalGeneration`'s `config.json`. The vision
tower (`vision_config`) is intentionally not loaded — arbi-serve
is text-only by design.

Reconstruct the residual-stream output under cross-layer fusion.

A fused block returns only ``mlp_out``; the true residual-stream
output of layer ``i`` is ``residual_buf[:n] + out`` at this point (the
next layer's fused input_layernorm mutates ``residual_buf`` in place,
so the add must materialize here, not lazily). ``residual_buf`` rides
in ``extras[2]`` and is rung-padded — slice to the real token rows. On
the unfused path (``residual_buf is None``) the block already returns
``residual + h``, so ``out`` passes through.

Check out the per-(num_tokens) persistent residual buffer.

A CHECKOUT, not a lookup: the slab is marked in flight under
``num_tokens`` until :meth:`_release_residual_buf` returns it, and a
second checkout of the same key while the first is out raises by
name. Every layer's fused norm accumulates into the slab in place,
so two forwards of one width holding it together — two sub-chunks
pipelined against each other — would sum both residual streams into
the same rows with no shape error and wrong logits. The registry is
what makes that a loud failure instead. Sequential reuse of a width
returns the same tensor: check out, release, check out again.

First checkout at each ``num_tokens`` allocates inside the engine's
``graph_buffers_pool`` (passed to the model by
:func:`arbi_serve.engine.build.build` as ``_residual_buf_pool``);
captured layer-graphs at that bucket reference the resulting
``data_ptr``; later checkouts at the same shape reuse it. The
slab is per-num_tokens — a single buffer is shared across
``(num_seqs, is_prefill)`` variants at the same num_tokens since
the layer-piecewise capture buckets share the residual stream
(both attention and GDN layers see ``(num_tokens, hidden_size)``
flat residuals; they don't differ by num_seqs).

Cost on the served path: one dict membership test and one insert
(the release is one pop). No device work, no sync. Called from
inside the compiled forward, so the checkout is Dynamo-traceable
under ``fullgraph=True``: the dict mutations are side effects
Dynamo replays after the graph, and the raise is an untaken branch
guarded on the registry — which is empty at every forward entry,
because every forward releases in a ``finally`` — so the guard
holds and nothing re-traces on the first live request.

Sleep/resume: the buffer's stable VA is preserved across
sleep-phase2 release/resume because it lives in
``graph_buffers_pool``, which the existing
:class:`SleepableTensorPool` covers via the buffers_pool mempool
membership. The dict itself is not cleared on sleep; on resume
the same VAs map back to the same data_ptrs and the captured
graphs are bit-identical.

Compile-on whole-model fullgraph. The ``with
pool.use()`` ctx manager (a generator-based contextmanager) is
not Dynamo-traceable under ``fullgraph=True`` — the inliner
bails on ``__enter__`` with ``SKIPPED INLINING <code object
__enter__ at .../contextlib.py:132>``, breaking the whole-model
compile + cudagraph capture sweep. The ``if buf is not None:
return buf`` early-out is traceable, so pre-allocating every
``num_tokens`` we will ever see at boot — via
:meth:`prewarm_residual_bufs` — guarantees the in-forward call
always takes the early-out branch. Lazy alloc is kept here as
a backup for paths that pre-warm wasn't aware of (per-layer
compile, eager-only); they pay the graph-break cost
(per-layer compile already breaks on the call-site eager
boundary, so the additional break is invisible there).

Bounded fallback slab for OFF-RUNG / ABOVE-CAP serving forwards.

:meth:`_ensure_residual_buf` caches one permanent buffer per distinct
``num_tokens`` and never evicts — safe for the bounded set of captured
rungs (pre-warmed at boot + allocated during the capture sweep), whose
``data_ptr`` the captured layer-graphs bake. But a serving-time forward
whose token count lands off every captured rung — an above-cap eager
prefill (``num_tokens`` above the top piecewise rung, run eager by
design) or an eager off-rung decode — would otherwise allocate a fresh
permanent buffer for every distinct token count it ever sees. Under
sustained varied concurrency that grows ``_residual_bufs`` without
bound until the fixed-reservation ``capture.io_buffers`` named MemPool
is exhausted → ``torch.OutOfMemoryError`` on a ~MiB alloc while GiBs of
device VRAM sit free (the named pool can't grow past its reservation).

Such forwards run eager — no captured graph references their
``residual_buf`` by ``data_ptr`` — so a single shared, growable slab,
sliced to the live token count by the caller, is correctness-equivalent
to a per-shape buffer. This bounds the residual-buf footprint to the
captured-rung set plus one overflow slab (sized at the largest off-rung
width seen, ≤ the max prefill token count). Growing replaces the slab;
the previous one is referenced only by the just-finished eager forward,
so freeing it is safe. The slab lives in the same ``graph_buffers_pool``
(stable VA across sleep/resume) as the pinned buffers.

Mark slab ``key`` in flight for a forward at ``num_tokens``.

Raises when ``key`` is already out. The message names the slab and
both widths so the caller that overlapped two forwards on one slab
reads the collision, not a downstream logits mismatch.

Return the slab checked out under ``key``.

``key`` is the ``num_tokens`` a pinned slab was checked out at, or
:data:`_RESIDUAL_OVERFLOW_SLAB`. Returning a slab that is not out is a
bookkeeping fault of the caller and raises.

The ``num_tokens`` to size the cross-layer ``residual_buf`` at.

The per-layer dispatcher pads an off-rung prefill split-attn /
split-gdn forward up to the smallest captured rung ``>=`` its real
token count and replays that rung's captured graphs (see
:func:`arbi_serve.runtime.capture.dispatch.dispatch.dispatch_layer_args`).
Those captured pieces apply the fused ``input_layernorm`` /
``post_attention_layernorm`` in place against the ``residual_buf``
whose ``data_ptr`` was baked at the rung width — so the model
forward must hand every layer (and the final-norm tail) the
residual_buf sized to that same rung. Otherwise the captured
layers accumulate the cross-layer residual into the rung-width
buffer while the model forward zeroed + reads a different
real-width buffer → the entire residual stream is lost → garbage
under concurrency (the num_seqs>1 packed prefills are the only
forwards that land off-rung; single-seq prefills hit exact rungs).

Returns the real ``num_tokens`` unchanged when there is no pad-up:
decode steps, non-split / non-prefill forwards, exact-rung
prefills, above-cap (too-large-eager) prefills, or when piecewise
capture is disabled. In all those cases the captured layers (if
any) bake the real-width residual_buf, so it already matches.

Pre-allocate every ``residual_buf`` shape we will see at runtime.

Called from the boot orchestrator (engine.cudagraph_admin) before
:func:`precapture_compile_warmup` so that the trampoline's first
compile (and every per-shape recompile that the cudagraph
capture sweep triggers) finds an entry in
``self._residual_bufs[num_tokens]`` and takes the
Dynamo-traceable early-out branch in :meth:`_ensure_residual_buf`.

Must run outside any active compile / cudagraph window — the
``pool.use()`` ctx manager allocates regular memory; doing it
inside a captured-graph context would bake the alloc into the
graph and mis-route on replay. The boot orchestrator is the
canonical safe site (see :func:`precapture_compile_warmup`).

Idempotent — :meth:`_ensure_residual_buf` already short-circuits
on cached entries. Each slab is released as soon as it exists; a
prewarm holds nothing.

Run the interleaved GDN/attention layer stack + final norm.

With cross-layer (residual + input_layernorm) fusion enabled,
allocates a persistent ``residual_buf`` of shape (num_tokens,
hidden_size) in ``graph_buffers_pool`` (stable data_ptr across
captured replays), zeroes it once at step start, threads it
through every layer's input_layernorm fusion. ``hidden`` (the
embed) is the layer-0 pending-add; subsequent layers receive
the previous layer's ``mlp_out`` as pending-add. The cross-
layer add migrates into the next layer's input_layernorm
fusion. Final tail: ``residual_buf + last_mlp_out → norm``.
Without fusion: plain positional walk + ``self.model.norm``.

``lora_state`` rides the ``extras`` tuple (slot 3) to every layer
so MLP + full-attention LoRA applies. ``getattr`` keeps the CPU
test stubs (``_Batch`` without the field) working; production
``ScheduledBatch`` always carries it (``None`` when no adapter).

``arena`` rides slot 4 (appended, never inserted — the capture
dispatch reads ``attn_ops`` at ``block_args[4]``). The per-layer
dispatcher replaces it with each attention layer's output buffer
and blanks it for the GDN layers, so no block sees an allocator.

MTP layer KV-fill pass. When the bundled head is present we
run its decoder layer over the same flat batch the main model
just saw, reusing ``batch.attn_meta`` and ``batch.positions``.
Effect: the MTP layer's per-layer KV slab is populated at
exactly the same slot_mapping the main model wrote to (page-
id basis is shared across layers; the per-layer slab is
indexed by the request's pages + offsets), so the next draft
step reads real history. Only runs when the model carries an
MTP head and ``batch.attn_meta`` is populated (i.e. PAGED_KV
path); pure-prefill profiling forwards built outside the
engine path skip this safely. Ordering: norm(embed) +
norm(hidden) -> cat -> fc -> decoder layer.

Embed tokens, run the interleaved GDN/attention layer stack (with
optional cross-layer residual fusion), and return last-token logits
(optionally with the full hidden state).

``compute_logits=False`` skips the final ``lm_head(last_hidden)``
GEMV and returns ``(None, hidden)``. The MTP verify pass wants only
``hidden`` (it recomputes ``lm_head`` over all K+1 flat tokens
itself), so the last-token head here is pure waste — a full-vocab
GEMV (+ a TP all-gather) discarded every spec step. Only honoured
together with ``return_hidden_state=True``.

Qwen 3.6 MoE — `Qwen3_5MoeForConditionalGeneration` (Qwen3.6-35B-A3B).

Same hybrid backbone as the dense Qwen 3.5 / 3.6 arch — GDN
linear-attention + gated full-attention per ``text_config.layer_types``,
``(1 + weight)`` RMSNorm convention, partial RoPE — with the dense gated
MLP replaced by a sparse MoE FFN on every layer
(:class:`Qwen3_5MoeSparseMoeBlock`):

  * router ``mlp.gate`` (dense bf16): softmax over all ``num_experts``
    logits → top-``num_experts_per_tok`` → always renormalize (HF
    ``Qwen3_5MoeTopKRouter`` has no ``norm_topk_prob`` knob — the top-k
    weights are unconditionally renormalized to sum 1);
  * routed experts — the representation is the pack's to choose
    (:func:`~arbi_serve.models._moe_dispatch.resolve_expert_dispatch`).
    When the backend can serve the expert set stacked, the checkpoint's
    per-expert ``mlp.experts.{E}.{gate,up,down}_proj`` tensors bind into
    one :class:`arbi_serve.models.moe.FusedMoE` weight pair per layer and
    the fused Triton kernel runs a single binned GEMM per stage instead of
    one GEMM per touched expert; the reference pack is block-wise FP8
    (``weight_scale_inv`` @ 128×128), consumed in-kernel off the fp8 bytes
    with no bf16 dequant mirror. A payload the fused kernels cannot read
    (an EXL3 trellis, a W4A8 pre-scale) keeps per-expert modules and rides
    the grouped-expert component instead;
  * always-on shared expert ``mlp.shared_expert`` gated by
    ``sigmoid(mlp.shared_expert_gate(x))`` (Qwen2-MoE / Qwen3-Next
    convention), a dense gated MLP outside the fused path.

**Expert parallelism.** At ``--ep-size N`` each rank holds a contiguous
slice of ``num_experts // ep_size`` experts (the weight map binds only
the local slice); the router runs replicated, ids are shifted into the
local frame, slots belonging to another rank contribute zero, and one
trailing all-reduce over ``_EP_GROUP`` sums every token's full top-k
mixture. Attention / GDN / embeddings stay replicated within each EP
shard.

**Capture.** The fused dispatch is sync-free and fixed-shape — the
token→expert binning stays on device end to end — so decode and
piecewise capture apply to this arch like any other.

**Speculative decode.** The checkpoint bundles a one-block MTP draft
head under ``mtp.*`` whose decoder block is itself MoE (router + full
routed-expert set + gated shared expert) on top of the same gated
full-attention mixer. It is built as
:class:`~arbi_serve.spec_decode.mtp_head.Qwen3_5MtpHead` with a
:class:`Qwen3_5MoeSparseMoeBlock` FFN, so its experts ride the same
:class:`~arbi_serve.models.moe.FusedMoE` stack, EP slice and trailing
all-reduce as a main layer's. Under ``--ep-size`` every rank runs the
drafter in lockstep (the driver broadcasts ``draft``), which is what
makes the head's EP all-reduce well-formed; no out-of-band collective
is introduced.

Not yet wired (fail-soft, text-only): the vision tower
(``model.visual.*``) — its checkpoint keys are ignored by the weight map
and the quant swap.

Per-arch construction bundle for :class:`Qwen3_5MoeModel`.

Args:
    gdn_cfg: shared per-layer GDN shape (one config across the stack).
    rotary_dim: partial-RoPE channel count
        (``head_dim * partial_rotary_factor``).
    moe: routed-expert shape (``num_experts`` / top-k /
        ``moe_intermediate_size``) shared by every layer.
    shared_expert_intermediate_size: width of the always-on shared
        expert MLP (0 → no shared expert).
    experts: routed-expert representation resolved from the pack by
        :func:`~arbi_serve.models._moe_dispatch.resolve_expert_dispatch`
        — the stacked :class:`FusedMoE` and its quant parameters, or
        per-expert modules.
    mtp_attn: attention shape (``num_heads`` / ``num_kv_heads`` /
        ``head_dim``) of the bundled MTP draft head, or ``None`` when
        the checkpoint ships no ``mtp.*`` keys. The head mirrors the
        main model's full-attention layers but is constructed
        standalone, so it needs the shape spelled out.

Routed experts as per-expert modules — the non-stackable-pack path.

Integration point for the shared grouped-expert component: per-expert
gated MLPs keyed by GLOBAL expert id, driven by one grouped trellis
GEMM per stage on an EXL3 pack and by the per-expert module walk on any
other payload the fused kernels cannot read.

Fail loud unless the bundled ``mtp.*`` head is the head we build.

Called only when the checkpoint carries ``mtp.*`` keys. Every check
guards a way the head could differ from
:class:`~arbi_serve.spec_decode.mtp_head.Qwen3_5MtpHead` +
:class:`Qwen3_5MoeSparseMoeBlock` such that loading it would still
"succeed" while drafting from the wrong weights: a multi-block head,
dedicated drafter embeddings, a dense FFN, a different expert count /
expert quant format, or a shared expert on one side only. Spec-decode
silently falling back to non-spec decode is not an option — an
unrecognised head is a checkpoint we do not know how to serve.

Router + routed experts + gated shared expert.

Numerics mirror HF ``Qwen3_5MoeSparseMoeBlock`` exactly::

    probs   = softmax(gate(x), dim=-1, dtype=float)
    w, ids  = topk(probs, top_k);  w /= w.sum(-1, keepdim=True)
    routed  = Σ_k w_k · expert_{ids_k}(x)
    shared  = sigmoid(shared_expert_gate(x)) · shared_expert(x)
    y       = routed + shared

``experts`` is a single :class:`FusedMoE` holding this rank's local
expert slice as one stacked ``(E_local, 2N, H)`` / ``(E_local, H, N)``
weight pair, driven by the fused Triton kernel. The per-expert
checkpoint tensors are bound straight into their slot of that stack
(see :func:`arbi_serve.models._qwen3_5_weights.build_moe_weight_map`),
so an fp8 pack is served off its fp8 bytes with the checkpoint's
``weight_scale_inv`` grid handed to the kernel unchanged. The router
(``gate``) and ``shared_expert_gate`` stay dense bf16
(``modules_to_not_convert`` in the reference pack).

Holds the ``model.language_model.*`` subtree (as ``model.*``).

``layer_specs`` must exclude any MTP-marker spec — the bundled draft
head owns its own decoder block, and a marker here would create a
main-model layer the checkpoint has no keys for.

Bind the MTP head block's sparse-MoE FFN keys.

Handed to :class:`Qwen3_5MtpHead` as ``mlp_weight_map`` so the
object that owns the FFN module also owns its binding: a head that
constructed cannot then silently fail to load, because there is no
second site an arch could forget to update.

Routed through the same builder the main decoder layers use, so
the head's expert stack lands on the identical fused-buffer layout
and the identical per-rank EP slice — at ``--ep-size 2`` each rank
binds only experts ``[rank*128, +128)`` of the head, exactly as it
does for a backbone layer.

Bound late (called at ``weight_map()`` time, not construction), so
it reads the parallel/quant config in force when the load happens.

MTP-layer KV-fill pass — run the head's decoder block over the
same flat batch the main model just saw.

Reuses ``batch.attn_meta`` and ``batch.positions`` but a different
``state_view`` / ``attn_op`` (the MTP marker spec's slab), so the
head's per-request KV history is written at exactly the slots the
main model wrote to and the next draft step reads real history.
The block's output is discarded — only the K/V scatter matters.

Runs only when the model carries a head and ``batch.attn_meta`` is
populated (the PAGED_KV path); profiling forwards built outside the
engine skip it safely, and the capture layer selects the baked
variant through ``batch.mtp_fill_enabled``.

EAGLE/MTP next-token shift: the block at flat slot ``i`` is driven
with ``embed(token[i+1])`` fused with ``hidden(i)`` — the pairing
the head was trained on and the one
:meth:`Qwen3_5MtpHead.forward` uses at draft time. Each request's
last in-span slot has no in-batch successor, so it stays
self-paired; that frontier cell is repaired from the corrected
``(embed(bonus), hidden)`` pair by
``MtpDriver._repair_kv_frontier`` before every draft chain.
Tensor-only — no host round trip, so the fill captures.

Embed, run the interleaved GDN/attention MoE layer stack, and
return last-token logits (optionally with the full hidden state).

``compute_logits=False`` (honoured only together with
``return_hidden_state=True``) skips the final
``lm_head(last_hidden)`` GEMV and returns ``(None, hidden)``: the
MTP verify pass recomputes ``lm_head`` over all K+1 flat rows
itself, so the last-token head would be a full-vocab GEMV
discarded every spec step.

Build a Qwen3.6-MoE model from an HF checkpoint directory.

Parses the nested ``text_config`` (same GDN/attention interleave
parser as the dense arch, ``model_type="qwen3_5_moe"``), attaches
one shared :class:`MoEConfig` to every layer spec, appends the
bundled MTP head's marker spec when the checkpoint carries one,
then runs the standard quant-swap + meta-materialize +
weight-load sequence. Each EP rank binds only its expert slice —
for the main layers and the MTP head (see
:func:`arbi_serve.models._qwen3_5_weights.build_moe_weight_map`).

Qwen3-MoE — ``Qwen3MoeForCausalLM`` (Qwen3-30B-A3B, Qwen3-235B-A22B).

Dense Qwen3 attention (per-head q/k RMSNorm, full RoPE, no qkv bias)
with the gated MLP replaced by a sparse MoE FFN on the sparse layers::

    probs   = softmax(gate(x), dim=-1, dtype=float)
    w, ids  = topk(probs, num_experts_per_tok)
    if norm_topk_prob:  w /= w.sum(-1, keepdim=True)
    y       = Σ_k w_k · expert_{ids_k}(x)

There is **no shared expert** (that is Qwen2-MoE), no routed-expert
score bias, no routed scaling factor and no grouped routing.

Layer sparsity follows HF: layer ``i`` is MoE iff ``i not in
mlp_only_layers`` and ``(i + 1) % decoder_sparse_step == 0``. Every
released Qwen3-MoE checkpoint sets ``decoder_sparse_step=1`` and
``mlp_only_layers=[]``, so every layer is MoE; the dense-layer branch
exists because the config expresses it.

``head_dim`` is read from ``config.json`` and NOT derived — Qwen3-30B-A3B
has ``hidden_size=2048``, ``num_attention_heads=32`` and
``head_dim=128``.

**Expert parallelism.** Each rank holds a contiguous
``num_experts // expert_shard_count`` slice; the router is replicated,
ids outside the local frame contribute zero, and the trailing TP/EP
all-reduce pair sums every token's full mixture. Same fused stacked
dispatch as every other MoE arch here, so decode capture applies.

Fail loud on an expert pack the stacked loader cannot bind.

The routed experts are read as one stacked tensor per layer, so the
generic per-linear quant swap never sees them — an unrecognised expert
format would otherwise load quantized bytes as if they were bf16.

Serving model for ``Qwen4ExpForConditionalGeneration``.

The vision half is the shared Qwen-VL tower: Qwen4-Exp's vision
config and vision model are zero-override subclasses of Qwen3.5-MoE's,
so ``self.visual`` is built by
:func:`arbi_serve.models.vision.base.build_vision_tower` and image
inputs enter through the modality-agnostic ``mm_bindings`` seam.
It is bound only when the checkpoint carries a ``vision_config`` AND
``ARBI_ENABLE_VISION`` is on; a text-only run leaves ``self.visual``
``None`` and allocates none of it.

One :class:`LayerSpec` per MTP stage the config declares.

``mtp_num_hidden_layers`` and ``mtp.layer_types`` name the stack; the
checkpoint carries it under ``mtp.layers.{i}``. Binding the intersection
would leave a stage's weights unread, so a disagreement raises.

Bind one ``mtp.layers.{stage}`` decoder stage.

The MTP stack carries its own ``mtp.layers.{i}`` namespace whose indices
restart at zero, so the stage number — not the decoder layer index —
names the source keys.

Tokens per KV page this checkpoint's sparse attention requires.

QSA selects whole compressed index blocks, and a block addresses its
tokens through the page table, so one index block has to BE one page.
Any other page size makes a selected block span or subdivide pages and
the selected token set stops being expressible as a block table.

Derive PLE row IDs on host and start the deduplicated gathers.

Each row's n-gram context comes from its request's fed token stream
(:meth:`arbi_serve.engine.request.Request.fed_token_ids`), keyed by
the row's first RoPE position. That makes the derivation identical
whatever the scheduler did to reach this row — chunked prefill,
preemption and re-prefill, a prefix-cache or savepoint hit, a resumed
suspension, or a mid-decode context injection — and it raises rather
than approximate a history it cannot name.

Derive PLE rows for a request-free batch, every row starting fresh.

Only the synthetic boot forwards reach this. The token ids come off
the batch itself, and the host copy it costs is paid once at boot
rather than on a serving step.

Partial-RoPE tables for the QSA queries and for every block start.

The indexer rotates its query at the token's own position and each
pooled block key at the position of the block's FIRST token. With
1-D positions that first position is the block column times the
compression ratio, one table for every row. Under M-RoPE a position
is three channels that depend on where the image sits, so the block
start's position is whatever was fed at that page's first slot: the
step records each token's position under its slot, and the block
table's pages read their start positions straight back out. Rows
therefore get a table each.

Attend every query token over its own QSA-selected pages.

The step is one sparse call: the chunk's K/V and index keys are
scattered once, the selector scores every query token's complete
blocks in one batched pass, and each token becomes one row of a
``block_topk + 1`` wide sparse block table whose ``last_page_len``
caps its trailing partial page at its own position.

Re-pool every QSA layer's pooled key for the given slots' pages.

The per-page pooled index key is derived from the page's index keys,
and the layer only re-pools what its own per-token write touches. A
pool-level mover that rewrites index-key slots underneath the layer
— a tree accept compacting its scattered path — therefore has to say
so here, or the selection scores those pages with keys they stopped
holding and nothing raises.

Bind the payloads the generic weight loader does not carry.

The PLE table binds on EVERY path, warm flat-weight restore
included. Its shards stay mmap-backed in host memory and are
deliberately never written into a flat dump — a warm boot restores
device weights and would otherwise leave the n-gram table unbound,
which surfaces later as a failure to size the offload staging.

The ``context_len`` tokens fed immediately before position ``start``.

PLE hashes each token together with the ``ngram_size - 1`` tokens that
preceded it, so a row beginning at position ``start`` needs the fed token
ids at positions ``[start - context_len, start)``. The request's fed
stream — prompt, output, and consumed context injections in forward
order — is the only exact source for that map: chunk boundaries,
preemption rewrites, prefix-cache hits, and mid-decode context injection
all move ``start`` without changing it.

Left-pads with ``eos_token_id`` when fewer than ``context_len`` tokens
precede ``start``, matching the EOS padding applied to a first chunk
shorter than the n-gram context.

Raises when the request cannot name every token it has already fed
(a parked speculative or async-output advance), because the only
alternative is an approximate history.

Multimodal half of the checkpoint: tower plus its token contract.

``tower`` is the shared Qwen-VL :class:`VisionConfig` — Qwen4-Exp's
vision tower is the Qwen3.5-MoE tower with no overrides, so it is
built by :func:`arbi_serve.models.vision.base.build_vision_tower`
from the same parse the other Qwen-VL checkpoints use. The token ids
mark the placeholder rows the tower output overwrites and the
``<vision_start>``/``<vision_end>`` delimiters that bracket them.
``mrope_section`` is the ``(t, h, w)`` split of the rotary
frequencies driving interleaved M-RoPE.

The bundled multi-token-prediction stack the checkpoint declares.

``layer_types`` holds one normalised kind per ``mtp.layers.{i}`` stage,
so the loader binds exactly the stages the config names instead of
assuming the released checkpoint's single QSA stage.

Parse the multimodal half of the checkpoint, or ``None`` for text-only.

Returns ``None`` when the checkpoint carries no ``vision_config`` or
when ``ARBI_ENABLE_VISION`` is off — the same gate every other
multimodal arch uses, so a text-only run never allocates the tower.
A checkpoint that DOES carry a tower but describes it in terms this
serving path cannot honour raises instead of dropping vision.

Qwen Sparse Attention index selection reference.

Block scores pass through ``torch.relu`` before the head sum, so every block
a query suppresses on every head scores exactly ``0.0`` and a context wider
than the budget routinely offers more exactly-tied blocks than there are
slots. Selection resolves a tie by index: :func:`tkv.runtime.sparse_topk_ids`
orders score descending and index ascending, and the lowest-index block is
the one kept when a tie straddles the budget. Upstream selects with
``scores.topk(...)``, whose tie-break follows its own partitioning and
therefore differs between CPU and CUDA and with the number of tied blocks
competing. A tie is a group the indexer scored identically, so index order
resolves it as faithfully as any other rule and, unlike the others, it makes
the selection reproducible and safe to replay from a captured graph.
``tests/test_qwen4_exp_qsa.py`` pins the contract on both selectors.

Return selected token IDs padded with ``-1``.

``query`` has shape ``(B, Q, H, D)``, ``raw_keys`` is
``(B, K, D)``, and ``visible`` is the causal/sequence-validity mask
``(B, Q, K)``. Complete visible blocks compete for the budget and the
incomplete trailing block is always retained. Selected blocks are laid
out score descending and index ascending among equal scores, and the
retained tail's tokens follow them.

Return the pooled index key of every page in ``pages``.

``index_view`` is the paged index-key stream
``(pages, compress_ratio, head_dim)`` and the result is
``(len(pages), head_dim)``: a page's whole index-key block averaged
and normalized, which is what a complete block scores with.

Re-pool every page the KV slots in ``slots`` write into.

``pooled_keys`` is the layer's ``(pages, head_dim)`` per-page pooled
index keys and is derived state: it is only ever read by
:meth:`select_paged_blocks`, and it is only correct while it equals
:meth:`pool_pages` of the pages it covers. Every write into
``index_view`` therefore has to pass its slots through here.

Slots sharing a page recompute the same row from the same already
written page, so the duplicate scatter is value-identical. Shapes
follow ``slots``, which makes this safe to capture.

Return the selected complete-block columns for every query token.

The batched twin of :meth:`select` on the paged index-key layout,
where one index block IS one KV page: ``block_table`` is
``(num_rows, width)`` and ``pooled_keys`` is ``(pages, head_dim)``,
so block ``w`` of batch row ``r`` scores with the pooled key of page
``block_table[r, w]``.

A page's pooled key is a function of the page's index keys alone, so
it is pooled once by :meth:`repool_written_pages` when the page is
written and read back here — a full page's key never changes again,
and the trailing partial page is repooled by the same write that
extends it.

``query`` is ``(tokens, heads, head_dim)`` and every other per-token
argument is ``(tokens,)``: ``complete`` is a token's count of
complete visible blocks and ``pad_slot`` its slot in the
``(num_rows, query_stride)`` grid the score matmul groups queries
into (``num_rows * query_stride`` is the spare slot for a token
outside the grid). ``block_cos`` /
``block_sin`` are the partial-RoPE tables at every block's START
position: ``(width, rotary_dim)`` when one position per block serves
every row, and ``(num_rows, width, rotary_dim)`` under M-RoPE, where
a block's three position channels depend on whose page it is.

The result is ``(tokens, block_topk)`` block columns, score
descending and index ascending among equal scores, padded with
``-1``. The always-retained incomplete tail is NOT part of it: it is
the trailing partial page, which the caller appends.

Batched sparse paged-attention metadata for Qwen Sparse Attention.

QSA pins the paged-cache page size to the indexer compression ratio, so one
index block IS one KV page: complete block ``w`` of a batch row is the page
``block_table[row, w]``, and the always-retained incomplete tail is exactly
the trailing partial page. A query token's selected token set is therefore
``block_topk`` whole pages plus that partial page, which makes the whole
step one decode-shaped sparse call: every query token is its own row, one
query per row, ``causal=False``, and a fixed ``block_topk + 1`` wide block
table. Each row's ``last_page_len`` caps its partial page at its own
position, so a prefill chunk may scatter its K/V once up front and still
attend causally.

Every per-row length is derived from the step's device tensors with no
device-to-host copy, and the published block table, CSR triplet and length
vectors live in buffers sized once and reused, so a captured decode graph
replays against stable addresses.

Per-query-token geometry of one QSA step, all device-resident.

``row_seq`` maps a query token to its batch row, ``visible`` is the
token's causally visible KV length, ``complete`` the number of whole
index blocks inside it, and ``tail`` the length of the trailing partial
page. ``pad_slot`` places the token in the ``(num_rows, query_stride)``
grid the batched selector groups queries into; ``num_rows *
query_stride`` is the spare slot reserved for a token that falls
outside the grid.

Persistent buffers for the QSA sparse call, plus its no-sync build.

One instance serves every QSA layer of a step: the layers run serially
and each publishes its metadata to the attention op before the next one
builds. Buffers grow to the widest step seen and are never reallocated
afterwards; a growth request while a CUDA graph is capturing is refused
rather than silently invalidating the addresses that graph baked in.

Publish one decode-shaped sparse call over the selected pages.

``chosen`` is the ``(tokens, block_topk)`` selector output: complete
block columns in score-descending order, padded with ``-1``. Each
row's block table is those columns' pages followed by its trailing
partial page, which is what ``last_page_len`` then caps.

Shared recurrent-mixer host-side helpers.

These helpers resolve per-step metadata fields — and advance the
masked-replay conv window — the same way for every recurrent block
(Mamba-2 :class:`~arbi_serve.models.mamba2_block.Mamba2Block`, ShortConv
:class:`~arbi_serve.models.short_conv_block`, and Gated Delta Net
:class:`~arbi_serve.models.gdn_block.GDNBlock`). They live here — not on any
one block — because they are arch-agnostic and cross-imported by all of them.

Return a per-row ``has_initial_state`` mask, defaulting to all-True
when ``meta.has_initial_state`` is ``None``.

The default-True path covers the existing test fixtures that
pre-allocate state via :meth:`RecurrentStatePool.alloc_for_request`
and expect the block to read ``state_view.X[slab_row]`` as their
seed state. Production callers (the engine's metadata builder) set
the field explicitly per row.

Return the int64 slab-index tensor for ``index_select`` /
``index_copy_`` over the slab.

Prefers ``meta.state_indices_long`` (the persistent capture buffer
pinned by the engine on the captured-decode + piecewise-prefill
paths). Falls back to ``meta.state_indices.to(int64)`` on the eager
path where no persistent buffer exists.

Capture-safety: a per-call ``.to(int64)`` allocates a fresh tensor on
every forward. Under stream capture, the captured kernel launches
bake the fresh tensor's ``data_ptr`` at capture time; the tensor is
freed immediately after, and replays reallocate at different
addresses. The captured kernels then read invalid memory —
illegal-address crash or silent state corruption (also explains
the temp=0 non-determinism observed under cudagraphs). Threading
a persistent int64 buffer through ``meta.state_indices_long``
pins the ``data_ptr`` for the captured kernels' lifetime.

In place, set every REJECTED step of each replay staging buffer.

``step_mask`` is the ``(N, T)`` bool ``t <= nacc`` prefix mask the pool
builds once per tick; each target is ``(buffer, fill)`` where the
buffer is ``(N, T, …)`` and ``fill`` is the value that makes a masked
step an exact identity for the kernel that consumes it (``0.0`` for
the delta-rule key — it kills the rank-1 update and the v-correction;
``-inf`` for the gate pre-activation — ``softplus`` goes to ``0`` so
the decay is ``exp(-0.0) == 1.0``).

One inversion of the mask serves every target, and each write is a
single ``masked_fill_``: the alternative spelling
(``copy_(where(mask, buf, fill))``) costs a scalar fill, an
out-of-place select and a copy per buffer, which at one call per
recurrent layer is launch cost rather than work. Writing the fill
directly also cannot turn a stale NaN in a dead slot into a NaN in
the recurrence, where a multiply by a zero mask would.

Fixed shapes, no host reads — stream-capture-safe.

Commit each row's post-accepted-prefix conv window by pure gather.

The rolling conv buffer after committing tokens ``0..n`` is window
``n`` of ``[base_conv[..., 1:] ‖ x_0 .. x_{T-1}]`` — element copies of
the values a per-token replay would have rolled in, so no conv is
re-run. ``nacc_by_row`` is the per-SLAB-ROW accepted offset
(``-1`` = leave the row untouched).

Shapes: ``conv_slab`` / ``base_conv`` ``(N, C, Kw)``,
``replay_x_conv`` ``(N, T, C)``, ``nacc_by_row`` ``(N,)``.

Fixed shapes, no host reads — stream-capture-safe. Shared by
:meth:`GDNBlock.replay_masked_commit` and
:meth:`Mamba2Block.replay_masked_commit`: both slabs use the same
channels-first ``causal_conv1d_update`` rolling-buffer convention.

On CUDA this dispatches to a single-launch Triton kernel that does
the same addressing and reads only the ``Kw`` columns it emits
(:mod:`arbi_serve.kernels.conv_window_commit`). Both arms copy the
same source elements, so they agree byte-for-byte; the torch body
below stays the reference and serves CPU, the recompute-mode base
frame, and any dtype the kernel is not defined for.

Commit each row's conv window at an accepted TREE node, by gather.

The tree generalization of :func:`advance_conv_window_masked`, and the
same operation: both read a ``Kw``-wide window out of
``[base_conv[..., 1:] ‖ x_0 .. x_{T-1}]``. A chain's accepted set is a
PREFIX, so its window is ``Kw`` consecutive columns and a scalar offset
per row addresses it. A tree's accepted path is scattered through node
order, so the window is ``Kw`` columns picked along the node's own
ancestor chain and each row needs an index VECTOR — which is what
``window_index`` supplies, precomputed per geometry.

``accepted_node_by_row`` is the per-SLAB-ROW accepted node id
(``-1`` = leave the row untouched); ``window_index`` is the ``(N, Kw)``
table from ``arbi_serve.models._gdn_tree_scan.tree_conv_window_index``.

Shapes: ``conv_slab`` / ``base_conv`` ``(N_rows, C, Kw)``,
``replay_x_conv`` ``(N_rows, N_nodes, C)``,
``accepted_node_by_row`` ``(N_rows,)``.

Fixed shapes, no host reads — stream-capture-safe, like its sibling.

Register the persistent ``cos``/``sin`` buffers on ``module``.

When ``table_pool`` is given, only the final tables are allocated
inside it (via a clone). The fp32 trig temporaries the caller built
to produce ``cos_tbl``/``sin_tbl`` stay in the default allocator, so
they are not stranded in the pool — a private MemPool never returns
its freed segments (pytorch#145168), so anything transiently
allocated inside it inflates the pool's resident size for the process
lifetime. The clone is bit-identical to the input tables.

ShortConv mixer block — LFM2 / LFM2-MoE.

Architecture:

  in_proj : Linear(dim -> 3 * dim)         # produces (B, C, x)
  conv1d  : causal 1-D conv along token axis (kernel = conv_kernel)
  out_proj: Linear(dim -> dim)

Forward, per request, per layer (decode token):

    BCx       = in_proj(hidden)           # (N, 3*conv_dim)
    B,C,x     = chunk(BCx, 3, dim=-1)
    Bx        = causal_conv1d(B * x)      # state-aware causal conv
    y         = C * Bx                    # SiLU-style gating via C
    out       = out_proj(y)

When the ``causal-conv1d`` extension is available we dispatch to its
varlen prefill (``causal_conv1d_fn``) and decode-token-update
(``causal_conv1d_update``) kernels. When it is not, we fall back to a
reference torch path (slower; supports CPU for unit testing).

Per-layer state is owned by :class:`ShortConvStatePool`
(``StateKind.SHORT_CONV``) — the model never allocates a conv buffer of
its own. The block reads its layer view + per-request row indices off
:class:`arbi_serve.engine.batch.ShortConvMeta`.

Reference per-channel causal 1-D conv.

Returns ``(out, final_state)`` where ``final_state`` is the last
``width-1`` activations of the input — i.e. what the next call's
``initial_state`` should be.

This is a teaching / unit-test path; production runs use
``causal_conv1d_fn`` from the extension.

LFM2-style ShortConv mixer.

The block reads its per-request conv state from a
:class:`ShortConvLayerView` (handed in via the model forward) and
its per-request row / has-initial-state metadata from
:class:`ShortConvMeta`. It does not allocate state itself.

``in_proj`` and ``out_proj`` are TP-aware: ``in_proj`` is
column-parallel so the per-rank conv operates on a slice of the
channel dim; ``out_proj`` is row-parallel and all-reduces back to
the full hidden_size on the way out. At ``tp_size == 1`` both
short-circuit to plain matmuls.

Split the fused in_proj output into ``(B, C, x)`` per-rank shards.

At tp_size > 1, ``in_proj`` returns a per-rank fused tensor
whose layout is ``[B_local | C_local | x_local]`` along the
last dim (each shard is ``conv_dim_local`` wide).

Run one forward step for this ShortConv layer.

Returns ``(N_tokens, hidden_size)`` post-output-projection.

Dispatches:
  - ``meta.is_prefill`` and the cuda kernel is available →
    ``causal_conv1d_fn`` over the varlen prefill (with per-row
    ``initial_states`` seeded from ``state_view.conv_state``).
  - ``not meta.is_prefill`` and the kernel is available →
    ``causal_conv1d_update`` per-decode-token, with
    ``conv_state_indices`` pointing at each request's row.
  - kernel missing → :func:`_ref_causal_conv1d` per-row
    (slower; CPU + unit-test path).

GPU + kernel-available paths are routed through
``torch.ops.arbi_serve.short_conv_step`` so Inductor sees the
causal-conv1d call as a stable boundary at trace time. The
gate-split (``B``, ``C``, ``x``), the ``Bx = B * x`` multiply,
and the final ``out_proj`` stay outside the op — they're plain
ATen ops Inductor can fuse. CPU / kernel-missing paths stay on
the direct reference path; Inductor never sees those.

Call through ``torch.ops.arbi_serve.short_conv_step``.

The per-step ``cu_seqlens`` / ``state_indices`` /
``has_initial_state`` are threaded as real tensor args (not a
``self`` side-channel) so they flow through ``torch.compile`` /
piecewise capture correctly — see the op docstring in
:mod:`arbi_serve._custom_ops`. Returns the conv output
(shape == ``Bx.shape``).

Real-impl side of ``arbi_serve::short_conv_step``.

``Bx`` is the gated input (``B * x``) — the conv kernel input.
``cu_seqlens`` / ``state_indices`` / ``has_initial_state`` are
the live per-step tensors threaded through the op signature
rather than a ``self._call_meta`` side-channel, which would go
stale under compile/capture (see the op docstring).

Returns the conv output of shape ``(N, conv_dim_local)``,
matching :meth:`_short_conv_step_fake`'s
``torch.empty_like(Bx)`` shape contract.

Run varlen prefill via :func:`causal_conv1d_fn`.

Shared by the B=1 capture-safe path and the B>1 eager loop.
Writes per-row final activations back into the slab via
``index_copy_`` (capture-safe at B=1) or per-row ``copy_``
(eager B>1). ``cu_seqlens`` / ``state_indices`` /
``has_initial_state`` are the live per-step tensors (op args),
not a stashed meta.

Step-Audio-2 — ``StepAudio2ForCausalLM`` (stepfun-ai/Step-Audio-2-mini).

The LM backbone is a stock Qwen2 dense decoder: Llama topology (no
per-head q/k-norm) **plus a bias on the Q/K/V projections** (``o_proj``
and the MLP stay bias-free). ``config.json`` nests the LM shape under
``text_config`` and the audio-encoder shape under
``audio_encoder_config``.

Topology:
  Embedding → [PreNorm → AttentionBlock(qkv_bias) + residual → PreNorm →
              GatedSiLUMLP + residual] × 28 → final RMSNorm → lm_head.

The checkpoint also carries a Whisper-style audio encoder
(``encoder.*``) and a Conv1d+MLP adapter (``adapter.*``) that project
16 kHz audio into LM embedding space at ≈12.5 tokens/s, scattered onto
``<audio_patch>`` placeholder positions. The audio tower loads only when
``ARBI_ENABLE_AUDIO`` is set (mirrors the vision gating); text-only
boots skip those tensors entirely.

Output side: the model emits one interleaved AR stream — text ids
(< 151688), control ids (151688–151695) and speech codes
(id − 151696 ∈ [0, 6560], 25 codes/s). :meth:`output_modalities`
declares the id ranges so the engine's demux can split the stream; the
CosyVoice2 token2wav decoder turns the code stream into 24 kHz audio.

Step-Audio-2 LM shape config (the nested ``text_config``).

The HF config nests the Qwen2 shape fields under ``text_config``, so
:meth:`from_hf_config` lifts that sub-dict into the shared
:class:`DenseDecoderConfig` fields. Uniform full-context attention;
Q/K/V projections carry a bias (wired in the decoder layer, not
here — this is purely the shape lift).

One Qwen2 transformer block: pre-norm + attn(+QKV bias) + pre-norm + MLP.

Identical to :class:`LlamaDecoderLayer` except the attention block
carries a Q/K/V projection bias (``qkv_bias=True``); like Llama it
has no per-head q/k-norm.

Lift the nested ``text_config`` into the shared dense fields.

Fails loud on knobs this path has not been validated against
(rope scaling, sliding window) so a mis-served checkpoint
surfaces at boot.

Declare the interleaved output stream's token geometry.

Step-Audio-2 emits one AR stream: text ids (< 151688), control
ids (151688–151695) and speech codes (id − 151696 ∈ [0, 6560])
at 25 codes/s; a speech turn is seeded with ``<tts_start>``.
The engine's per-request demux and the CosyVoice2 token2wav
decoder consume this spec.

Route-side input bindings, without a materialized model.

Placeholder expansion (``multimodal.openai.build_media_prompt``) reads
only ``placeholder_token_id`` and ``extra``; ``encode`` runs in the
engine process, where the tower lives. So the process-mode API child —
which holds no model object — resolves the wiring from ``config.json``
alone and ships preprocessed features over the wire.

``encode`` / ``num_placeholder_tokens`` are ``None`` here on purpose: a
binding that reached the runner must raise, never silently mis-encode.

The output-stream geometry, without a materialized model.

The spec is a property of the architecture, not of the loaded weights,
so the process-mode API child (which holds no model object) resolves it
straight from the registry — see ``models.omni_output_spec_for_dir``.

``model_dir`` is unused here (Step-Audio-2's id geometry is fixed by the
architecture) but is part of the registry contract: a future omni model
whose ranges depend on its config reads them from the dir.

Map each parameter to its safetensors source key and TP shard spec.

Qwen2 HF naming — Llama layout plus the q/k/v projection biases.
The checkpoint's ``encoder.*`` / ``adapter.*`` (audio tower)
tensors are covered by the tower's own sub-map when the audio
path is enabled; a text-only boot leaves them unread.

Vision towers — one shared encoder, per-family heads/tails.

The vision analogue of the text side's ``ModelBase`` / ``AttentionBlock``
/ ``LayerStack`` seam: a single :class:`~arbi_serve.models.vision.config.VisionConfig`
drives the shared :class:`~arbi_serve.models.vision.layers.VisionEncoderLayer`
stack, and :func:`~arbi_serve.models.vision.base.build_vision_tower`
dispatches to the per-family tower (Qwen-VL, Gemma). Towers
satisfy the :class:`~arbi_serve.models.vision.base.VisionTower` protocol
and contribute their weights through the host text model's weight_map.

``VisionTower`` protocol — the one seam every vision encoder implements.

A vision tower turns preprocessed pixels into LM-embedding-space
tokens. Concrete towers (Qwen-VL, Gemma) assemble the shared
:class:`arbi_serve.models.vision.layers.VisionEncoderLayer` stack with
an arch-specific patch head, positional scheme, and merge tail; they
all satisfy this protocol so the text model integrates against the
protocol, not the family.

The contract is intentionally narrow — ``encode`` plus ``weight_map``
— mirroring :class:`arbi_serve.models.base.ModelBase` on the text
side. Weights load through the *host text model*'s ``weight_map`` (the
tower contributes its sub-map under the checkpoint's ``model.visual``
/ ``model.vision_tower`` prefix), so the tower never owns a loader.

Run the tower.

Args:
    pixel_values: packed patch features for all images in the
        step. Qwen layout: ``(sum_patches, C·tps·p·p)``; Gemma
        layout: ``(num_images, C, H, W)`` (or pre-flattened
        ``(num_images·num_patches, patch_feat)``) — the tower
        documents which it expects.
    grid_thw: ``(num_images, 3)`` int — per-image ``(t, h, w)``
        in patch units. Drives per-image attention windowing,
        positional embeddings, and the merge tail's token
        count.

Returns:
    ``(sum_vision_tokens, out_hidden_size)`` — vision tokens in
    LM embedding space, concatenated over images in input
    order, ready to scatter into the text embedding stream.

Per-parameter shard spec for the tower's weights.

``dst_prefix`` is the attribute path of the tower on the host
model (e.g. ``"visual"``); ``src_prefix`` is the checkpoint key
prefix (e.g. ``"model.visual"``). Returns replicated specs
(vision towers are not TP-sharded in the checkpoints we
target).

Generic, arch-agnostic vision-tower config — the ``LayerSpec`` of vision.

One :class:`VisionConfig` drives the shared
:class:`arbi_serve.models.vision.layers.VisionEncoderLayer` for every
vision family, the same way :class:`arbi_serve.models.layer_spec.LayerSpec`
drives the shared :class:`arbi_serve.models.attn.AttentionBlock` across
the text decoders. A new vision tower contributes a ``from_hf_*``
parser that fills these flags — never a fork of the encoder body.

The arch-specific *head* (patch embedding), *positional* scheme
(learned interpolated table and/or 2-D RoPE), and *tail* (patch
merger / multimodal projector) are selected by the string-tag fields
(``patch_embed_kind`` / ``merger_kind``) and built by the per-arch
tower in :mod:`arbi_serve.models.vision`. Everything between — the
transformer stack — is shared.

Backend-agnostic vision-encoder shape + behaviour flags.

Fields group into: *patch head*, *transformer stack*, *positions*,
and *merge tail*. The boolean / string-tag fields are the vision
analogue of :class:`LayerSpec`'s ``attn_output_gate`` /
``use_qk_norm`` / ``skip_kv_proj`` knobs — they let one encoder
body serve Qwen-VL (pre-norm-2, LayerNorm, 2-D RoPE, sequential
GELU MLP) and Gemma (4-norm sandwich, RMSNorm, q/k-norm, gated
GELU MLP) without branching the class.

Args:
    family: short tag used by the registry / weight-map dispatch
        (``"qwen_vl"`` | ``"gemma"``).
    in_channels: input image channels (3 for RGB).
    patch_size: spatial patch edge in pixels.
    temporal_patch_size: temporal patch depth (Qwen video = 2;
        images repeat the frame to fill it). 1 for still-only
        towers.
    patch_embed_kind: ``"conv3d"`` (Qwen: Conv3d over the
        ``(t,p,p)`` patch cube), ``"linear"`` (Gemma: Linear over
        a pre-flattened patch feature vector), or ``"conv2d_cls"``
        (CLIP: Conv2d patch grid + class token + interpolated
        absolute pos-embed — DeepEncoder's CLIP half; the head
        lives on the tower, this tag documents the scheme).
    hidden_size: encoder residual width.
    depth: number of :class:`VisionEncoderLayer` blocks.
    num_heads: attention heads in the encoder.
    head_dim: per-head width (``hidden_size // num_heads`` unless
        the checkpoint overrides it).
    intermediate_size: MLP inner width.
    mlp_kind: ``"gelu_seq"`` (fc1 → act → fc2) or ``"gelu_gated"``
        (down(act(gate) * up)).
    hidden_act: activation tag (``"gelu_tanh"`` | ``"gelu"`` |
        ``"silu"``).
    norm_kind: ``"layernorm"`` (affine LN with bias) or
        ``"gemma_rms"`` (Gemma plain ``* weight`` RMSNorm).
    block_layout: ``"prenorm2"`` (norm1→attn→+res ; norm2→mlp→+res)
        or ``"gemma_sandwich"`` (the four-norm Gemma block).
    use_qk_norm: per-head Q/K RMSNorm before attention (Gemma
        vision = True; Qwen vision = False).
    attn_bias: include bias on q/k/v/o projections (Qwen vision
        qkv carries bias).
    use_2d_rope: apply Qwen-style 2-D rotary to the encoder Q/K.
    use_pos_embed: add a learned, bicubically-interpolated absolute
        position embedding (Qwen3-VL).
    num_position_embeddings: row count of the learned ``pos_embed``
        table; its base grid edge is ``isqrt``. 0 when unused.
    spatial_merge_size: edge of the square patch group fused by the
        merge tail (Qwen = 2 → 2×2 patches per output token).
    out_hidden_size: tower output width — MUST equal the text
        tower ``hidden_size`` so merged tokens drop into the LM
        embedding stream unprojected.
    merger_kind: tail tag (``"qwen_merger"`` | ``"gemma_projector"``).
    soft_tokens_per_image: Gemma fixed token budget per image after
        average pooling. 0 for towers whose token count is
        data-dependent (Qwen: ``t*h*w / merge²``).

Parse a Qwen-VL ``config.vision_config`` block into a :class:`VisionConfig`.

Covers Qwen2-VL / Qwen2.5-VL / Qwen3-VL towers. Qwen3-VL adds the
learned ``pos_embed`` table (``num_position_embeddings`` present);
earlier towers omit it and rely on 2-D RoPE alone — both fall out
of the flags here.

Parse a Gemma ``config.vision_config`` block into a :class:`VisionConfig`.

Gemma's vision encoder layer is the same four-norm sandwich + per-
head q/k-norm + gated-GELU MLP as its *text* decoder layer, run
bidirectionally with no RoPE — so it maps onto the shared encoder
with ``block_layout="gemma_sandwich"`` and ``use_2d_rope=False``.
The patch head is a Linear over pre-extracted patch features and
the tail average-pools to a fixed ``soft_tokens`` budget per image.

DeepEncoder — the Unlimited-OCR (DeepSeek-OCR) dual vision tower.

``deeplip_b_l`` fuses two encoders over the same pixels:

  * **SAM-ViT-B** (``model.sam_model.*``): 16-px conv patch embed +
    interpolated absolute pos-embed, 12 blocks with 14×14 windowed
    attention (full attention at ``global_attn_indexes``) and
    decomposed relative-position bias, then a conv "neck" + two
    stride-2 convs downsampling 4× to 1024 channels — for a 1024²
    input the output is a ``(B, 1024, 16, 16)`` feature map.
  * **CLIP-L/14** (``model.vision_model.*``): consumes the SAM feature
    map as its patch embedding (its own conv patch head is bypassed at
    inference), prepends the class token, adds the bicubically
    interpolated absolute pos-embed, and runs 24 shared
    :class:`VisionEncoderLayer` blocks (fused qkv + bias, quick-GELU
    sequential MLP, pre-norm-2).

The two outputs fuse per patch — ``cat(clip_tokens[1:], sam_map)`` →
2048 channels — and a single linear projector maps into the LM
embedding space (1280). Per-image token assembly appends the learned
``image_newline`` embedding after each patch row and the
``view_seperator`` embedding after the global view; tiled (cropped)
views precede the global view in the output stream:

    [tile rows + newlines] [global rows + newlines] [view_seperator]

matching the reference ``masked_scatter_`` order in
``modeling_unlimitedocr.UnlimitedOCRModel.forward``.

Shape bundle for the dual tower, lifted from ``config.json``.

``vision_config.width`` carries the per-encoder geometry;
``projector_config`` sizes the fusion projector. The CLIP FFN width
is fixed by the reference ``vit_model_cfg`` (4096 — not present in
``config.json``; validated against the checkpoint tensor shapes at
load time by the loader's shape check).

CLIP absolute pos-embed for ``tgt_len`` tokens (incl. class token).

Port of the reference ``get_abs_pos``: the class row passes through;
the square patch grid is bicubic-antialias resampled in fp32 to the
target grid when sizes differ. Returns ``(tgt_len, hidden)``.

Relative-position rows for a (q_size, k_size) attention — SAM port.

Linearly resamples the ``(L, head_dim)`` table when ``L`` does not
match ``2·max(q,k) − 1`` (e.g. global blocks at a non-1024 input),
then gathers by relative coordinate.

SAM windowed/global attention with decomposed rel-pos bias.

Operates on ``(B, H, W, C)`` token grids (already window-partitioned
by the block when windowed). The additive attention bias is
``rel_h + rel_w`` from the per-axis relative-position tables —
materialised as an explicit SDPA ``attn_mask`` (the reference path).

SAM + CLIP dual encoder → fused, projected LM-space vision tokens.

``encode_views`` is the per-request entry the model's
:class:`~arbi_serve.multimodal.registry.MediaBinding` closure calls:
it consumes the preprocessed global views (+ optional crop tiles)
and returns the assembled ``(sum_tokens, n_embed)`` embedding
stream, newline/separator rows included, matching the placeholder
expansion the preprocessor declared.

Encode one request's views → ``(sum_tokens, n_embed)``.

Per image the output stream is ``[tiles, global, view_seperator]``
(tiles absent when the image was not cropped) — the order the
reference scatters embeddings into the placeholder positions.

Run the worst-case dummy encode; return the vision-token count.

The DeepEncoder worst case per image is fixed by the layout: one
``sam_img_size²`` global view plus the maximum tile grid at the
640² tile resolution (``dynamic_preprocess`` max 32 tiles). The
dual-encoder micro-batching bounds the per-chunk transient, so
one full-tile image alongside ``max_images`` global views covers
the serving-time peak (globals at 1024² dominate the per-chunk
transient; extra tile images only grow the small output
accumulation, which the assembled return covers).

Gemma-4 vision tower.

A faithful, parametrized reimplementation of ``transformers``'
``Gemma4VisionModel`` on our stack — verified to HF numerical parity
(``tests/test_vision_gemma.py``). The attention kernel is the only piece that
is not model code: the served path runs it through the owned mainline
(``VisionAttention`` / ``tkv``'s ``arbi_bf16_varlen``); the eager path here uses
SDPA (bit-comparable, and what the parity test pins against HF).

The Gemma vision layer is the same four-norm sandwich + per-head q/k/v-norm +
gated-GELU MLP as its *text* decoder, run bidirectionally. What is Gemma-4
specific — and implemented here rather than reused from the shared vision
machinery — is faithfully mirrored from the reference:

  * ``_ClippableLinear`` — QAT: clamp activations to learned in/out ranges;
  * ``_GemmaPatchEmbedder`` — ``2*(x-0.5)`` scale, a linear over pre-extracted
    patch features, and a ``(2, N, hidden)`` position-id embedding table;
  * ``_GemmaVisionRotary`` + ``_apply_multidim_rope`` — each spatial axis
    rotates its own ``head_dim // ndim`` channel band;
  * attention ``scaling = 1.0`` (the q/k RMSNorm carries the scale), v-norm
    with NO learned scale;
  * ``_GemmaVisionPooler`` — average-pool patches to a fixed soft-token budget,
    then ``* sqrt(hidden)``.

Linear that clamps its input and output to learned ranges (Gemma QAT).

The clamp bounds are buffers loaded from the checkpoint; unset they are
``±inf`` (a no-op), so a non-QAT checkpoint round-trips unchanged.

Rotate each spatial axis over its own channel band, then concat.

``x`` is ``(B, seq, heads, head_dim)``; ``cos``/``sin`` are ``(B, seq,
head_dim)``. Split all three into ``ndim`` equal bands and apply standard
RoPE per band — the reference's ``apply_multidimensional_rope``.

Bidirectional attention over the valid (non-padding) patches.

Gemma's is the only vision-specific step here — its image processor emits a
padded batch, so this packs the valid patches into the varlen form the rest
of the stack uses (strip padding; ``cu_seqlens`` = per-image valid counts,
keeping attention block-diagonal per image) and then defers to the shared
:func:`~arbi_serve.models.vision.layers.varlen_bidir_attention` — the same
arbi prefill kernel dispatch Qwen uses. The result is scattered back into the
padded layout; padding rows stay zero and the pooler drops them.

``q/k/v`` are ``(B, N, heads, head_dim)``; ``valid`` is ``(B, N)`` bool.

The CUDA served path defers to the shared kernel dispatch; the CPU/fp32
reference stays a batched SDPA with the padding mask, matching HF's exact
(batched) reduction so parity is bit-tight — the packed/unbatched form
differs by a ~1e-6 reduction-order step that the QAT clamp discontinuities
amplify across the stack.

Vision soft tokens → text embedding space (HF ``Gemma4MultimodalEmbedder``).

A scale-free RMSNorm precedes the bias-free projection: the pooled tower
output carries a large magnitude (the pooler's ``·√hidden`` scale), so the
pre-projection norm is load-bearing — a bare Linear would pass that
magnitude straight through and swamp the text stream. Module/param names
(``embedding_pre_projection_norm`` scale-free → no param;
``embedding_projection.weight``) mirror the checkpoint so the weight map is
a straight rename.

Gemma-4 vision encoder → per-image soft tokens in the text embed space.

``forward(pixel_values, pixel_position_ids)`` mirrors HF's
``Gemma4VisionModel.forward`` and is what the parity test pins. ``encode``
is the :class:`VisionTower` protocol entry the multimodal path calls.

Run the worst-case dummy encode; return the vision-token count.

Gemma's per-image worst case is the processor's fixed patch ceiling —
``soft_tokens_per_image · pooling_kernel_size²`` patches, every one
valid (no padding), which maximises the encoder's attention transient.
A near-square valid ``(x, y)`` grid (both axes multiples of the pooling
edge ``k``) keeps the pooler's block reshape exact; the actual pixel
values and position values are otherwise immaterial to the peak.

HF ``model.safetensors`` vision keys → our module tree (unsharded).

HF names: ``<src>patch_embedder.*``, ``<src>encoder.layers.<i>.*``,
``<src>pooler.*``. Ours drop the ``encoder.`` segment (layers hang off
the tower directly). Every vision weight is replicated (no TP shard) —
the tower runs on one rank and its output is broadcast.

Shared vision-encoder primitives — one body for every vision tower.

This is the vision analogue of :mod:`arbi_serve.models.attn` /
:mod:`arbi_serve.models.layers`: a single :class:`VisionEncoderLayer`
(plus the attention / MLP / norm / RoPE helpers it composes) covers
every family. A new tower selects behaviour through
:class:`arbi_serve.models.vision.config.VisionConfig` flags — it does
not subclass or copy the encoder body.

Design choices vs the text decoder primitives:

  * **Bidirectional, un-cached attention.** Vision encoders attend over
    a fixed set of patches with no causal mask and no KV cache, so
    :func:`varlen_bidir_attention` calls turbo-attn's arbi prefill
    kernel on flat (non-paged) K/V with ``causal=False`` and one
    ``cu_seqlens`` block per image, rather than going through the paged
    :class:`arbi_serve.backends.base.AttnOp` seam. On CPU it falls to
    SDPA with an additive block-diagonal mask.
  * **Plain ``nn.Linear`` / ``nn.LayerNorm``.** Vision towers are small
    (≈0.1–0.2 B), run in the model dtype, and are never TP-sharded or
    weight-quantized in the checkpoints we target, so they skip the
    ColumnParallel / quant machinery the text linears carry.

Return ``(turbo_prefill, BypassLoader)`` from turbo-attn.

``turbo_prefill`` takes ``q`` as ``(1, total_q, heads, head_dim)`` plus
``cu_seqlens_q``/``cu_seqlens_k``; ``BypassLoader.from_bf16`` wraps the
raw ``(1, total_k, heads, head_dim)`` K/V into the loader that picks the
kernel's no-decompression SMEM-fill body.

Gemma RMSNorm (fp32 interior). ``with_scale=False`` is the vision
v-norm (a bare normalize, no learned weight).

Local copy (not imported from :mod:`arbi_serve.models.gemma4`) to
keep the vision package import-cycle-free — gemma4 imports vision
to build its tower. Uses ``pow(ms, -0.5)`` (not ``rsqrt``) to match the
reference's Torch/JAX-portable form bit-for-bit.

Qwen-VL vision 2-D rotary embedding.

Each head channel range is split in half between the patch's row
(h) and column (w) coordinate. Half of ``head_dim`` carries the
``[h_freqs | w_freqs]`` rotation, then is duplicated to the full
``head_dim`` (NEOX layout). Matches
``Qwen2VLVisionRotaryEmbedding`` + ``apply_rotary_pos_emb_vision``.

Bicubically resample a square learned position table to ``(grid_h, grid_w)``.

``pos_table`` is ``(base*base, hidden)``; it is reshaped to its
square base grid, resampled to the target patch grid, and returned
row-major as ``(grid_h*grid_w, hidden)``. Reordering into the
encoder's merge-contiguous patch order (when the tower uses one) is
the caller's job via :func:`to_merge_order`.

Permutation: row-major ``(h*w)`` patch index → merge-contiguous order.

Qwen-VL groups each ``merge × merge`` spatial block of patches into
one output token, and the patch stream is laid out so those blocks
are contiguous. This returns the index vector that, applied to a
row-major ``(grid_h, grid_w)`` flattening, yields that
block-contiguous order. Used as the single source of truth for the
patch order so the preprocessor, the learned pos-embed, and the
2-D RoPE coordinates all agree.

Per-patch ``(hpos, wpos)`` coordinates in merge-contiguous order.

Mirrors ``Qwen2VLVisionTransformer.rot_pos_emb``: row / column
indices over the patch grid, reordered into the same
block-contiguous patch stream the encoder consumes, repeated over
the temporal dimension.

Bidirectional multi-head attention over a patch sequence.

Two projection layouts, selected by ``cfg.family``:

  * **fused** (Qwen, DeepEncoder CLIP): one ``qkv`` Linear
    (out = ``3·heads·head_dim``) + ``proj`` output Linear.
  * **split** (Gemma): separate ``q_proj`` / ``k_proj`` / ``v_proj``
    + ``proj`` output Linear, with per-head q/k RMSNorm.

A ``cu_seqlens`` argument makes attention block-diagonal so patches
of different images in the same packed batch never attend across
image boundaries (Qwen-VL windowing degenerates to per-image full
attention here).

One shared vision transformer block — pre-norm-2 or Gemma sandwich.

``prenorm2`` (Qwen-VL)::

    x = x + attn(norm1(x))
    x = x + mlp(norm2(x))

``gemma_sandwich`` (Gemma vision; same shape as the Gemma *text*
decoder block)::

    x = x + post_attention_layernorm(attn(input_layernorm(x)))
    x = x + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(x)))

Additive block-diagonal attention mask from per-image boundaries.

``cu_seqlens`` is the cumulative patch-count boundary tensor
(``[0, n0, n0+n1, …]``). Returns an additive ``(total, total)``
mask with ``0`` inside each image block and ``-inf`` across blocks,
or ``None`` when there is a single block (no masking needed).

Shared block-diagonal (per-``cu_seqlens``-segment) bidirectional attention.

The one attention-kernel dispatch for every vision tower — Qwen calls it
directly on its packed patches; Gemma packs its padded batch to varlen
(strip padding, ``cu_seqlens`` = per-image valid counts) and calls it too.
Two explicit, device-selected paths — no silent fall-through:

  * **CUDA** (production): turbo-attn's arbi prefill kernel
    (``turbo_prefill`` + ``BypassLoader``, ``causal=False``). It raises
    rather than degrading to masked SDPA on GPU.
  * **CPU** (fp32 reference / parity): SDPA with an additive block-diagonal
    mask.

Returns ``(seq, heads·head_dim)``.

Block-diagonal varlen attention → ``(seq, heads·head_dim)``.

Thin wrapper over the shared :func:`varlen_bidir_attention` so every
vision tower (Qwen here; Gemma after packing its padded batch to
varlen) dispatches attention through one kernel path.

Qwen-VL vision tower (Qwen2-VL / Qwen2.5-VL / Qwen3-VL).

Assembles the shared :class:`VisionEncoderLayer` stack with:

  * a Conv3d **patch head** over ``(temporal, patch, patch)`` cubes;
  * an optional learned, bicubically-interpolated **position
    embedding** (Qwen3-VL — earlier towers omit it);
  * 2-D rotary on the encoder Q/K;
  * a 2×2 **PatchMerger** tail (LayerNorm → group ``merge²`` patches →
    MLP) projecting to the LM ``out_hidden_size``.

The patch stream is laid out in merge-contiguous order (see
:func:`merge_order_perm`) so the merger's ``reshape(-1, hidden·merge²)``
groups true spatial neighbours; the pos-embed and RoPE coordinates use
the same order.

Interpolated learned pos-embed for one image, in merge order.

Corner-aligned bilinear gather from the square ``pos_embed``
grid, matching ``Qwen3_5VisionModel.fast_pos_embed_interpolate``:
sample at ``linspace(0, S-1, grid)`` points, blend the four
integer neighbours, then reorder row-major → merge-contiguous.

Encode packed Qwen-VL patches → ``(sum_vision_tokens, out_hidden)``.

Images are encoded **one at a time** and the merged outputs
concatenated in order. This is numerically identical to running
the whole packed batch through one block-loop: attention is
block-diagonal per image (``cu_seqlens`` never bridges an image
boundary — no cross-image attention) and every other op
(patch_embed, pos-embed add, RoPE, MLP, norms, merger) is purely
per-patch, so an image's output depends only on its own patches.

The win is the activation peak: the all-at-once residual stream
is ``(Σ_patches, hidden)``, so the peak scaled with the sum of
patches across the request (the multi-image OOM). Per-image
encoding frees each image's activations before the next, so the
peak collapses to the largest single image. The single-image
case is unchanged (one iteration, no concat).

Multimodal input plumbing — image/video preprocessing + merge contracts.

This package is the backend-agnostic seam between an OpenAI-style
request carrying images and the model's text-token embedding stream:

  * :mod:`arbi_serve.multimodal.inputs` — the dataclasses that carry
    preprocessed pixels (``MultiModalFeatures``) and the per-step
    merge plan (``MultiModalBatch``) that rides on
    :class:`arbi_serve.engine.batch.ScheduledBatch`.
  * :mod:`arbi_serve.multimodal.preprocess` — per-family image
    processors (Qwen-VL smart-resize/patchify, Gemma fixed-grid).
  * :mod:`arbi_serve.multimodal.rope` — Qwen M-RoPE 3D position index.
  * :mod:`arbi_serve.multimodal.merge` — scatter image embeddings into
    the embedded text stream at placeholder positions.
  * :mod:`arbi_serve.multimodal.registry` — arch → (processor, knobs).

The vision *towers* (the weight-bearing encoders) live under
:mod:`arbi_serve.models.vision`; this package never holds weights.

ASR / S2TT prompts and transcript post-processing for Step-Audio-2.

Two things live here because both are properties of the model's trained
convention, not of any one route.

The prompt matters more than it looks. Step-Audio-2 is an omni chat model;
"transcribe this" is not a mode, it is an instruction it was trained to
follow in one specific phrasing. The canonical Chinese instruction produces a
transcript; a reasonable-sounding English paraphrase can make the model
answer the content of the audio instead of transcribing it, silently turning
a transcription request into a question-answering one whose answer still
gets shipped to the client in a field labelled "transcript".

With the canonical prompt the model prefixes a language tag — ``<英语>``,
``<中文>``, or ``<非语音>`` for non-speech. These are ordinary text tokens
(below ``OmniOutputSpec.text_id_ceiling``), not control ids, so the omni
demux correctly routes them to the text lane and nothing strips them.
Callers should get a clean transcript plus a structured language field
instead.

Verbatim ASR instruction. ``None``/unknown → the generic canonical prompt
(auto-detect, best-effort); a supported ``source_language`` (en/zh/ja, or an
alias like ``"Chinese"``) → a language-pinned prompt that reliably stops SA2
translating and mis-tagging.

Canonical translation instruction for ``target_language`` (en/zh/ja).

Raises :class:`UnsupportedTranslationTarget` for anything else. There is
no silent fallback here: the model does not error on an unknown target,
it emits Chinese, so a fallback would ship confidently-wrong output under
the caller's requested language label.

``"<英语>Hello"`` -> ``("en", "Hello")``.

Returns ``(language, clean_text)``. ``language`` is ``None`` when there is no
recognized tag (the text is returned untouched), and ``""`` for ``<非语音>``
(non-speech detected -- a real answer, distinct from "unknown").

Only strips a tag we recognize. An unknown ``<...>`` prefix is left in place:
silently deleting model output we do not understand is how information gets
lost, and a caller seeing a stray tag will file a bug we can act on.

True when ``text`` could still grow into a language tag.

A streaming caller must hold back emission while this is True, or a partial
tag leaks: at ``text == "<"`` nothing matches yet, so a naive
"emit whatever is clean" loop ships the bare ``"<"`` to the client.

Dataclasses carrying multimodal media through the request pipeline.

Three surfaces:

  * Per-modality feature carriers — preprocessed, tower-ready media for
    *one request*, produced by the modality's preprocessor at submission
    time: :class:`MultiModalFeatures` (vision pixels) and
    :class:`AudioFeatures` (log-mel chunks). The engine ``Request``
    carries them as a ``dict[modality_name, features]`` in
    ``Request.multimodal`` (``None`` for text-only requests).
  * :class:`MultiModalBatch` — the per-step merge plan that rides on
    :class:`arbi_serve.engine.batch.ScheduledBatch`: the already-encoded
    embeddings for every media placeholder in the step plus a boolean
    mask marking which flat token positions are placeholders. The model
    reads it in ``forward`` to scatter media embeddings — the same plan
    shape for every modality (and for a mix), because towers run
    encode-once per request upstream and the scatter is placeholder-
    position-ordered regardless of which tower produced each row.

All plain data — no weights, no device assumptions beyond the tensors
they hold.

Tower-ready pixels for the images of one request.

Args:
    pixel_values: packed patch features for every image in the
        request, concatenated in prompt order. Qwen-VL layout:
        ``(sum_patches, in_channels·temporal·patch·patch)``.
    image_grid_thw: ``(num_images, 3)`` int — per-image
        ``(t, h, w)`` in patch units.

Tower-ready views for the images of one Unlimited-OCR request.

Every image contributes one padded global view; images larger than
the tile size additionally contribute an aspect-matched grid of
local tiles. The per-image placeholder count is layout-derived at
preprocess time (global rows + newlines + view separator
[+ tile rows + newlines]) — see
``DeepseekOCRImageProcessor._token_count``.

Args:
    global_views: ``(num_images, 3, base, base)`` float — padded
        global views, prompt order.
    crops: ``(sum_tiles, 3, tile, tile)`` float — local tiles of
        every cropped image, concatenated in prompt order; ``None``
        when no image was cropped.
    crops_per_image: tile count per image (0 = global view only).
    spatial_crop: per-image ``(w_num, h_num)`` tile-grid shape.
    per_image_token_counts: per-image LM placeholder-token count.

Tower-ready log-mel chunks for the audio clips of one request.

Each audio content part is split into fixed-length (≤25 s) chunks by
the preprocessor; every chunk becomes one ``<audio_start>…<audio_end>``
placeholder group in the prompt and one row of the packed mel batch.

Args:
    mels: ``(num_chunks, n_mels, T_max)`` float — log-mel
        spectrograms, right-padded to the longest chunk, prompt
        order.
    mel_lens: ``(num_chunks,)`` int — the encoder-mask length per
        chunk (Step-Audio-2 convention: ``true_frames − 2``).
    chunk_token_counts: per-chunk LM placeholder-token count (the
        adapter's output length; ``compute_token_num`` of the mel
        frame count).
    chunks_per_item: how many consecutive chunks each original
        audio content part contributed — the placeholder expansion
        groups ``chunk_token_counts`` by it.

One media item's placement in a request's expanded token stream.

Used by the request builder to record where each item's placeholder
block starts and how many tokens it spans, so logging /
introspection can map tokens back to media items.

Per-step multimodal merge plan attached to a :class:`ScheduledBatch`.

Carries **already-encoded** media embeddings for the placeholder
positions in *this step's* flat token stream. The towers run once
per request (cached on the ``Request``); the runner slices the
cached embeddings to whatever placeholder tokens fall in the current
(possibly chunked) prefill step — so the model forward only
scatters, never encodes, and a single media item can span prefill
chunks.

Args:
    embeds: ``(num_media_tokens_this_step, hidden)`` — the tower
        outputs for the placeholders present in this step, in
        stream order (already interleaved across modalities when a
        request mixes them).
    token_mask: ``(N_tokens,)`` bool over the step's flat token
        stream — ``True`` exactly at media-placeholder positions.
        ``token_mask.sum() == embeds.shape[0]``.

Media-content salting for prefix-cache keys.

The radix prefix cache (and the recurrent savepoint store) key on token
ids. Multimodal prompts expand every media item to runs of a single
placeholder id, so two requests with identical text but different
images / audio produce identical id streams — the cache would serve KV
computed from the wrong media. :func:`salted_cache_key_ids` builds a
same-length cache-key view of the prompt in which each media item's
placeholder positions carry a value derived from a content hash of that
item's preprocessed features:

  * identical media → identical salt → the prefix still shares;
  * different media → different salt at the first placeholder → the
    radix walk diverges there and the request re-prefills;
  * per-item hashing keeps multi-turn sharing: a new clip appended to a
    conversation leaves earlier items' salts (and thus the shared
    prefix) untouched.

Salt values are ≥ 2**63, far outside any vocab, so they can never
collide with real token ids. The view is used only for cache keying —
scheduling, prefill and the model always see the real
``prompt_token_ids``.

→ (per-item placeholder-token counts, per-item salt values).

Hashes the unpadded per-item feature bytes so the salt is stable
regardless of how the item packs against others in the request
(audio chunks are right-padded to the request's longest chunk;
identical clips must hash identically across requests).

Return the same-length cache-key view with salted placeholders.

Walks each modality's placeholder positions in prompt order,
assigning the i-th item's salt to its placeholder-token span. The
span lengths come from the features' own token accounting; a count
mismatch with the prompt raises loudly (the same class of bug the
encode-time guard catches).

Scatter media embeddings into the text embedding stream.

The one merge primitive shared by every multimodal model and modality:
take the token embeddings the text tower produced for ``input_ids``
(with media placeholder ids embedded to throwaway rows), and overwrite
the placeholder rows with the media tower's output, in order. This
keeps the merge identical across families and modalities — the only
things that vary are which token ids mark placeholders and which
tower(s) made the embeddings.

Overwrite placeholder rows of ``inputs_embeds`` with ``media_embeds``.

Args:
    inputs_embeds: ``(N_tokens, hidden)`` text embeddings.
    token_mask: ``(N_tokens,)`` bool — ``True`` at the placeholder
        positions to replace, in stream order.
    media_embeds: ``(num_media_tokens, hidden)`` tower output,
        ordered to match the ``True`` positions left-to-right.

Returns:
    ``inputs_embeds`` with placeholder rows replaced (out-of-place;
    the input is not mutated).

Raises loudly when the placeholder count and media-token count
disagree — a silent mismatch would shift every media token and
poison the whole sequence.

OpenAI chat → multimodal request glue.

Turns OpenAI ``content`` parts (text + ``image_url`` + ``input_audio``)
into the two things the engine needs: a ``prompt_token_ids`` stream with
each media item expanded to its full block of placeholder tokens, and
the per-modality preprocessed features dict that rides on the engine
``Request``. Supports multiple items per request, in prompt order.

Image sources accepted in ``image_url.url``:
  * ``data:image/...;base64,<...>`` data URLs (the common API form);
  * a local filesystem path (handy for tests / batch jobs).

Audio sources accepted:
  * ``{"type": "input_audio", "input_audio": {"data": <b64>,
    "format": "wav"}}`` — the OpenAI audio-in form;
  * the shorthand ``{"type": "audio", "audio": <path|bytes>}`` used by
    the Step-Audio-2 reference clients.

Network URL fetching is intentionally not done here — pulling remote
bytes is a policy decision (egress, SSRF) the server layer owns; pass
already-fetched bytes as a data URL / base64 payload.

Collect audio payloads from chat ``content`` parts, in prompt order.

Returns raw container bytes (for base64 ``input_audio``) or a local
filesystem path string (shorthand form) per item; decoding to
waveform is the preprocessor's job.

Expand each single placeholder token into its full vision block.

The chat template emits one ``image_token_id`` per image
(``<|vision_start|><|image_pad|><|vision_end|>``); the model needs
``per_image_token_counts[i]`` placeholder tokens for image ``i`` (its
post-merge vision-token count). Replaces the i-th placeholder with
that many copies, in order.

When ``boi_token_id`` / ``eoi_token_id`` are given (Gemma-4: the chat
template emits a bare ``image_token`` and the reference processor
wraps the run), each expansion is emitted as
``boi + image_token*n + eoi`` — the begin/end-of-image delimiters the
LM needs to read the soft-token run as an image. Delimiters are not
``image_token_id`` so they add no placeholder rows to the merge mask.

Raises if the number of placeholders in the stream does not match
the number of images — a mismatch would misalign every image.

Expand each single audio placeholder into its full per-chunk blocks.

The chat template emits one ``<audio_start><audio_patch><audio_end>``
group per audio content part. Long clips split into fixed-length
chunks upstream, each needing its own group; the i-th part's single
``patch_token_id`` marker expands to::

    patch×n₁ [end start patch×n₂ [end start patch×n₃ …]]

which, inside the template's existing ``start … end`` frame, yields
one complete group per chunk.

Raises if the number of markers does not match the number of
preprocessed audio items — a mismatch would misalign every clip.

Render → preprocess → expand into ``(token_ids, features_dict)``.

``rendered_token_ids`` is the chat-template output already tokenized
(one placeholder marker per media item). Walks the model's
``bindings``: for each modality with a processor supplied and media
present in the messages, preprocesses the items and expands the
markers. Returns the expanded token stream plus the per-modality
features dict, or ``(rendered_token_ids, None)`` when the request
carried no media.

Output-modality seam — per-model token-range declaration + per-request demux.

An omni model emits a single interleaved autoregressive stream mixing text
tokens, control tokens and speech codes. The model declares its id
geometry once (:class:`OmniOutputSpec`, via ``model.output_modalities()``
— text-only models return ``None`` / lack the method), and every request
served by such a model carries an :class:`OmniStreamState` that
``post_token`` routes each committed token through before the text
detokenizer sees it:

  * **text** ids append to the state's ``text_ids`` — the view the
    incremental detokenizer consumes (``Request.detok_ids()``);
  * **speech codes** append to ``audio_codes`` (already rebased by
    ``audio_code_base``) — the buffered code stream an AudioDecoder
    consumes, at the model's fixed code rate;
  * **control** ids flip the tts-segment flag and are swallowed.

The seam is deliberately wider than the first (interleaved-stream)
pattern it serves:

  * a *thinker–talker* model (Qwen3-Omni style) produces codes from a
    second model consuming the thinker's hidden states — its adapter
    feeds the same ``audio_codes`` buffer from a per-step hook instead
    of from the sampled-token demux, and everything downstream
    (decoder, SSE audio frames) is unchanged;
  * an RVQ detokenizer (LFM2-Audio style) emits tuples of codes per
    audio frame — ``codes_per_frame`` declares the arity so a decoder
    can regroup the flat buffer; interleaved single-codebook models use
    ``1``.

A model's output-stream token geometry.

Args:
    text_id_ceiling: ids strictly below this are ordinary text.
    control_ids: swallowed ids (never detokenized, never decoded)
        that structure the stream — audio/tts start/end/pad markers.
    audio_code_base: first speech-code id; ``code = id − base``.
    audio_code_count: number of speech-code ids (codes ≥ this are
        invalid and dropped with a warning by the demux).
    tts_start_id / tts_end_id: the control ids delimiting a speech
        segment (``tts_start`` also seeds a speech-output turn).
    speech_seed_text: string appended after the generation prompt to
        elicit speech output (``"<tts_start>"``).
    codes_per_frame: how many consecutive codes form one audio
        frame (1 = single-codebook interleaved; >1 = RVQ stacks).
    code_rate_hz: audio frames per second of the code stream.
    sample_rate: the decoder's output waveform rate.

The served model's :class:`OmniOutputSpec`, or ``None`` if text-only.

Two engine topologies, one answer. In thread/inline mode ``eng.model`` is the
live model and declares the spec. In process mode the API child holds no
model object (``ProcEngine.model is None``), so it carries a
registry-resolved ``omni_output_spec`` instead.

Without this fallback every speech-out request 400s (HTTP) / raises (realtime
WS) on an omni model whenever ``ARBI_ENGINE_PROC`` is on — reading ``model``
off a ProcEngine yields ``None``, and a missing spec is indistinguishable
from "this model has no audio out".

Per-request output demux state (engine thread only).

``route`` classifies one committed token id and updates the
buffers; the caller decides what to do with the class (text goes on
to the detokenizer, everything else stops here).

Per-family image preprocessors (PIL + torch; no torchvision dep).

Each processor turns decoded RGB images into tower-ready
``MultiModalFeatures`` (packed ``pixel_values`` + ``grid_thw``),
replicating the corresponding HuggingFace *slow* image processor so
the patches line up with the vision tower bit-for-bit.

Build the image processor for a model's declared ``family``.

Data-driven: resolves via :data:`_PROCESSOR_REGISTRY` and constructs with
``from_pretrained(model_path)``. Raises ``KeyError`` (naming the known
families) for an unregistered / ``None`` family — a mismatched preprocessor
silently garbles every image, so guessing is never right.

Unlimited-OCR (DeepSeek-OCR) image preprocessing — ``PIL.Image`` →
DeepEncoder-ready views.

Replicates the reference ``UnlimitedOCRForCausalLM.infer`` pipeline
(crop mode): every image gets a **global view** padded onto a
``base_size²`` canvas (mean-grey fill); images larger than the tile
size in either dimension additionally get **local tiles** from an
aspect-ratio-matched ``tile_size²`` grid split (``dynamic_preprocess``).
Pixel scaling is ``ToTensor`` (1/255) + mean/std 0.5 normalize.

The per-image placeholder-token count is fixed by the layout the tower
emits (see ``DeepEncoderVisionTower.encode_views``)::

    global:  (q_base + 1) · q_base + 1            # rows + newlines + view sep
    tiles:   (q_tile · w_num + 1) · (q_tile · h_num)

with ``q = ceil((size // patch) / downsample_ratio)`` — 16 for the
1024 global view, 10 for 640 tiles.

Stateless Unlimited-OCR image processor.

Construct from a checkpoint with :meth:`from_pretrained` (reads
``processor_config.json``), or pass the knobs directly.
``tile_size`` / ``min_tiles`` / ``max_tiles`` mirror the reference
``infer`` defaults (``image_size=640``, ``dynamic_preprocess``
min 2 / max 32) — they are not present in ``processor_config.json``.

Gemma-4 image preprocessing — ``PIL.Image`` → SigLIP2-style patch soup.

Replicates ``transformers.models.gemma4.image_processing_gemma4.Gemma4ImageProcessor``
(the torchvision backend the Gemma-4 checkpoint ships): aspect-ratio
preserving *resize into a patch budget* (largest size whose patch count
fits ``max_soft_tokens · pooling_kernel_size²`` and whose sides are
multiples of ``pooling_kernel_size · patch_size``), torchvision
bicubic-with-antialias resampling in the **uint8** domain, ``1/255``
rescale (identity normalize), patchify into the
``(num_patches, patch·patch·C)`` SigLIP2 layout, then pad every image to
the shared ``max_patches`` budget (pixels → 0, positions → ``(-1, -1)``).

Depends only on torch / torchvision / PIL — not on ``transformers`` at
runtime. The uint8-domain resize is load-bearing: resizing in float
drifts each normalized pixel by ~1/255.

One request's Gemma-4 image features — tower-ready, per-image padded.

``pixel_values`` is ``(num_images, num_patches, 3·patch·patch)`` and
``pixel_position_ids`` is ``(num_images, num_patches, 2)`` holding
per-patch ``(x, y)`` grid coordinates; padding patches carry
``(-1, -1)``. Both are padded to the shared ``max_patches`` budget,
but each image keeps its OWN valid patch grid — after the tower's
``k×k`` pooling drops padding, an image yields
``valid_patches // k²`` soft tokens, which VARIES with the image's
aspect-preserved resize (``soft_tokens_per_image`` is only the
ceiling, reached when an image fills the full patch budget).
``per_image_soft_tokens`` records the actual per-image count so the
prompt-side placeholder expansion matches what the tower emits.

Largest patch-budget-fitting, pool-aligned target size — matches HF.

Mirrors ``image_processing_pil_gemma4.get_aspect_ratio_preserving_size``:
scales to the ``max_patches`` pixel budget, floors each side to a
multiple of ``pooling_kernel_size · patch_size``, and rescues a side
that floored to zero.

Per-image placeholder counts — the ACTUAL pooled soft-token count
for each image (``valid_patches // k²``), not the fixed ceiling.

Gemma's placeholder expansion is layout-derived (like the DeepEncoder
path) and needs no ``merge_unit`` / ``grid_thw``, but the count is
per-image (aspect-dependent), so it is carried explicitly rather than
broadcast from a single budget.

Nemotron VoiceChat perception preprocessor: 16 kHz wav -> 128-bin log-mel.

Faithful port of NeMo's ``AudioToMelSpectrogramPreprocessor`` +
``FilterbankFeatures`` at THIS checkpoint's
``model.stt.model.perception.preprocessor`` config:

  * 16 kHz, ``window_size=0.025`` (400-sample Hann window, NeMo builds it
    with ``periodic=False`` — NOT the ``torch.hann_window`` default),
    ``window_stride=0.01`` (160-sample hop), ``n_fft=512``, 128 mel bins;
  * pre-emphasis 0.97 (``FilterbankFeatures`` default — not overridden in
    this checkpoint's preprocessor config, so the default applies);
  * ``dither=1e-5`` is a TRAINING-only knob in the reference
    (``if self.training and self.dither > 0: ...``) — inert at inference,
    so this module never dithers;
  * power spectrogram (``mag_power=2.0``), Slaney-normalized mel
    filterbank (computed here in closed form — bit-compatible with
    ``librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, norm="slaney")``
    to float32 noise, avoiding a librosa dependency on the serving path,
    same approach as ``step_audio2_audio.py``);
  * ``log(x + 2**-24)`` (``log_zero_guard_type="add"`` default);
  * ``normalize="NA"`` — NOT ``"per_feature"``: the reference's
    ``normalize_batch`` falls through to a no-op for any type string it
    does not recognize, so this checkpoint applies NO per-feature
    normalization (confirmed against the real config, not guessed);
  * ``pad_to=0`` — no padding of the mel time axis to a stride multiple.

Own feature dataclass (:class:`NemotronVoiceChatAudioFeatures`) — per
``docs/adding-a-model.md`` §7, this never lives in the shared
``multimodal/inputs.py``.

The mel-frame -> LM-placeholder-token count formula (:func:`token_count`)
duplicates the tower's subsampling arithmetic in pure Python (needed to
size prompt placeholders before any model forward exists) and MUST stay
in lockstep with ``_calc_length``/``_calc_length_int`` in
``dw_striding`` subsampling) — the tower's own ``encode()`` cross-checks
its actual output length against this file's precomputed count and
raises on any drift, so the two can never silently disagree.

``FilterbankFeatures.get_seq_len`` for this checkpoint (``exact_pad=False``).

``pad_amount = n_fft // 2 * 2`` cancels against the STFT's own
``center=True`` padding, leaving ``floor(num_samples / hop_length)`` —
NOT the raw STFT frame count (which is one frame larger near the
edge); the reference explicitly masks the extra edge frame to zero
rather than treating it as valid content, and this module follows suit.

Post-subsampling (factor 8, causal ``dw_striding``) frame count.

Must match ``_calc_length``/``_calc_length_int`` in
``models/audio/nemotron_voicechat_perception.py`` bit-for-bit
(``all_paddings=3, kernel_size=3, stride=2``, 3 repetitions —
``add_pad = all_paddings - kernel_size = 0``, so this is exactly
``floor(n/2) + 1`` applied three times).

Reference log-mel: ``(N_MELS, T)`` from a 16 kHz mono float waveform.

The tail beyond ``mel_length(wav.numel())`` is zeroed to mirror the
reference's ``x.masked_fill(mask_beyond_seq_len, pad_value=0.0)`` —
the tower's own subsampling re-masks from ``mel_lens`` too, so this
is defense-in-depth, not load-bearing, but keeps a standalone call to
this function bit-faithful to the reference on its own.

Incremental log-mel front-end for ONE live 16 kHz

The streaming counterpart of :func:`log_mel_spectrogram`, feeding
:meth:`~arbi_serve.models.audio.nemotron_voicechat_perception
.NemotronVoiceChatPerceptionTower.encode_stream`. Unlike the Conformer
stack above it, the mel front-end needs no learned cache: its receptive
field is BOUNDED and does not compound — one ``n_fft``-wide STFT frame
plus a single sample of pre-emphasis — so retaining a small trailing
raw-PCM carry reproduces the offline result exactly.

``center=True`` puts frame ``t`` at padded samples ``[160t, 160t+512)``,
i.e. raw ``[160t-256, 160t+256)``, so a frame is emitted only once raw
sample ``160t+255`` has arrived — a fixed 16 ms front-end delay on top
of the 10 ms frame grid, inherent to not knowing the future.

Output is byte-identical to :func:`log_mel_spectrogram` over the same
audio to float32 round-off (~1 ULP) and — importantly for a WS stream
whose chunk boundaries are arbitrary — is INVARIANT to how the caller
splits the audio: the same PCM produces the same frames whether pushed
in 30 ms or 200 ms pieces. Trailing frames that offline would
synthesize from the utterance's own zero padding are NOT emitted;
a live stream has no end to pad against.

One instance per connection; not thread-safe, not shareable.

Decode an audio payload -> 16 kHz mono float32 waveform ``(T,)``.

16-bit PCM WAV decodes natively (stdlib ``wave``); a non-16 kHz WAV
or any other container needs ``torchaudio`` — absent that, the error
says exactly what to install. Self-contained (no cross-family import)
per this codebase's per-vendor preprocessor convention.

Tower-ready log-mel features for the audio items of one request.

Each audio content part becomes one row of the packed mel batch and
one placeholder-token group in the prompt — the whole item goes
through the tower in a single forward call (no chunking; see
``models/audio/nemotron_voicechat_perception.py``'s module docstring
for why ``att_context_style="chunked_limited"`` makes manual chunking
unnecessary for offline/batch inference).

Args:
    mels: ``(num_items, n_mels=128, T_max)`` float — log-mel
        spectrograms, right-padded to the longest item, prompt order.
    mel_lens: ``(num_items,)`` int64 — valid mel-frame count per item
        (the encoder's pre-subsampling input length).
    token_counts: per-item LM placeholder-token count (the encoder's
        post-subsampling output length; :func:`token_count` of the
        mel-frame count).

Preprocess OpenAI audio parts -> :class:`NemotronVoiceChatAudioFeatures`.

Stateless (the filterbank is a module-level cache); safe to share
across requests. Runs on CPU in the HTTP process, mirroring the
image preprocessors and ``StepAudio2AudioProcessor``.

Qwen-VL image preprocessing — ``PIL.Image`` → tower-ready patches.

Replicates ``Qwen2VLImageProcessorFast`` (the processor the Qwen3.5
checkpoint ships): aspect-preserving *smart resize* to a patch-aligned
grid, ``torchvision`` bicubic-with-antialias resampling, rescale +
mean/std normalize, then patchify into the merge-contiguous
``(num_patches, C·temporal·patch·patch)`` layout the Conv3d patch
embed consumes. Depends only on torch / torchvision / PIL — not on
``transformers`` at runtime.

Step-Audio-2 audio preprocessor: 16 kHz wav → 128-bin log-mel chunks.

Byte-faithful port of the reference pipeline (the checkpoint's
``modeling_step_audio_2.log_mel_spectrogram`` + the reference client's
25 s chunking):

  * 16 kHz mono waveform, right-padded by 479 samples per chunk;
  * ``torch.stft(n_fft=400, hop=160, hann)`` — 100 mel frames/s, the
    last STFT frame dropped;
  * Slaney-scale mel filterbank (128 bins, fmin 0, fmax 8000,
    area-normalized) — computed here in closed form; bit-compatible
    with ``librosa.filters.mel(sr=16000, n_fft=400, n_mels=128)``
    within float32 arithmetic (guarded by a gated test);
  * ``log10`` → clamp to ``max − 8`` dynamic range → ``(x + 4) / 4``;
  * clips longer than 25 s split into consecutive 25 s chunks, each its
    own mel (and its own placeholder group in the prompt).

Token accounting: the encoder front-end (conv stride 2 → avg-pool 2)
plus the adapter conv (stride 2) put one LM audio token per ~80 ms;
:func:`compute_token_num` is the reference formula.

Container decoding accepts 16-bit PCM WAV natively (stdlib ``wave``)
at any sample rate (a vendored windowed-sinc polyphase resampler
handles non-16 kHz input); other containers route through
``torchaudio`` when it is importable and fail loud otherwise.

Reference formula: mel frame count → LM placeholder-token count.

Encoder: conv1 (k3 s1 p1, length-preserving) → conv2 (k3 s2 p1,
halves) → avg-pool (k2 s2, halves); adapter conv (k3 s2 p1). The
leading ``− 2`` mirrors the reference's mask-length convention
(``mel.shape[1] − 2`` is what feeds the encoder as ``x_len``).

Slaney-normalized triangular mel filterbank ``(n_mels, n_fft//2 + 1)``.

Closed-form equivalent of ``librosa.filters.mel(sr, n_fft, n_mels)``
at the library defaults (``htk=False``, ``norm="slaney"``, fmin 0,
fmax sr/2); float64 internally, float32 out — matching librosa's
arithmetic to ≤1e-8 per weight.

Decode an audio payload → 16 kHz mono float32 waveform ``(T,)``.

16-bit PCM WAV decodes natively (stdlib ``wave``); a non-16 kHz WAV
or any other container needs ``torchaudio`` — absent that, the
error says exactly what to install.

Windowed-sinc polyphase resample ``rate`` → ``target_rate``, pure torch.

Same kernel family as ``torchaudio.functional.resample`` at its
defaults (``lowpass_filter_width=6``, ``rolloff=0.99``, Hann
window), so mels computed from a resampled clip agree with the
reference pipeline (which resamples via torchaudio) to float noise
— with no torchaudio dependency on the serving path.

Per-model modality bindings — the modality-agnostic tower seam.

A multimodal model exposes ``mm_bindings: dict[str, MediaBinding]``
(one entry per input modality it accepts: ``"image"``, ``"audio"``, …).
Everything downstream — the OpenAI content-part glue, the runner's
encode-once + per-chunk scatter, the TP embed broadcast — iterates the
bindings instead of hard-coding a modality:

  * the binding's ``placeholder_token_id`` marks the positions in
    ``prompt_token_ids`` that the tower output overwrites;
  * ``encode(features, device) -> (n_tokens, hidden)`` runs the tower
    once per request over that modality's preprocessed features (the
    matching entry of ``Request.multimodal``), returning embeddings in
    placeholder order;
  * ``uses_mrope`` opts the modality into Qwen-style 3-D M-RoPE
    positions (vision); flat-position modalities (audio) leave the 1-D
    positions untouched.

Text-only models simply have no ``mm_bindings`` attribute (or an empty
dict) and skip every multimodal branch.

One input modality's wiring on a model.

Args:
    modality: stable name — the key in ``Request.multimodal`` and
        in the model's ``mm_bindings``.
    placeholder_token_id: the token id whose positions in the
        expanded prompt receive this modality's embeddings.
    encode: run the tower over this modality's per-request features
        → ``(n_tokens, hidden)`` in placeholder order. Called once
        per request by the runner; the result is cached on the
        ``Request``.
    num_placeholder_tokens: features → total placeholder-token
        count, used by the loud count guard at encode time (must
        equal the number of ``placeholder_token_id`` occurrences in
        the expanded prompt).
    encode_in_activation_arena: route tower allocations through the
        ``scratch.forward_arena`` named POOL
        (``EagerModelRunner._activation_arena_pool_ctx``) so they do
        not fragment against KV. Independent of
        ``--enable-activation-arena``, which is the per-layer
        :class:`~arbi_serve.runtime.activation_arena.ActivationArena`
        inside the same pool. Models opt in only when the cached
        output plus the following text forward fit within the
        profiled maximum lifetime.
    uses_mrope: modality participates in Qwen M-RoPE 3-D positions.
    extra: modality-specific knobs the API glue reads (e.g. vision's
        ``merge_unit``); kept out of the dataclass fields so new
        modalities don't grow this seam.

The engine's input bindings, whichever engine mode we are in.

In-process: ask the live model. Process mode: the API child has no model
object, so ``ProcEngine.mm_bindings`` resolves them from the served model
dir. ``{}`` for text-only models — callers skip every multimodal branch.

Mirrors :func:`arbi_serve.multimodal.output.resolve_output_spec`. Reading
``eng.model.mm_bindings`` directly yields ``None`` on a ProcEngine, which a
caller cannot distinguish from "text-only model" — so an audio-in request
would silently drop its audio and answer from an empty user turn.

M-RoPE 3-D position indices for Qwen-VL families.

Qwen-VL spreads RoPE over three position channels — temporal, height,
width — so a 2-D image grid keeps spatial structure under rotation.
Text tokens get the same value on all three channels (so M-RoPE
degenerates to standard 1-D RoPE on text), while image tokens get
distinct ``(t, h, w)`` coordinates and the running position advances by
``max(grid_h, grid_w) / merge`` past each image.

:func:`qwen_mrope_position_ids` reproduces
``Qwen3_5Model.get_rope_index`` for a single (unpadded) sequence — the
shape arbi-serve schedules one request at a time as. The result is a
``(3, seq_len)`` int tensor consumed by the M-RoPE-aware partial-rotary
cache in :mod:`arbi_serve.models.qwen3_5`.

3-D positions for one image's vision tokens (post-merge grid).

Mirrors ``Qwen3_5Model.get_vision_position_ids`` with
``temp_merge_size=1`` / ``time_interval=1``: temporal index per
frame, ``height + start`` / ``width + start`` over the merged
spatial grid, in ``(t, h, w)`` token order.

``(3, seq_len)`` M-RoPE positions for one sequence.

``input_ids`` is a 1-D token tensor; runs of ``image_token_id`` are
replaced by the per-image 3-D vision positions (consuming images
from ``image_grid_thw`` in order), and text runs get a shared
ascending counter. Equivalent to ``get_rope_index`` for batch
size 1 with no padding.

OpenAI Realtime API-compatible voice session layer.

The WebSocket route (:mod:`arbi_serve.server.routes.openai_realtime`) is a
thin adapter; the session state machine, VAD, audio codec helpers, and the
engine-driven turn generator live here so they stay unit-testable without
FastAPI / torch at import (the route module must remain torch-free — see
``routes/__init__.py``).

Turn-based caveat: Step-Audio-2-mini is NOT natively full-duplex — it
answers one committed user turn at a time. "Barge-in" here is therefore
VAD-orchestrated: fresh user speech detected during an active response
CANCELS that response (``eng.cancel``) and opens a new turn; the model is
never attending to overlapping audio. See :class:`~.session.RealtimeSession`.

Duplex connection lifecycle: build, publish (admit), and tear down a
:class:`~arbi_serve.realtime._duplex_connection.DuplexConnection`.

This module owns the three-step handoff a caller
(``RealtimeSession``, a test harness) drives a connection through:
:func:`attach_duplex_connection` (build + register), then
:func:`publish_duplex_connection` (admit — ``run_forever`` starts driving
it), and eventually :func:`close_duplex_connection` (tear down). It also
owns the engine-loop marshaling primitive (:func:`_run_on_engine`) both
the publish and close paths cross the WS/engine thread boundary through.

Resolve ``(hard_cap_ticks, soft_cap_ticks)`` for a new connection.

KV/state cap is a
SERVER-LEVEL setting (``BatchConfig.duplex_max_session_s``/
``duplex_soft_warning_s``, converted to ticks via
``duplex_frame_interval_s``), not a per-connection knob a caller
quietly loosens — so ``max_ticks=None`` (the default) always resolves
to the server's configured cap. A caller that DOES pass an explicit
``max_ticks`` (every current test, and a bounded/turn-shaped session's
own short safety valve) keeps that value verbatim for the hard cap,
but the soft-warning threshold still scales proportionally
(``duplex_soft_warning_s / duplex_max_session_s``, i.e. 75% by
default) so "soft warn" keeps meaning "most of the way to the hard
cap" regardless of which value produced it.

``BatchConfig.duplex_max_input_lag_s`` in FRAMES, or

``None`` — never bound — for a BOUNDED (turn-shaped) session: it is
handed a whole turn's audio up front by construction, so its queue
depth carries no information about lag and capping it would discard
the turn's own input. A genuine real-time connection gets the
configured bound, on this engine's own frame grid, so the limit means
a fixed amount of WALL CLOCK however the deployment paced its duplex
sessions. ``duplex_max_input_lag_s <= 0`` disables the bound.

``BatchConfig.duplex_bounded_preroll_silence_s`` in FRAMES, or


``None`` — never trim — unless this connection can enter a scheduler-
parked state. Turn-rearmable connections park between bounded turns;
semantic duplex connections park after sustained idle silence. Both
retain the newest silent run-up before recognized speech.

``duplex_bounded_preroll_silence_s <= 0`` disables the trim, on this
engine's own frame grid otherwise — so the margin is a fixed amount of
WALL CLOCK however the deployment paced its duplex sessions.

Refuse to attach a duplex connection to a MULTI-RANK engine.

The WS path's early copy of the one refusal
(:func:`~arbi_serve.realtime.nemotron_voicechat_duplex_admission
.duplex_multi_rank_refusal`, which owns the reasons): it fires here,
at attach, so nothing is registered on the lane before the scheduler
admission refuses the same engine. Both sites share the predicate, so
they can never disagree (design doc §7.18).

Build a :class:`DuplexConnection`, register it on ``eng``'s lane
(creating the lane on first use), and return it.

``request`` must ALREADY be admitted (``admit_duplex_request``) — this
function does not admit anything, mirroring the admission/drive
separation the rest of this design keeps.

``max_ticks``: the connection's KV/state cap, in ticks (design doc
§3.5) — ``None`` (the default) resolves it from the server's
configured ``BatchConfig.duplex_max_session_s``/
``duplex_frame_interval_s`` (see :func:`_resolve_duplex_caps`); an
explicit value overrides it (a bounded/turn-shaped test session's own
short safety valve, e.g.). Either way this is the ONLY thing that ever
stops a genuine (``bounded=False``) duplex connection on its own — see
``new_duplex_request``'s ``ignore_eos=True``/``max_tokens=None``.

``semantic_eou``: routes the EOU-gated BOS admission tracker to the
per-frame semantic verdicts riding this connection's frames (the
served checkpoint's own RNNT endpoint detector, design doc §7.32)
instead of the acoustic VAD — pass ``True`` only for a connection
whose frames actually carry verdicts (a genuine duplex connection fed
through ``DuplexAudioIn`` on a model with the bundled RNNT head); see
:class:`~arbi_serve.realtime.duplex_tick_pump.DuplexTickPump`.

``silent_grace_ticks``: the silent-turn stop condition for a BOUNDED
session (:meth:`DuplexTickPump.post_stt`). ``None`` (the default)
resolves it from :data:`_BOUNDED_SILENT_GRACE_S` on this engine's own
frame grid for a bounded session, and stays ``None`` — inapplicable —
for a genuine duplex connection, which has no turn to end. Pass ``0``
to end a silent bounded turn on the very first tick past its audio,
or a large value to keep the pre-existing "run to ``max_ticks``"
behaviour.

``turn_rearmable``: ``False`` (the default) preserves the original
bounded-session contract exactly — this connection's end-of-turn
tears the whole connection down
(``DuplexLane._end_connection``). ``True`` marks a repointed
turn-based session (design doc §7.24): end-of-turn instead calls
``DuplexLane._end_turn``, which leaves the Request live for
:meth:`DuplexConnection.begin_turn` to re-arm.

``start_awaiting_turn``: only meaningful alongside ``turn_rearmable=
True`` — the connection starts already waiting for its FIRST
:meth:`DuplexConnection.begin_turn` call rather than ticking
immediately, so a repointed session's caller drives turn 1 through
the exact same handoff as every later turn instead of a special case
(``False`` is correct for every other caller, including a bounded
one-turn test, which wants to start ticking immediately).

Post-condition on a duplex teardown: the row is really gone.

Every duplex close path funnels through
:func:`~arbi_serve.engine.lifecycle.cancel_by_id`, which is
deliberately tolerant of an unknown id — a cancel is allowed to race
its own publish, so a miss writes a tombstone and returns quietly.
That tolerance is right for the submission intake and was ``eng.requests`` at all, so every teardown took the
quiet-miss branch and left the row in ``Scheduler.running`` forever,
holding the EXCLUSIVE duplex lane against every later session — with
not one log line to say so.

``_publish_inner`` now registers the request, so the miss cannot
happen; this check is what makes a future regression of that class
loud instead of silent. It runs once per connection close (never on
the tick path) and scans only ``running``/``waiting``, which hold at
most ``max_batch`` + queue entries. Never raises: it is called from
teardown paths that must not fail.

Run ``fn(*args)`` on the engine loop, marshaling if needed.

Exactly the pattern :func:`arbi_serve.engine.lifecycle.cancel_by_id`
uses (and for the same reason): the scheduler's queues and
``eng._wakeup`` are engine-loop-owned state that a WS handler on the
HTTP loop must never mutate directly. Falls through to an inline call
when the bridge is disabled (single-loop deployments,
``ARBI_ENGINE_OWN_THREAD=0``) or when already on the engine loop.

Engine-loop body of :func:`publish_duplex_connection`.

This is the duplex lane's INTAKE BOUNDARY — the exact analogue of the
submission intake's ``request_factory.publish_and_admit``, and it owes
the engine the same bookkeeping that function does.

``eng.requests`` is the engine's id -> ``Request`` map, and it is the
ONLY way to get from a request id back to the live object.
``publish_and_admit`` populates it for every ordinary request; this
path did not, so a duplex request existed in ``Scheduler.running``
with no entry in it. The consequence was not cosmetic:
:func:`~arbi_serve.engine.lifecycle.cancel_by_id` — which EVERY duplex
teardown path funnels through (:func:`_close_inner` on WS disconnect,
``DuplexLane._end_connection`` on the KV/state cap and on adapter
errors) — looks the id up in exactly this map, and on a miss treats it
as a cancel-that-raced-its-own-publish: it writes a tombstone and
RETURNS. So every duplex teardown was a silent no-op. The request was
never removed from ``Scheduler.running``, never released its pages,
and — because the duplex lane is EXCLUSIVE — kept being re-selected by
every future duplex tick forever, wedging the lane against every later
session on the box. Reproduced live (design doc §7.13):
``connection closed``, still ``1 active`` over 30 s later, with
the process only recovering on restart.

Registering here also makes the request visible to the machinery that
is SUPPOSED to see every live request and silently skipped duplex
sessions before: the SIGTERM drain (``server/lifecycle.py``, which
reported "0 in-flight" while a duplex session was running), the
critical-section drain that gates model reload / sleep
(``engine/critical.py``), and admin introspection. The timeout sweep
is unaffected — it skips any request with no ``submit_time`` /
``effective_timeout_s``, and ``new_duplex_request`` sets neither
(a duplex session is bounded by the §3.5 KV/state cap, not by a
request deadline).

Symmetric teardown already exists: ``lifecycle._cancel_inner`` pops
the id, as do the ordinary finish paths (``run_step/finish.py``,
``run_step/terminal.py``) — so the ``cache.max_context`` backstop
cleans up through the same map.

THE handoff: admit ``conn``'s request so ``run_forever`` starts
driving it.

Call this only once the connection's inbound queue has whatever audio
is already available, so the first autonomous tick has a real frame
rather than a silence fill. From this point on the engine's own loop
owns every forward for this request; the caller only pushes frames in
and drains events out.

Tear a connection down: deregister it from the lane and release its
scheduler-owned ``Request`` (pages, recurrent state) through the SAME
``cancel``/``_finish_request`` path every other request uses.

This is the explicit-close entry point (client disconnect, server
shutdown, an application-level error) — the WS-side counterpart of the
lane's own internal teardown paths (``DuplexLane._end_connection``/
``DuplexLane._close_already_finished``), reusing the SAME underlying
``cancel_by_id``/``Scheduler.remove`` machinery so a connection is torn
down identically regardless of which side initiated it.

``reason`` is stamped on ``conn.finished_reason``/``Request
.finish_reason`` (mirrors ``cancel_by_id``'s own ``reason`` contract)
— but only if the connection does not already have one, so a
connection the lane already closed for its own reason (a KV/state-cap
hit, an adapter error) keeps that reason even if the WS layer's own
disconnect handler calls this afterward, unaware the connection was
already gone. Idempotent either way: closing an already-closed
connection is a safe no-op (``cancel_by_id`` tombstones an
unknown/already-finished request id; ``DuplexLane.unregister`` on an
already-unregistered id is a no-op).

Marshaled onto the engine loop, so the audio-out generator is never
closed underneath a tick that is mid-flight on the engine thread.

One live duplex WS connection's lane-side state.

Built by :func:`~arbi_serve.realtime.duplex_lane.attach_duplex_connection`,
which also registers it. The WS handler side only ever touches
:meth:`push_audio_frame` and :attr:`outbound`/:attr:`wakeup`; the engine
thread only ever touches the rest.

Hand one already-encoded audio frame to the lane (WS side).

Decoupled from when the engine actually schedules the tick that
consumes it: the handler pushes as audio arrives, the engine pops
on its own 80 ms grid, and a tick that finds the queue empty is
fed silence rather than skipped.

``semantic_blank``: the frame's RNNT endpoint-detector verdict
(design doc §7.32), when the encode path that produced the frame
computed one — ``None`` for a frame without a verdict (the
pre-encoded/turn-based paths).

**Overflow is dropping the OLDEST audio, deliberately** (design
doc §7.45). On a genuine real-time connection :attr:`inbound`
carries a ``maxlen``; appending to a full one discards the frame
at the left, i.e. the stalest. That is the correct direction:
this queue's whole purpose is to let the model hear the
microphone AS IT IS NOW, and a client that delivers faster than
the engine can step (a fast audio clock, an
``input_audio_sample_rate`` lower than what it really sends)
otherwise pushes the model arbitrarily far into the past —
measured live at 200+ frames and climbing, with the user's real
speech still queued when the session hit its KV cap. Losing the
oldest 80 ms is a real loss and is counted and reported as one;
losing the live conversation is not recoverable at all.

Drop banked SILENCE from the front of a between-turns queue,
keeping the :attr:`preroll_silence_keep` frames that immediately
precede the user's speech (design doc §7.48).

Safe to run from the WS side because both parked states disable
the scheduler row, making this thread the queue's sole consumer
for that window. It uses only atomic deque operations.

**What is droppable is decided by the audio, not by the VAD event
stream.** Every queued frame already carries the checkpoint's own
RNNT endpoint-detector verdict (design doc §7.32), computed off
the same encoder forward that produced the frame, at append time.
A frame is trimmable only when that verdict is an explicit
``True`` — no verdict at all (``None``: a model with no
``rnnt_head``, or a pre-encoded/turn-based frame) is treated as
speech and stops the scan cold, so a deployment without the
detector keeps every frame rather than dropping audio it has no
evidence about. This is deliberately NOT the
``speech_started``/``speech_stopped`` latch §7.44.4 refused to
build on: that one is a silero edge detector over the raw PCM,
and it is the signal that misfired there.

The kept margin is the frames JUST BEFORE the first non-blank one,
not the oldest ones — the queue's whole value here is the acoustic
run-up to what the user actually said, and a detector that flags
the opening frames of a word as blank must not be able to clip the
word. Dropping the interior of a silence changes nothing the model
could have learned from it.

**The scan stops at the first non-blank frame, with no tolerance
for an isolated one, and that is not an oversight.** This detector
is SPARSE inside real speech: §7.46.4 measured an 11-20% non-blank
duty cycle with blank runs "up to 79 ticks inside one speech
window". So a rule that stepped over an isolated non-blank as
noise — treating it as a cough or a keystroke, which reads
perfectly sensible in the abstract — would step straight through
the middle of an utterance and trim the user's own speech. That is
the exact "fast but corrupting" trade §7.44.4 declined to make,
and the asymmetry here settles it: a spurious non-blank on true
silence only makes this trim do LESS (the connection degrades to
the pre-§7.48 late answer), while a missed one would delete
speech. Only the first kind of mistake is affordable.

Bounded work per call: after one pass the leading blank run is at
most ``keep``, so every later call scans at most ``keep + 1``
frames before it stops.

Report a connection that is losing input audio — once when it
starts, then once per further second's worth of dropped frames.

Unconditionally loud (``WARNING`` + an OTel counter) because this
is never normal: the bound it hit sits far above every healthy
backlog this integration has measured, so reaching it means the
client is delivering audio faster than the engine can consume it
and the conversation is degrading. Rate-limited rather than
per-frame — an over-delivering client trips this on EVERY frame,
and a log line per 80 ms would bury the very signal it exists to
raise.

Queue context tokens to be fed into the running sequence (WS side).

The tool-call resume hook: ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>``
fed back into a paused turn. The next duplex tick becomes an
injection tick — the engine forwards these tokens as one wide
``1 + N`` row and resumes generation from its last position. It is
still an ordinary frame tick otherwise: this tick's audio frame is
consumed and fused into the row's first slot, and the tick's
sampled text/function tokens feed the frame-lockstep recurrence
exactly as on any other tick (design doc §7.9.15; see
``DuplexLane._prepare_tick_inner``).

The tokens are context, not output: they take KV slots and advance
positions but never enter ``output_token_ids`` or the transcript.

Staged here rather than written straight onto the request: the
engine reads ``pending_context_token_ids`` TWICE per step (the
slate builder sizes the row from it, then ``_gather_slate``
re-derives the width to validate), so a WS-thread write landing
between the two would desync them. The pre-schedule hook — which
runs on the engine thread, before ``schedule()`` — moves the
staged ids onto the request, so the field is only ever mutated
where the step cannot be mid-flight. ``deque.append`` is the same
thread-safe handoff :meth:`push_audio_frame` uses.

Stage a request to re-arm this connection's pump for another
bounded turn (WS side) — design doc §7.24's turn re-arm primitive,
the duplex analogue of building a fresh ``EngineSttStep`` for the
next turn without tearing this connection's KV/Mamba state down.

Only valid on a :attr:`turn_rearmable` connection whose current
turn has already ended (:attr:`awaiting_turn`); raises loudly
otherwise rather than silently staging a re-arm nothing will ever
honour, or racing an in-flight turn's own termination check.

Staged rather than applied here for the same reason
:meth:`push_context_tokens` stages: :class:`DuplexTickPump`'s
fields are read mid-tick on the engine thread, so the actual
mutation happens in the pre-schedule hook, on that thread, before
``schedule()`` can observe a half-updated pump.

Discard every queued-but-unconsumed inbound frame (WS side).

The duplex analogue of ``input_audio_buffer.clear``: "forget the
audio I sent that you have not listened to yet". Frames the engine
has ALREADY consumed are not (and cannot be) unheard — they are in
the request's KV/Mamba state — which is exactly the same bound the
turn-based path has, where ``clear`` only ever drops the
not-yet-committed tail of the buffer.

Deferred to the engine thread rather than clearing here. The
inbound ``deque`` is a lock-free single-producer/single-consumer
handoff (``duplex_lane``'s own module docstring, threading
section): ``append`` from this side and ``popleft`` from the
engine side are each atomic under the GIL, but a WS-side
``clear()`` is neither of those and can race
``DuplexConnection._next_frame``'s ``if self.inbound`` /
``popleft`` pair into an ``IndexError`` on the engine thread —
which the lane would (correctly, and uselessly) turn into a
dropped connection. Setting a flag the consumer honours keeps the
queue single-consumer.

Stop the agent mid-utterance and terminate the response it is
producing (WS side) — the duplex analogue of ``response.cancel``.

Four effects, deliberately the SAME the VAD barge-in AND-gate
already produces (``duplex_tick_pump``/``duplex_barge_in``),
since "the human interrupted" and "the client asked me to stop"
want identical behaviour from the model:

  * ``mark_agent_idle()`` closes the turn at the CLIENT-VISIBLE
    level;
  * ``force_eos_feedback`` closes it at the MODEL level too
    (design doc §7.26) — ``mark_agent_idle`` alone leaves the
    model's own recurrent conditioning (``prev_text_token_id``)
    following whatever it actually said the instant before this
    call, so without this a cancelled turn can keep sampling
    content from that stale mid-utterance state indefinitely,
    immune to the flag this call just forced. ``getattr``: a
    minimal session stand-in in a CPU test predates this hook and
    stays a plain flag flip;
  * **for a ``turn_rearmable`` connection, ``awaiting_turn`` is
    forced ``True`` directly** — the SAME effect
    ``DuplexLane._end_turn`` produces, applied here instead of
    waiting for one. This is the load-bearing fix (design doc
    §7.26): the ONLY other way a bounded turn ever reaches
    ``awaiting_turn=True`` is a full active-then-idle cycle
    observed by ``DuplexTickPump.post_stt`` on a REAL tick
    (``agent_was_active and agent_idle_this_tick``), which
    requires the agent to have gone active at least once. SAME utterance) reads as a
    barge-in on a turn that had not yet produced anything and
    cancels it — would otherwise leave ``agent_was_active``
    permanently ``False`` and ``awaiting_turn`` permanently
    ``False``, so ``DuplexConnection.begin_turn``'s own
    precondition can never be satisfied again and the connection
    is stuck refusing every future commit. Directly setting the
    flag needs no additional tick to run: ``_end_turn``'s own body
    is exactly this one assignment (plus a reason string and a
    consumer wakeup, both applied here too) — nothing about it
    requires the engine thread;
  * a ``TurnDone("client_cancelled")`` on the outbound queue gives
    the WS bridge a terminal event to close the SYNTHETIC response
    with, so the client sees ``response.done`` with
    ``status="cancelled"``/``reason="client_cancelled"`` — the
    same status vocabulary ``RealtimeSession._finalize_response``
    stamps on a cancelled turn-based response, rather than the
    ``"completed"`` an ordinary agent-went-idle transition means.

A no-op beyond the idle flip when the agent is already idle: the
bridge drops a terminal event with no response open.

``append`` (never ``extend``) so this shares the outbound deque
with the engine thread's own per-tick publish under exactly the
atomic single-op discipline the rest of this module uses.

Queue an ``ack_messages`` "on hold" phrase to be SPOKEN while a
tool call is pending (WS side) — design doc §7.9.18.

The duplex equivalent of the turn-based path's
``_maybe_speak_ack``. It cannot reuse that mechanism (a standalone
TTS-only thread over the phrase's own token ids) because duplex
mode has no second TTS thread — everything goes through the one
frame-lockstep tick loop. So the phrase is expressed as forced
TEXT tokens through the very gate that would otherwise be forcing
PAD: :func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
pops one id per 80 ms frame, and because this checkpoint's
STT->TTS coupling is a bare id passthrough that is the same
per-frame ``generate_tts_frame`` drive the turn-based path
performs — at the same rate, from the same ids.

Staged rather than written straight onto the session's FC state
for the same reason :meth:`push_context_tokens` stages: the field
is read on the engine thread inside the step, and the
pre-schedule hook is the one place a WS-originated write can land
without racing it.

Signal the WS-side consumer that this connection has state to
look at — exactly ONE hop per tick, through the marshaler seam
(``LoopBridge.notify_http``), never a raw cross-loop reach-in from
here (``tests/test_transport_boundary.py``'s invariant). A closed
consumer loop is not an error: the events stay in the bounded
deque for a final drain.

This tick's frame: the oldest real one, or silence.

Also latches :attr:`prepared_frame_embed` — set to the dequeued
frame when it came from :attr:`inbound`, and back to ``None`` for a
synthesized silence frame — so a tick that is prepared but never
stepped can return a REAL frame to the queue without ever
re-queueing filler
(``DuplexLane._requeue_prepared_frame``).

===

The opt-in in :mod:`arbi_serve.realtime._session_duplex_optin` is a CLIENT
choice: a client that never sends ``duplex: true`` stays on ORDINARY
turn-based dispatch at the wire level — ``input_audio_buffer.commit``/
``response.create`` keep meaning exactly what they meant before,
``duplex_enabled`` stays ``False``, every event type stays available. What
changes for a served model that declares ``SUPPORTS_DUPLEX_MODE``
(:meth:`_BoundedDuplexMixin._bounded_duplex_available`) is which BACKEND
generates the response: instead of building a fresh ``EngineSttStep`` per
turn (:mod:`arbi_serve.realtime.nemotron_voicechat_turn`'s
``stream_nemotron_voicechat_turn``, which re-feeds the ENTIRE conversation
history from scratch every turn), this class lazily builds ONE persistent,


  * :meth:`_BoundedDuplexMixin._maybe_begin_bounded_turn` (awaited right
    after ``_commit_user_turn``, i.e. on EVERY commit — auto ``server_vad``
    or manual) closes the "don't speak until asked" gate and calls
    :meth:`~arbi_serve.realtime.duplex_lane.DuplexConnection.begin_turn`
    with this utterance's real-audio frame target — the connection starts
    consuming it, muted.
  * :meth:`_BoundedDuplexMixin._start_bounded_response` (the
    bounded-backend half of ``_start_response``, reached from both the
    auto path (``_commit_and_respond``) and the manual one
    (``_on_response_create``)) opens the gate — the model may now speak.
    From here the ALREADY-RUNNING background pump task
    (:func:`~arbi_serve.realtime.duplex_ws_bridge.pump_duplex_connection`,
    started once in :meth:`_ensure_bounded_duplex_conn`, exactly the
    mechanism ``_enable_duplex_mode`` already uses) streams the same
    ``response.*`` wire events ``_run_response`` used to build by hand,
    including its already-built ``ensure_response`` exactly-one-
    response-per-bounded-turn framing (design doc §7.21.2's silent-turn
    fix) and its already-built function-calling round trip
    (:meth:`_on_duplex_function_calls` et al. below, generalized to read
    :attr:`_active_duplex_conn` instead of hardcoding ``_duplex_conn``, so
    it now serves BOTH connection kinds for free).

Audio itself streams into the connection incrementally, the same way it
would in full duplex mode (``_on_append`` routes it through
:class:`~arbi_serve.realtime.duplex_audio_in.DuplexAudioIn` whenever the
bounded backend is in use), NOT batched at commit time — this is the
concrete reason the speak gate has to exist at all here rather than being
redundant with commit-time batching: the connection is genuinely
"listening" to audio before the client has asked it to answer, and
without the gate Site-3's BOS could fire on that audio unprompted, which
is duplex's contract, not turn-based's. ``_input_buf``/``_history`` stay
exactly as they were — the wire-visible commit/item bookkeeping, the
transcription echo and live voice-clone reference audio are all
unaffected; what the bounded connection's own prompt renders ONCE, at
first use, is ``instructions``/``tools`` (mirroring
``_enable_duplex_mode``'s own seeding), and every turn after that is
audio frames plus the model's own token stream — no re-rendered
``messages`` at all. This is why multi-turn coherence through this
backend is STRICTLY STRONGER than ``EngineSttStep``'s: nothing is ever
thrown away and re-fed.

**Deliberately deferred, stated rather than silently dropped** (each is
independent of the two primitives themselves and does not block the
repoint):

  * **Speculative endpointing** (``turn_detection.speculative_hold_ms``)
    is built on ``_run_response``'s own event-buffering (``_emit_turn``/
    ``_held``); the bounded backend's events flow straight through
    ``pump_duplex_connection`` to the socket with no equivalent buffer to
    hold. ``_speculative_hold_ms`` reads as ``0`` whenever the bounded
    backend is in use, falling back to ordinary immediate commit+respond
    — exactly what a client that never set the knob already gets, not a
    broken feature.
  * **Live voice selection per turn** (including ``CLONE_VOICE``) is not
    read for the bounded backend, matching the SAME limitation full
    duplex mode already has (``_enable_duplex_mode`` does not call
    ``_resolve_response_voice`` either) — voice is whatever the
    connection's adapter was built with at first use.
  * **Mid-session ``instructions``/``tools`` changes** are not re-seeded
    into an already-built bounded connection the way
    ``_restart_duplex_session`` re-seeds a full duplex one — the
    ``session.updated`` echo reflects the new values, but the connection's
    own prompt (rendered once) does not. * **A manual-mode client that commits and then waits an unusually long
    time before calling ``response.create``** may find its turn already
    closed as an empty ``no_response`` by the bounded (turn-shaped)
    silent-turn grace (design doc §7.21.2) before ever asking — the
    grace's premise ("the model is free to answer and chooses not to")
    does not strictly hold while the speak gate is still closed. Mitigated
    by resolving a generous ``silent_grace_ticks`` for this backend
    specifically (:data:`_BOUNDED_BACKEND_SILENT_GRACE_TICKS`) rather than
    the ordinary few-second default, bounding a genuinely abandoned turn
    by the connection's own overall KV/state cap instead. The common
    cases — ``server_vad`` auto-respond, and manual mode's immediate
    commit-then-create — are unaffected either way.
  * **Mid-response barge-in racing an immediate re-commit** faster than
    one duplex tick (80 ms) is a narrow, believed-unreachable-in-practice
    race (every default VAD's own ``silence_duration_ms`` is well over
    80 ms) between ``_on_speech_started``'s cancel and the next commit's
    :meth:`_maybe_begin_bounded_turn`; if hit, it surfaces as an ordinary
    ``internal_error`` event, not a crash or silent misbehavior.

Whichever duplex connection is actually driving this session's
audio right now — the client's own full duplex opt-in
(:attr:`_duplex_conn`) if active, else the bounded backend
transparently driving ordinary turn-based responses
(:attr:`_bounded_conn`, design doc §7.24), else ``None`` (a served
model with no duplex support at all). The two are mutually
exclusive by construction (:meth:`_bounded_duplex_available`
refuses once :attr:`duplex_enabled` is set), so this is never
ambiguous.

Whether ORDINARY (non-opted-in) turn-based responses on this
connection should be driven through the bounded duplex backend
(design doc §7.24's repoint) instead of
``self._turn_streamer``/``EngineSttStep``.

Mutually exclusive with the client's own explicit full duplex
opt-in, which already owns this connection's whole audio path —
checked first so the two can never both try to build a
connection. A served model that does not declare
``SUPPORTS_DUPLEX_MODE`` (Step-Audio-2-mini today) always answers
``False`` here, leaving that model's turn-based dispatch entirely
untouched by this backend's existence.

Lazily build this connection's persistent BOUNDED, re-armable
duplex connection (design doc §7.24) the first time ordinary
audio/a response needs to be driven through it. A no-op returning
``True`` if already built. On failure, sends an ``error`` event
and returns ``False`` — the same failure shape
``_enable_duplex_mode`` uses.

Mirrors ``_enable_duplex_mode``'s own connection-setup recipe
almost exactly (same adapter resolution, same seed-prompt
rendering from ``instructions``/``tools``), with three
differences: ``bounded=True``/``turn_rearmable=True``/
``start_awaiting_turn=True`` (this connection lives across many
turns, one at a time, and its first turn goes through the same
:meth:`~arbi_serve.realtime.duplex_lane.DuplexConnection
.begin_turn` handoff every later one does — see
:meth:`_maybe_begin_bounded_turn`); the speak gate starts CLOSED
(a repointed session may not speak until asked, unlike genuine
duplex); and :attr:`~arbi_serve.realtime.duplex_lane.DuplexConnection.pump`'s
silent-turn grace is widened
(:data:`_BOUNDED_BACKEND_SILENT_GRACE_TICKS` — see this module's
own docstring's "deliberately deferred" section for why).

Hand a just-committed utterance over to the bounded duplex
backend (design doc §7.24) — awaited right after
``_commit_user_turn`` (which is itself synchronous) by both its
callers, ``_on_commit`` and ``_commit_and_respond``, for EVERY
commit (auto ``server_vad`` or manual), so the gate closes and the
pump's own real-audio target advances whether or not a response
follows immediately.

A no-op when the bounded backend is not in use. Lazily builds the
connection on the very first call (mirrors ``_on_append``'s own
lazy build for the audio-in side — whichever event reaches this
connection first is the one that creates it).

The bounded-backend half of ``_start_response`` (design doc
§7.24): open the speak gate so the connection may now answer.

Gated on :attr:`_bounded_turn_committed` (a plain, WS-thread-
synchronous flag), deliberately NOT ``conn.awaiting_turn``: the
engine thread only clears that once it drains the STAGED
``begin_turn`` on its own next tick (up to one 80 ms tick of lag)
— gating the actual gate-open on it would risk opening a mute
connection that never gets un-muted before its own silent-turn
grace elapses. ``open_speak_gate()`` is a plain attribute write
with no such race: whenever the pump does start consuming this
turn's audio, it reads the flag fresh, whether that is before or
after this call landed.

Tear down the bounded duplex backend (design doc §7.24) —
idempotent, the bounded-backend counterpart of ``_aclose_duplex``
for a client's own explicit opt-in. No restart/drain shape here:
this backend has no mid-session re-seed path (see this module's
own docstring's "deliberately deferred" section), so a plain
client-disconnect-shaped close is the only one it ever needs.

One tick's detected tool calls, handed over by the duplex pump
on THIS loop — the WS-side half of the FC pause/resume flow.

Synchronous and non-blocking by contract (see
:func:`~arbi_serve.realtime.duplex_ws_bridge.pump_duplex_connection`):
it registers each call's pending future immediately, so the client
cannot answer faster than the bridge can accept the answer, and
defers the actual waiting — up to
:attr:`ToolResultBridge.timeout_s`, ten seconds by default — onto
its own task. Blocking here would stall the consumer that is
supposed to keep delivering the connection's audio WHILE the tool
runs.

Speak one of the called tool's ``ack_messages`` while the client
executes it — the duplex half of the "on hold" UX the turn-based
path gets from :func:`~arbi_serve.realtime.nemotron_voicechat_turn
._maybe_speak_ack` (design doc §7.9.18).

Selection is that function's rule verbatim (its own
:func:`_tool_ack_messages` lookup over the raw tool dicts, then a
uniform :func:`random.choice` among the FIRST called tool that
defines any) — this differs only in how the phrase is spoken.
Turn-based synthesizes it on a standalone TTS thread while
``run_turn`` is parked; duplex mode has no second thread and never
parks, so the phrase goes out as forced TEXT tokens through the
pause gate itself (``DuplexConnection.push_ack_tokens``), one id
per 80 ms frame. Same ids, same per-frame TTS drive, no new
machinery.

Best-effort throughout, exactly like the turn-based one: no
``ack_messages`` configured, a connection without the hook, or an
encode failure yields silence — the gate then just holds PAD, i.e.
the pre-§7.9.18 behaviour. A missing filler phrase must never fail
a tool call.

Reopen the duplex session's ``_fc_in_progress`` PAD gate so the
agent's text/audio stream resumes. Tolerant of a session object
that has no gate (a non-NemotronVoiceChat duplex model, or a test
double): this is a release, so failing to find one is not an
error.

Await every call's client result, then resume the model with
one ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>`` injected into the
RUNNING sequence.

The whole batch is answered by ONE block, matching both the
reference chat template's documented client-response shape and
``stream_nemotron_voicechat_turn``'s own grouping — a
``<TOOLCALL>`` block naming N calls is answered by a
``<TOOL_RESPONSE>`` array of N results.

The resume is ``1 + N``
injection row (made correct on the lane by design doc §7.9.15) —
NOT a new request and not a re-prefill. The persistent request's
KV, Mamba state and frame-lockstep recurrence all carry straight
through, which is the entire reason duplex mode does tool calls
this way instead of the finish-and-resubmit shape vLLM/SGLang use
(§7.9.12's comparison).

The PAD gate is released in a ``finally``: a timeout, a
disconnected connection or an encode failure must never leave the
agent permanently mute.

THIRD tier) — one of :class:`~arbi_serve.realtime.session.RealtimeSession`'s
mixins (see that module's docstring for the overall session design and
why this concern is split out).

=== THIRD tier) ===

A THIRD dispatch tier, additive and OFF by default: a client that sends
``session.update {"session": {"duplex": true}}`` switches THIS connection
onto the duplex lane (:mod:`arbi_serve.realtime.duplex_lane`) instead of
the turn-based path above — the fundamentally different session shape
design doc §2's duplex lane was built for (no turns, no VAD-driven commit,
continuous audio in both directions).

**Why a dedicated top-level ``duplex`` boolean, not a new
``turn_detection.type`` value** (design doc §5 flagged this choice as
genuinely undesigned): ``turn_detection`` configures HOW a discrete turn
boundary is detected — it presupposes turns exist. Duplex mode has no
turns at all (design doc §5: "every audio delta the model emits is
unprompted, continuous" — there is no boundary for ANY ``turn_detection``
value, including a hypothetical ``"duplex"`` one, to describe). Modeling
duplex as a value of a field whose whole job is "pick a turn-boundary
algorithm" would be a category error, not a convenience. A dedicated
top-level field instead follows this module's own existing convention —
``modalities``/``voice``/``instructions``/``tools`` are all flat,
independent top-level session knobs, not nested inside one another — and
keeps ``turn_detection`` itself completely untouched (still meaningful,
still configurable, for a client that stays turn-based).

**Scope, a deliberate judgment call**: ``duplex`` may only be set on the
FIRST ``session.update`` that turns it on, before any turn-based activity
(a committed user turn or a response) has happened on this connection —
switching back to turn-based mid-connection, or switching TO duplex after
turn-based activity already started, is rejected with an ``error`` event
asking the client to open a new connection instead. The two session
shapes are close to disjoint state (design doc §1's own framing); trying
to reconcile a live turn-based response/history with a duplex connection
majority-owned by a persistent scheduler ``Request`` is real, unscoped
complexity this task does not need to take on to deliver the opt-in
itself. Once active, only the events in :attr:`RealtimeSession
._DUPLEX_ALLOWED_EVENTS` are accepted — see that attribute's own comment
for what each one means without a turn boundary, and why
``input_audio_buffer.commit`` / ``response.create`` (the two that ASSERT a
turn boundary) are the ones that stay rejected.

**Mid-session ``session.update``** (design doc §7.9.14's "irreducible
residue"): ``temperature``, ``output_audio_sample_rate`` and
``turn_detection`` are pushed onto the LIVE connection in place
(:meth:`_DuplexOptInMixin._apply_duplex_live_config`), matching how the
turn-based path simply re-reads them next turn. ``instructions``/``tools``
are rendered at prompt position 0, which an append-only injection row
cannot edit, so a change to either re-seeds the connection —
:meth:`_DuplexOptInMixin._restart_duplex_session`, the finish-and-re-admit
path, transparent at the protocol level.

``_on_duplex_function_calls``/``_resolve_duplex_function_calls``, shared
with the bounded duplex backend below.

**Audio routing**: once duplex is active, ``input_audio_buffer.append``
stops feeding the turn-based VAD/commit buffer above and instead routes
through :class:`~arbi_serve.realtime.duplex_audio_in.DuplexAudioIn` onto
the connection's inbound frame queue — see that module's docstring for
the real, currently-model-limited state of PCM-to-frame-embedding
encoding.

**Lifecycle**: :meth:`_DuplexOptInMixin._enable_duplex_mode` builds the
persistent duplex ``Request``/session/connection (mirroring
``tests/test_duplex_lane_run_forever_live_gpu.py``'s own recipe) and
starts :func:`~arbi_serve.realtime.duplex_ws_bridge.pump_duplex_connection`
as a background task feeding the SAME ``pump_duplex_connection``'s own docstring names this call site
as the precedent for.

A comparable snapshot of everything baked into the live duplex
connection at open — ``instructions`` and ``tools``, the two fields
the chat template renders at prompt position 0, plus ``tool_choice``
— or ``None`` when no duplex connection is live. Compared
before/after a ``session.update`` to decide whether the connection
still describes this session (see
:meth:`_apply_duplex_session_update`).

``tool_choice`` is in here even though it is NOT rendered into the
seed prompt: it is compiled into the connection's
:class:`~arbi_serve.realtime.duplex_function_grammar.DuplexFunctionGrammar`
at the same one-shot moment the seed is rendered, and there is no
in-place path to swap a live matcher's grammar. So a client that
raises the guarantee mid-session (``"auto"`` → ``"required"``) must
take the same rebuild every other seed-baked field takes; leaving it
out would silently keep serving the OLD policy while the session
object reported the new one.

``json.dumps(..., sort_keys=True)`` rather than the raw ``tools``
list because it must be an equality test against a value the
client may have re-sent verbatim: two structurally identical tool
lists must compare equal (and NOT trigger a re-seed), while any
real edit — a renamed tool, a changed schema, a new
``ack_messages`` entry — must not. Falls back to ``repr`` for the
(unvalidated, client-supplied) shape JSON cannot encode.

Reconcile a live duplex connection with a ``session.update``
that has already been applied to this object's own config.

Two classes of change, and the split is not a convenience — it is
exactly where the engine's prompt-then-output model draws the line
(design doc §7.9.14's "irreducible residue"):

  * **Everything that is not in the seed prompt** —
    ``temperature`` (a field on the request's own
    ``SamplingParams``, read fresh by the sampler every step),
    ``output_audio_sample_rate`` (re-read by the response bridge
    per audio delta), ``turn_detection`` (the barge-in VAD object)
    — is applied to the RUNNING connection in place. No
    discontinuity, matching the turn-based path, where all three
    are simply re-read on the next turn.
  * **``instructions``/``tools``** are rendered at prompt position
    0 and cannot be edited under existing output at all, and
    **``tool_choice``** is compiled into the connection's function-
    channel grammar at that same one-shot moment (there is no
    in-place path to swap a live matcher's grammar). The honest
    resolution is to treat the change as what it semantically is —
    a new session — and :meth:`_restart_duplex_session` does
    exactly that.

Push the non-seed half of this session's config onto the LIVE
duplex connection (see :meth:`_apply_duplex_session_update`).

``temperature`` is written straight onto the persistent request's
``SamplingParams``. That is a plain attribute store of a float
from the WS loop against a field the engine thread only ever
READS, one step at a time — the same single-writer/single-reader
discipline ``NemotronVoiceChatDuplexSession.end_function_call``
already relies on for the FC gate, and the same reason neither
needs a lock: the worst interleaving is that the change takes
effect one step later than it could have.

The VAD is re-read rather than re-pushed because
``_rebuild_vad`` REPLACES ``self._duplex_vad`` with a new detector
object, and the lane's pump captured the old one at attach time;
without this a mid-session ``turn_detection`` change would be
accepted at the protocol level and silently invisible to the
barge-in AND-gate. It is ``_duplex_vad``, never ``_vad``: the
pump is fed the frame queue's model-rate PCM, and ``_vad`` is
built for the client's own rate (design doc §7.48).

Cleanly end this connection's duplex generation and re-admit a
fresh one seeded from the CURRENT ``instructions``/``tools``
(design doc §7.9.14/§7.9.17's finish-and-re-admit resolution).

Transparent to the client at the protocol level: no error event,
no reconnect, no ``duplex`` toggle — the socket, the session id
and the mode are all unchanged, and the only thing the client
observes is that any response the agent was mid-way through is
closed with ``status="cancelled"``, ``reason="session_restart"``
(the pump's own teardown path emits it) before the re-seeded
connection starts producing again. That IS what happened: the
agent was cut off. Reporting it as a clean completion would be
the lie.

What is genuinely lost, stated rather than glossed: the old
request's KV/Mamba state, i.e. everything the model had heard on
this connection so far, plus whatever audio was queued on the
inbound side at the moment of the swap. That loss is the entire
content of "the system prompt changed" — the new prompt's
position-0 KV is by definition not a continuation of the old
one's — and it is why this path is reserved for the two fields
that cannot be expressed any other way rather than used for
ordinary config.

Returns ``True`` on success. On failure the connection is left
torn down and an ``error`` event has been sent by
:meth:`_enable_duplex_mode`; ``duplex_enabled`` stays ``True`` so
the connection does not silently fall back to a turn-based mode
whose state was never built.

Apply a ``session.update {"duplex": ...}`` request.

Returns ``True`` if ``_on_session_update`` should proceed to send
its normal ``session.updated`` echo (a no-op re-affirmation of the
current mode, or a successful enable); ``False`` if an ``error``
event was already sent and the caller should stop there.

Build this connection's persistent duplex ``Request``/session,
register it on the engine's duplex lane, and start the WS-side
consumer pump — see module docstring's "Lifecycle" paragraph.

Mirrors ``tests/test_duplex_lane_run_forever_live_gpu.py``'s own
connection-setup recipe (the only other real, proven caller of
this exact sequence: ``resolve_duplex_adapter`` ->
``adapter.new_session()`` -> ``new_duplex_request`` ->
``attach_duplex_connection`` -> ``publish_duplex_connection``).

Tear down this connection's duplex lane registration + pump task
(idempotent) — the

``reason`` is stamped on the connection and decides the status of
any synthetic response the pump still has open (see
``duplex_ws_bridge._STATUS_FOR_FINISHED_REASON``). The default is
the socket going away; :meth:`_restart_duplex_session` passes
``"session_restart"`` so a re-seed is distinguishable from a
disconnect in both the wire event and the server log.

Order matters: close the engine-side connection FIRST (frees pages/
recurrent state through the same ``cancel_by_id`` path every other
request's teardown uses — see
:func:`~arbi_serve.realtime.duplex_lane.close_duplex_connection`),
THEN cancel the pump task directly rather than waiting for it to
notice ``conn.closed`` on its own poll cadence --
``pump_duplex_connection``'s own docstring names exactly this call
site as the precedent for that choice.

``drain``: give the pump :attr:`_DUPLEX_DRAIN_TIMEOUT_S` to notice
``conn.closed`` and emit its own closing ``response.*`` sequence
before it is cancelled. Cancelling a pump parked on its poll
raises straight through the loop body, so the trailing close after
that loop never runs — which is correct on a client disconnect
(there is no socket left to send it to) and wrong on a session
restart, where the client is still there and is owed the terminal
events for the response it was receiving.

``_ResponseMixin`` — the turn-based response lifecycle, its speculative
endpointing opt-in, and live voice cloning — one of
:class:`~arbi_serve.realtime.session.RealtimeSession`'s mixins (see that
module's docstring for the overall session design and why this concern is
split out).

=== Speculative endpointing ===

On a short VAD pause the session commits + generates SPECULATIVELY but
HOLDS all turn output for ``speculative_hold_ms``. If speech resumes
within the hold the turn is discarded silently (no audio ever reaches the
client — the false endpoint is inaudible) and the utterance continues; if
the hold elapses the turn is confirmed and the buffered output is
flushed. This overlaps the expensive prefill/decode with the
endpoint-confirm wait instead of paying the full VAD hangover THEN
generating.

=== Live voice clone ===

``voice == CLONE_VOICE`` builds a token2wav conditioning triple from the
incoming speaker's most-recently-committed audio
(:class:`arbi_serve.audio.voice_prompt.VoicePromptExtractor`) so the reply
is rendered in THEIR voice rather than a prepared one. This is the
substrate for the realtime translated out-channel — you speak, and your
words come back in a target language IN YOUR VOICE. Falls back to
``_DEFAULT_VOICE`` (with a warning) when there is no input yet or the
clone deps are unavailable.

Freeze the input buffer into a user turn WITHOUT mutating history.

Mirror of ``_commit_user_turn``'s message shape, but non-destructive:
the buffer, history, and cursors are untouched so a resumed utterance can
continue. The snapshot is committed for real only on confirm.

The ``voice`` argument for the speech-out turn.

A prepared voice name passes through unchanged. ``CLONE_VOICE`` builds
the incoming speaker's conditioning triple from ``_last_user_pcm`` (off
the event loop) so the reply is rendered in their own voice. Any failure
degrades to ``_DEFAULT_VOICE`` with a warning — the turn still speaks,
just not cloned — so cloning can never break the response path.

Emit the terminal turn events. ``failure_message`` set means the
turn generator raised (see ``_run_response``'s ``except Exception``
branch) — the response object must report ``status="failed"``, not
``"completed"``: an empty-looking but "completed" response is
indistinguishable, on the wire, from the model genuinely choosing to
say nothing, which silently hides a real server error from the
client.

``_TranscriptionMixin`` — the input-audio transcription echo (off the
hot path) — one of
:class:`~arbi_serve.realtime.session.RealtimeSession`'s mixins (see that
module's docstring for the overall session design and why this concern is
split out).

The UI "you said …" echo: when ``input_audio_transcription`` is configured
on ``session.update``, a committed user turn's words are streamed back as
``conversation.item.input_audio_transcription.{delta,completed}`` over a
second, off-hot-path ASR turn — never awaited by the response path, so it
cannot delay time-to-first-audio.

StepFun's canonical ASR / S2TT instructions — see :mod:`arbi_serve.multimodal.asr_text`.

``language`` is accepted for API compatibility but deliberately NOT appended:
the canonical instructions are fixed strings the model was trained on, and the
model reports the detected language itself via its ``<英语>``/``<中文>`` tag,
which :func:`~arbi_serve.multimodal.asr_text.split_language_tag` extracts.

Cancel any in-flight UI-echo ASR turns.

The echo runs at ``priority="batch"``, but a *superseded* turn's echo
(barge-in, client cancel, or a fresh commit) is moot and — left running
— piles up concurrent multimodal requests that thrash the engine and can
starve the live response. Cancelling bounds concurrency to one echo. The
tasks self-remove from ``_asr_tasks`` via their done callback.

PCM16 codec + resampling helpers for the Realtime WebSocket.

The Realtime wire carries base64 little-endian 16-bit mono PCM
(``pcm16``). This module converts between that wire form and the raw
bytes the model pipeline consumes:

  * INPUT, turn-based — appended ``pcm16`` chunks accumulate as raw
    bytes; a whole committed turn is wrapped in a stdlib WAV container at
    the session's input sample rate. The audio preprocessor decodes any
    PCM16 WAV and resamples to 16 kHz internally, so no resampling is
    needed here on that path.

  * ITS rate, chunk by chunk, forever. That needs
    :class:`Pcm16StreamResampler`: anti-aliased (a downsample without a
    low-pass filter folds everything above the output Nyquist back into
    the speech band) and *stateful* (a per-chunk resample restarts its
    interpolation phase at every chunk boundary, so the same audio would
    decode differently depending on how the client split it). See that
    class's docstring for the measured numbers.

  * OUTPUT — the token2wav decoder emits 24 kHz PCM16. If the session
    asked for a different output rate we resample with a small
    linear-interpolation resampler (documented caveat: linear, not
    polyphase — adequate for near-integer ratios like 24k->16k, cheap
    on the event loop). Same-rate is a zero-copy pass-through.

Everything here is numpy + stdlib only (torch-free) so the session layer
imports on CPU/CI without the model stack.

Resample raw mono PCM16 bytes from ``src_rate`` to ``dst_rate``.

Linear interpolation, ONE-SHOT: the whole buffer is mapped from its
own first sample to its own last, so this is only correct for a
complete buffer. A live stream must use :class:`Pcm16StreamResampler`
instead — see its docstring.

Same-rate (or empty) input is returned unchanged. Used on the OUTPUT
path, when the session's requested output rate differs from the
decoder's native 24 kHz.

Anti-aliased, chunk-boundary-invariant PCM16 resampler for

The streaming counterpart of :func:`resample_pcm16`, and the one the
duplex audio-in path needs. :func:`resample_pcm16` is a one-shot
linear interpolation over whatever buffer it is handed, which makes it
wrong twice over for a live mic stream:

  * **No anti-alias filter.** Linear interpolation is a poor low-pass.
    Downsampling 24 kHz -> 16 kHz drops the Nyquist from 12 kHz to
    8 kHz, and everything above 8 kHz — sibilants, room noise — folds
    back into the speech band instead of being filtered out. On a
    10 kHz tone, which must vanish, it reappears at 6 kHz at
    **-9.1 dBFS**; through this class's kernel, **-49.8 dBFS**. Those
    40 dB are what the perception encoder would otherwise be fed.
  * **Phase resets at every chunk.** ``linspace(0, n-1, n_out)`` pins
    each buffer's own endpoints, so chunk *k* restarts the
    interpolation phase from zero rather than continuing the stream's,
    and emits ``round(n * dst / src)`` samples rather than the exact
    ratio — a clock that runs fast by **73 ms over a 300 s session** at
    the admin console's 2048-sample-at-24 kHz cadence. The
    :class:`~arbi_serve.multimodal.preprocess.nemotron_voicechat.MelStreamer`
    this feeds goes to real lengths to be invariant to how a WS client
    splits its audio, and a stateless resample ahead of it would
    destroy that invariance before the mel front end saw a sample.

This class instead carries the raw input tail between calls and runs
the same windowed-sinc polyphase kernel
``torchaudio.functional.resample`` uses at its defaults
(``lowpass_filter_width=6``, ``rolloff=0.99``, Hann-squared window) —
the same kernel family
``arbi_serve.multimodal.preprocess.nemotron_voicechat._resample``
already applies on the offline path, reimplemented in numpy so this
module stays torch-free (see the module docstring).

Output sample ``j = b*new + p`` is a fixed dot product over raw input
``[b*orig - width, b*orig + width + orig)``, exactly as in the offline
kernel, so the streamed result is bit-comparable to resampling the
whole stream at once — invariant to chunking by construction, not by
tolerance. Only the tail differs: a block is emitted when its
rightmost tap has actually arrived, so the stream trails the offline
result by at most ``width + orig`` input samples (0.5 ms at 24 kHz).
A live stream has no end to pad against, the same asymmetry
``MelStreamer`` documents.

One instance per connection per rate pair; not thread-safe, not
shareable.

Append raw mono PCM16 bytes; return the output samples they
completed, as PCM16 bytes (``b""`` when not enough has arrived yet).

A chunk split mid-sample carries its dangling byte to the next
call rather than dropping it — dropping one byte would shift every
following sample by half a sample for the rest of the stream.

Online speaker diarization for the streaming transcriber.

A live hearing has a small, *a priori* cast — a witness, a examiner, maybe a
judge — and the transcript is far more useful when each utterance carries *who
said it*. Step-Audio-2 is a single-stream model with no notion of speaker
identity, so diarization is a separate, acoustic problem.

We solve it the standard way, reusing a model the checkpoint already ships for
voice cloning: **CAM++** (``token2wav/campplus.onnx``), a speaker-verification
network that maps an utterance to a 192-d x-vector where same-speaker vectors
are close in cosine space. Per finalized VAD segment we extract one x-vector and
assign it online to the nearest speaker centroid (cosine ≥ ``threshold``), or
open a new speaker. Centroids are running means, so a speaker's model sharpens as
they talk more. This is exactly how CAM++/3D-Speaker diarization pipelines work,
minus the offline global clustering — the price of emitting a label *now* rather
than after the recording ends.

Deliberate limits (be honest about them):
  * Short utterances give noisy x-vectors. Segments below ``_MIN_EMBED_MS`` inherit
    the previous speaker rather than risk a spurious new one.
  * Online assignment cannot re-label the past: an early mislabel stays. A courtroom
    product would reconcile offline; a live caption cannot.
  * Overlapping speech (two people at once) is one segment → one label. The VAD cuts
    on silence, not on speaker change.

The module degrades to a no-op (every ``assign`` → ``None``) if onnxruntime or the
campplus model is unavailable, so diarization is always optional and never fatal.

Assign a stable ``"S1"``/``"S2"``/… label to each utterance, online.

Reuse across a stream: construct once, call :meth:`assign` per finalized
segment (never on partials). Not safe for concurrent streams — one per
transcriber.

THIRD tier / §7 step 9): the WS-side half of
routing a live ``input_audio_buffer.append`` PCM stream into a
:class:`~arbi_serve.realtime.duplex_lane.DuplexConnection`'s inbound
frame queue.

=== What this module is, and is not ===

:class:`DuplexAudioIn` is the model-agnostic glue between
:meth:`~arbi_serve.realtime.session.RealtimeSession._on_append` (the WS
message handler) and
:attr:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter
.encode_audio_chunk` (the model-specific "raw PCM -> frame embedding(s)"
step): resample the client's PCM to the model's expected input rate, call
through the adapter, and push whatever frame embeddings come back onto
the connection's inbound queue
(:meth:`~arbi_serve.realtime.duplex_lane.DuplexConnection
.push_audio_frame`) — the SAME queue :mod:`arbi_serve.realtime.duplex_lane`
's ``prepare_tick`` hook already pops from every tick.

It does NOT implement any model's actual audio encoding — that is
squarely ``encode_audio_chunk``'s job, and for NemotronVoiceChat that
field is currently unset (see
:mod:`arbi_serve.realtime.nemotron_voicechat_duplex_adapter`'s own
docstring for the grounded reason). :meth:`DuplexAudioIn.push_pcm16`
fails loud — a clear ``RuntimeError``, not a silently-dropped or
silently-wrong frame — when the served model's adapter has no
``encode_audio_chunk``, exactly the "fail loud on what you did not wire"
discipline the rest of this codebase uses (``docs/adding-a-model.md``
§3, ``resolve_duplex_adapter``'s own ``SUPPORTS_DUPLEX_MODE`` gate).

Feeds one live duplex connection's inbound frame queue from raw,
still-arriving WS PCM16 audio.

One instance per connection, built alongside its
:class:`~arbi_serve.realtime.duplex_lane.DuplexConnection` (see
:meth:`~arbi_serve.realtime.session.RealtimeSession._enable_duplex_mode`).

===

``encode_audio_chunk`` is real model work — for NemotronVoiceChat a
24-layer streaming Conformer step costing **12.5 ms per 80 ms encoder
frame** on an idle RTX 4090 (72% of that kernel-launch overhead; design
doc §7.15). The engine loop owns its own thread
(``ARBI_ENGINE_OWN_THREAD``), so the thread calling in here is
uvicorn's HTTP/WS loop — the one that also writes
``response.audio.delta`` frames out to the client. Inline, that work
stalls the loop: a 5 ms service heartbeat measures p99 13.9 ms / max
21.1 ms with 11.1% of service opportunities overshooting 5 ms, against
p99 1.1 ms / max 1.4 ms on an idle loop. On a worker thread it measures
p99 1.1 ms / max 1.6 ms — idle behaviour, because torch releases the
GIL around each ATen op — so that is where it runs.

Ordering is safe: ``openai_realtime``'s receive loop awaits
``handle_client_event`` before reading the next frame, so a connection
never has two chunks in flight, and the per-connection
``MelStreamer``/``PerceptionStreamState`` (neither thread-safe) are
touched by one thread at a time.

Route one ``input_audio_buffer.append`` chunk (raw PCM16 at
``src_sample_rate``) into the duplex lane.

Resamples to :attr:`DuplexModelAdapter.duplex_input_sample_rate`
if it differs from the client's own rate — through a stateful,
anti-aliased :class:`~arbi_serve.realtime.audio.Pcm16StreamResampler`
carried across chunks, NOT a per-chunk one-shot (that would both
alias and reset its interpolation phase at every chunk boundary;
see that class's docstring) — then calls
:attr:`DuplexModelAdapter.encode_audio_chunk` — stateful per
connection (the adapter's own session-keyed bookkeeping decides
how many NEW frame embeddings, if any, this chunk completes) —
and pushes every returned ``(frame_embed, semantic_blank)`` pair
onto the connection's inbound queue, paired with the CONTIGUOUS
span of resampled PCM that frame stands for
(:meth:`_frame_spans`) — the acoustic evidence the barge-in
AND-gate and the The
per-frame ``semantic_blank`` verdict (design doc §7.32) is
computed per encoder frame and rides through untouched.

Resample + encode run on a worker thread (see class docstring);
the queue push happens back on the caller's loop, so the lane only
ever sees frames appended in arrival order.

Raises ``RuntimeError`` if the served model's adapter has no
``encode_audio_chunk`` (see module docstring) — surfaced to the
client as an ``error`` event by
``RealtimeSession.handle_client_event``'s existing exception
handling, not a crash.

Deal this chunk's resampled PCM out across the ``n_frames``
frames it completed, so the audio riding the frame queue is the
WHOLE stream, in order, with nothing dropped.

The naive pairing — hand every frame the whole chunk that
happened to complete it — silently discards most of the
microphone. A client streaming 30 ms chunks completes one 80 ms
encoder frame roughly every third chunk, and the two chunks that
completed no frame push no PCM at all: measured on this
checkpoint's own bundled sample conversation, **the acoustic VAD
received 37% of the audio**, in 30 ms fragments stitched across
50 ms holes. Silero's speech probability never once reached its
0.5 gate on any of the five real user utterances there (peak
0.385, and 0.017 on one of them), so ``eou_user_spoke`` could
never be set and every candidate BOS after a connection's first
turn was refused for want of confirmed speech — design doc §7.47.

Dealing the chunk out instead keeps the CONCATENATION exact,
which is the only property a streaming detector's hysteresis
actually depends on; where one frame's span ends and the next
begins wobbles by under a frame, because the encoder's own mel
carry is not visible from here. That is the same tolerance the
pairing always had, and it is real: the VAD needs
roughly-aligned acoustic evidence, not sample-exact alignment.

A chunk that completes NO frame is not dropped either — it is
held in :attr:`_pcm_carry` and dealt out with the chunk that
completes the next one.

The VAD x Site-3 barge-in

Extracted from
:func:`~arbi_serve.realtime.duplex_driver.drive_duplex_or_turn_session`'s
own per-tick body so BOTH per-tick execution paths reach the SAME
implementation rather than carrying a copy each:

  * the PULL driver (:mod:`arbi_serve.realtime.duplex_driver`), where an
    external caller drives ticks explicitly;
  * the PUSH hook (:mod:`arbi_serve.realtime.duplex_lane`), where
    ``run_forever``'s own always-on loop drives the tick and the barge-in
    check has to run in its post-step hook (design doc §7.9).

===

Two INDEPENDENT signals on two DIFFERENT channels, combined with an
AND-gate — not one mechanism doing double duty:

  * the VAD (:class:`~arbi_serve.realtime.vad.VadDetector`, the same
    model-agnostic infrastructure already driving Step-Audio-2's
    turn-based barge-in) reads the HUMAN's incoming audio (input
    channel): did the human just start making speech sounds, purely
    acoustically;
  * ``agent_idle`` reads the ASSISTANT's own generated tokens (output
    channel) via whatever model-specific mechanism produced this tick's
    result (NemotronVoiceChat: the Site-3

Only both together are a barge-in. If the assistant is already idle,
incoming speech is ordinary listening — no action.

=== Why the VAD is fed even when the gate cannot fire ===

:meth:`VadDetector.feed` is STATEFUL (it tracks speech/silence run
lengths across frames to decide when a ``speech_started`` edge happens).
Skipping the feed on ticks where the assistant happens to be idle would
desync that state from the real audio timeline and mis-time the NEXT
edge. So the feed is unconditional whenever a frame is available, and
only the RESULT is ANDed with ``not agent_idle`` — the exact ordering the
driver loop's original inline implementation had.

=== One feed, two consumers (design doc §7.30) ===

:func:`vad_tick_events` is now the ONE place a tick's PCM is fed into the
VAD — :func:`vad_barge_in_fires` takes the resulting event list rather
than feeding the VAD itself, so a second caller (the EOU-gated BOS
admission check, :meth:`~arbi_serve.realtime.nemotron_voicechat_duplex_session.NemotronVoiceChatDuplexSession.note_vad_events`)
can read the SAME tick's events without feeding the VAD a second time,
which would desync its hysteresis state exactly as skipping a feed
would. :meth:`DuplexTickPump.post_stt` is the one caller that does both.

Feed this tick's PCM into the VAD exactly once; return whatever
transitions it reports (``[]`` when there is no VAD/frame this tick).

The single source of truth for feeding the VAD each tick — every
consumer of this tick's VAD signal (the barge-in AND-gate, the
EOU-gated BOS admission check) reads the list this returns rather
than calling :meth:`VadDetector.feed` itself.

The real per-tick duplex/turn driver loop — model-agnostic.

STT decode step then one real audio-out (TTS) step,
until a termination condition fires. This module knows NOTHING about any
specific model — it calls entirely through
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter`,
resolved generically via
:func:`~arbi_serve.realtime.duplex_model_adapter.resolve_duplex_adapter`.
NemotronVoiceChat-specific, directly
importing ``step_duplex_tick``/``run_turn``/
``NemotronVoiceChatDuplexSession``; that coupling has been extracted into
the adapter seam so a future duplex-capable model needs no change here,
only a new adapter registration).

=== Where this module's per-tick body actually lives now ===

Everything that happens AFTER the STT step — the bridging described
below, the VAD barge-in SHARED with the push-driven duplex lane
(:mod:`arbi_serve.realtime.duplex_lane`), which is the production shape.
This loop is now just "await the
STT step, hand the result to the pump, yield what comes back" — kept for
standalone/offline drivers and tests that own the engine exclusively (see
this function's own warning about running it against a live
``run_forever``).

=== Bridging an async STT step into a sync audio-out generator ===

A model WITH an audio-out side supplies
``adapter.create_audio_out_driver``, which builds a plain, synchronous
generator shaped exactly like
:func:`~arbi_serve.runtime.nemotron_voicechat_turn.run_turn` (that
function's own docstring explains why it is sync, not async: "does no
I/O of its own"). ``adapter.step_stt`` is ``async`` (it drives the real
engine's own async step path). This module bridges the two with a
one-slot mailbox: for tick ``i``, :func:`drive_duplex_or_turn_session`
``await``s ``adapter.step_stt`` FIRST, converts its
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexStepResultLike`
into an :class:`~arbi_serve.runtime.nemotron_voicechat_turn.SttStepResult`,
stashes it in the mailbox, THEN advances the audio-out generator with
``next()`` — whose synchronous ``stt_step_fn`` callback just reads the
already-computed mailbox value back out. (A production caller wanting the
blocking ``next()`` calls off the event loop can wrap them in
``asyncio.to_thread`` — the same later-caller choice ``run_turn``'s own
docstring already flags.)

Per that generator's own shape, one call to ``next()`` yields
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta`, the next
yields :class:`~arbi_serve.runtime.nemotron_voicechat_turn.AudioChunk` —
i.e. TWO ``next()`` calls complete one tick's audio-out half (the FIRST
of the pair is also what triggers ``stt_step_fn`` for that same tick) —
see :meth:`~arbi_serve.realtime.duplex_tick_pump.DuplexTickPump._advance_audio_out_one_tick`.

A model WITHOUT an audio-out side (``adapter.create_audio_out_driver is
None`` — a real, supported case, e.g. a hypothetical ASR-only duplex
model) skips this bridge entirely: this loop yields a bare
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta` per tick
and synthesizes its own
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TurnDone` when the
termination condition fires, so callers never need to branch on whether
the served model happens to have a TTS side.

Note: :class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta`/
``AudioChunk``/``TurnDone``/``SttStepResult`` currently live in a
NemotronVoiceChat-named module for historical reasons but are themselves
model-agnostic in shape (plain ``token_id``/``text``/``waveform``/
``finish_reason`` fields, nothing NemotronVoiceChat-specific) — this loop
reuses them as the established generic event vocabulary rather than
duplicating them. Moving them to a neutral module
(e.g. ``arbi_serve/realtime/duplex_events.py``) is a clean, low-risk
follow-up, deferred here to keep this pass scoped to the adapter seam
itself.

=== Lifecycle / termination: duplex (unbounded) vs. turn-shaped (bounded) ===

Per the design doc's own framing (§2.2: "every tick always has a frame
[real or silence-padded], forever" for a genuine duplex connection) there
is no natural per-turn EOS in duplex mode — the WS connection's own
lifetime (or the KV/state cap, design doc §7 step 8, not yet built) is
what ends it, not this loop. A bounded, turn-shaped session (single-turn
mode, repointed onto this same machinery per the confirmed product
decision — see the call site's own docstring) DOES need a termination
condition, and this module supplies one deliberately modeled on
NemotronVoiceChat's own
``_FrameLockstepState.should_stop_turn``'s existing shape ("turn-stop
token sampled AND no audio frames remain") — generalized to read
``agent_idle`` off :class:`~arbi_serve.realtime.duplex_model_adapter
.DuplexStepResultLike` (whatever model-specific mechanism the adapter's
``step_stt`` used to decide it) instead of a raw per-step EOS-token
check:

    is_eos_this_tick = bounded
        and no_more_real_audio_frames_left
        and agent_has_gone_active_at_least_once
        and this_ticks_agent_idle_is_true

i.e. the turn only ends once the input audio is exhausted AND one full
active-then-idle cycle has completed. ``agent_has_gone_active_at_least_once``
guards against a degenerate turn that never produces a real response
being mistaken for "done" on tick 0, when ``agent_idle`` is still just
its initial ``True`` value.

A ``max_ticks`` safety valve (mirrors ``run_turn``'s own ``max_tokens``)
always applies, in BOTH modes — for ``bounded=False`` (genuine duplex) it
is the only stop condition this loop itself ever applies; a real
always-running connection is expected to pass an effectively-unbounded
value and rely on the connection lifetime /

``audio_frame_embeds`` must already be at least ``max_ticks`` frames long
(the SAME whole-connection precomputed-array contract NemotronVoiceChat's
own ``step_duplex_tick`` already has — a future adapter's own ``step_stt``
sets its own contract for this argument; this loop just passes it
through unchanged). For a bounded turn, the caller is expected to pad
real audio with trailing silence-frame embeddings out to at least
``max_ticks`` frames — ``num_real_audio_frames`` tells this loop where
the real input ends and padding begins, for the termination check above
only.

=== VAD-driven barge-in

Two INDEPENDENT signals on two DIFFERENT channels, combined with an
AND-gate — not one mechanism doing double duty:

  * The VAD (:class:`~arbi_serve.realtime.vad.VadDetector`, the SAME
    existing, model-agnostic, already-proven infrastructure that drives
    Step-Audio-2's cancel-and-restart barge-in via
    ``RealtimeSession._on_speech_started`` — see that class's own
    docstring) reads the HUMAN's incoming audio (input channel): did the
    human just start making speech sounds, purely acoustically. Fed the
    SAME per-tick raw PCM (``audio_frame_pcm[tick_idx]``) that
    ``audio_frame_embeds[tick_idx]`` is the already-computed feature
    projection of — ``audio_frame_pcm`` runs ALONGSIDE
    ``audio_frame_embeds`` through this loop, never instead of it; the
    ordinary STT/TTS ticking above is completely unchanged whether or
    not a VAD is supplied.
  * ``session.agent_idle`` (:class:`DuplexSessionLike`, this loop's own
    minimal structural contract) reads the ASSISTANT's own generated
    tokens (output channel) via whatever model-specific mechanism the
    adapter's ``step_stt`` used to derive ``tick_result.agent_idle`` this
    tick (NemotronVoiceChat: the Site-3 This
    loop never inspects how that value was produced — same
    Protocol-level agnosticism the rest of this module already has.

Evaluated every tick, right after ``tick_result`` is known and BEFORE it
feeds the rest of this tick's bookkeeping (mirrors this loop's own
existing "update-then-decide" ordering for ``agent_was_active``/
``is_eos_this_tick``): if the VAD reports a ``speech_started`` event for
THIS tick's PCM AND the assistant is not already idle
(``not tick_result.agent_idle``), force the escape hatch
(``session.mark_agent_idle()``) and treat this tick's effective
``agent_idle`` as ``True`` for everything downstream in this same tick
(``agent_was_active``/``is_eos_this_tick``/the mailbox result) — so the
forced interrupt has real, immediate effect rather than only being
visible starting next tick. If the assistant is already idle, incoming
speech is ordinary listening, not a barge-in: no call, no state change
(``mark_agent_idle`` is itself idempotent, but the gate does not even
reach it in this case — matching the design doc's own "no action").
Both ``vad`` and ``audio_frame_pcm`` default to ``None``; the gate is
fully inert (zero behavior change, zero extra work) unless a caller
supplies both, so every existing caller of this function is unaffected.

This mechanism is entirely model-agnostic (the VAD operates on raw
audio; ``agent_idle`` is a session-level flag any duplex-capable model's
adapter could set) and therefore lives here, in the generic driver, not
in any model-specific adapter — per the design doc's own "belongs in the
model-agnostic part of the driver loop" framing.

Drive one real, scheduler-owned duplex ``Request`` one tick at a
time — real STT decode (``adapter.step_stt``) THEN one real audio-out
tick (``adapter.create_audio_out_driver``'s generator, if the model
has one) — until the termination condition fires (see module
docstring). An async generator: callers iterate this function's
return value with ``async for`` the same way they would drive a
model's own audio-out generator with plain ``next()``, one
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta`
(/:class:`~arbi_serve.runtime.nemotron_voicechat_turn.AudioChunk`,
when the model has an audio-out side) pair per tick, finally a
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TurnDone`.

``adapter``: resolve via
:func:`~arbi_serve.realtime.duplex_model_adapter.resolve_duplex_adapter`
— this function itself never constructs one, keeping it agnostic to
which model ``eng``/``model`` actually are.

``session``: an already-admitted session (``session.request`` already
admitted into the scheduler via this model's own admission path) —
or ``None`` to have this call build a fresh one via
``adapter.new_session()`` (the caller is then responsible for
whatever admission that model's own session type needs before the
first tick, exactly as when a caller supplies its own session).

VAD-driven barge-in AND-gate
(see module docstring's own section) — both optional and ``None`` by
default (no behavior change for existing callers). When both are
given, ``audio_frame_pcm[tick_idx]`` is this tick's raw PCM16 bytes,
index-aligned with ``audio_frame_embeds`` (the SAME per-tick frame,
fed through the VAD alongside, not instead of, the feature path
``adapter.step_stt`` already consumes it through). Shorter than
``max_ticks`` is fine — ticks beyond ``len(audio_frame_pcm)`` simply
skip the VAD check (e.g. a caller that only has real client audio for
part of a bounded turn's padded tail).

``tool_choice`` enforcement for the duplex

What this constrains
--------------------

NemotronVoiceChat decodes two heads per 80 ms tick: a TEXT channel (what the
agent says) and a separate FUNCTION channel (where ``<TOOLCALL>[...]``
blocks appear). This module masks the FUNCTION channel only. The text
channel's decode is untouched — the two heads are separate tensors
(``Request.logits_heads``), so scoping the constraint is a matter of which
tensor the bitmask is applied to, not of a flag anyone can forget.

Why this is NOT wired through ``SamplingParams``
-----------------------------------------------

The obvious-looking wiring — put ``tools``/``tool_choice`` on the duplex
:class:`~arbi_serve.engine.request.SamplingParams` and let the ordinary
grammar path pick them up — is actively WRONG here, and quietly so.
:attr:`~arbi_serve.engine.request.SamplingParams.needs_grammar` is
``tool_choice_enforces``, and the grammar the ordinary path attaches is
applied by the sampler to the request's TEXT logits. A duplex connection
carrying tools would therefore have its spoken text forced into
``<TOOLCALL>`` syntax. So the tool grammar is held HERE, on the connection's
duplex session object (alongside ``fc_state``, the same per-connection
function-channel state this sits next to), and applied by hand at the one
place the function channel is committed
(:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`).

Arming: why ``"required"`` cannot simply run from connection open
-----------------------------------------------------------------

Measured against the real checkpoint tokenizer + xgrammar 0.2.3, the three
modes behave very differently at their FIRST masked step:

  * ``"auto"`` — ``fill_next_token_bitmask`` returns ``need_apply=False``
    and all 131072 tokens are allowed. It stays that way until the model
    itself opens a ``<TOOLCALL>``, at which point the mask forces the block
    to be well-formed and parseable, then goes permissive again.
  * ``"required"`` / named — ``need_apply=True`` with **2 of 131072** tokens
    allowed: the trigger is forced IMMEDIATELY.

The duplex function channel decodes every tick from connection open, so a
``"required"`` matcher armed at connection open would force a ``<TOOLCALL>``
on tick 0 — a tool call with invented arguments, emitted before the user has
said a word. Enforcing modes are therefore ARMED (see :meth:`note_user_spoke`)
only once the connection has confirmed real user speech, and the matcher is
:meth:`reset` at each agent turn boundary so the obligation is per-response
rather than once-per-connection — which is also how OpenAI scopes
``tool_choice`` (it is a per-response field there).

``"auto"`` needs no arming: its mask is a genuine no-op until the model
opens a tag on its own, so it runs from connection open and costs one
CPU-side bitmask fill per tick.

``"none"`` builds no grammar at all (:meth:`build` returns ``None``), which
is also what a session with no ``tools`` gets.

Per-connection xgrammar matcher over one duplex FUNCTION channel.

Built by :meth:`build` from the session's ``tools`` + ``tool_choice`` and
parked on the duplex session object; driven per tick by
:meth:`mask_logits_` (before the channel's argmax) and :meth:`accept`
(after it, with the token that actually entered the recurrence).

One matcher per connection, not per tick: the FUNCTION channel is one
unbroken token stream for the connection's life (the recurrence never
restarts), and a :class:`xgrammar.GrammarMatcher` is exactly the
automaton state for one such stream. The compiled grammar behind it is
shared and cached process-wide by
:class:`~arbi_serve.sampler.xgrammar_processor.XGrammarLogitsProcessor`,
so a second connection offering the same tools pays no compile.

Build the connection's function-channel grammar, or ``None``.

``None`` — meaning "decode this channel exactly as before" — for a
session with no ``tools``, an explicit ``tool_choice="none"``, no
engine grammar processor, or a tag/compile failure. A failure
DEGRADES rather than raising: an unenforceable grammar must not take
down a live voice connection, and ``"auto"`` (the default) loses
nothing real when it degrades, since the model already calls tools
correctly unenforced.

``xgrammar_proc`` is the engine's own
:class:`~arbi_serve.sampler.xgrammar_processor.XGrammarLogitsProcessor`
(``Engine.xgrammar``) — reused rather than building a second
compiler so the :class:`xgrammar.TokenizerInfo` is the SAME one the
chat-completions path enforces against, and the compiled-grammar
cache is shared.

Arm an enforcing mode — the connection has confirmed user speech.

Idempotent and one attribute store on the hot path. Deliberately a
latch, not a mirror of the session's own ``eou_user_spoke``: that
flag resets at every turn boundary, and a mask that disarmed mid-
block would leave a half-emitted ``<TOOLCALL>`` that never closes —
precisely the stale-open state
``NemotronFunctionChannelExtractor._abandon_stale_open`` exists to
clean up after. :meth:`reset` is what re-arms the next window.

KNOWN COUPLING: an enforcing ``tool_choice`` is only as reliable as
whatever calls this. Its one caller feeds it ``eou_user_spoke``, and
design doc §7.39 records that flag staying If that
regresses, ``"required"`` silently enforces nothing. It is still the
right signal — it is the same flag ``eou_admits_bos`` gates
turn-opening on, so a connection where it is stuck cannot take turns
at all — but it is a dependency, not an assumption, and it is why
the live proof runs over real speech audio rather than synthetic
frames (see design doc §7.65.2).

Start a fresh obligation window (called at agent turn boundaries).

Drops the matcher's automaton state back to the grammar's start, so
an enforcing mode owes a NEW call next window rather than being
permanently satisfied by one made an hour ago, and any abandoned
half-open tag from an interrupted turn is discarded. Enforcing modes
go back to disarmed and wait for the next :meth:`note_user_spoke`.

Mask the disallowed function-channel tokens IN PLACE.

``function_logits`` is this tick's ``(vocab_size,)`` function head
(``take_dual_head_row``'s second element). Returns whether a mask was
actually applied — ``False`` is the overwhelmingly common case and
means the tensor was not touched at all.

The ``False`` fast path is xgrammar's OWN, not a heuristic of ours:
``fill_next_token_bitmask`` returns ``need_apply``, documented as
"the bitmask is already all-true, so no need to apply it". For
``"auto"`` that is every tick until the model opens a ``<TOOLCALL>``,
so the steady-state cost of the default policy is one CPU-side fill
per tick and no device traffic whatsoever.

Advance the matcher with the token that entered the recurrence.

Called with the SAME id ``finish_duplex_tick`` feeds to
``session.observe`` and to the detector, so the automaton state and
the model's own history can never disagree.

A rejected token is logged, not raised: the function channel commits
by argmax over the (possibly unmasked, e.g. unarmed ``"required"``)
head, so a token outside the grammar is reachable by construction and
must not kill a live connection. Once the matcher and the stream have
diverged the matcher stops constraining until the next :meth:`reset`.

The duplex lane: per-connection state + the two hooks ``run_forever``
calls around its

=== Why hooks, and not a driver loop ===

Design doc §7.8 recorded the blocking finding this module resolves: the
scheduler ALREADY drives duplex-lane requests autonomously. ``Scheduler
.schedule()`` checks the duplex-lane deadline before any ordinary slate
work and, when due, returns ``_build_duplex_slate``'s duplex-only slate
(``scheduler.py``'s exclusivity gate + ``scheduler_selection.py``'s
``duplex_tick_due``/``_build_duplex_slate``); ``run_forever`` calls
``schedule()`` unconditionally on every iteration, on a dedicated engine
thread that every real server boot spawns. So for a duplex-admitted
``Request``, **the engine's own unmodified loop already runs the STT
forward pass every 80 ms with no external caller at all.**

That makes the PULL driver (:mod:`arbi_serve.realtime.duplex_driver`) the
wrong production shape, not merely a slower one: a WS handler that ALSO
called ``schedule()``/``run_step_async`` for the same request would
double-step its KV/Mamba state against ``run_forever``'s own stepping of
it. Nothing about thread-marshaling fixes that — ``run_forever`` would
still independently re-discover and re-drive the request.

So this module adds no loop. It adds the two things the autonomous step
is genuinely missing, as hooks AROUND it:

  * :meth:`DuplexLane.prepare_tick` — the PRE-SCHEDULE hook. For every
    live connection whose request is admitted, take the next audio frame
    (real, from the connection's inbound queue; or silence when the
    client is behind) and park its fused embedding on
    ``request.pending_embed_override``, so the slate ``schedule()`` is
    about to build carries this tick's real audio. Also arms that
    request's multi-head retention (design doc §7.42).
  * :meth:`DuplexLane.complete_tick` — the POST-STEP hook. For every
    connection whose row was actually in the slate, run the model's own
    per-tick bookkeeping (feedback tokens + turn-open/turn-close switch),
    then the shared :class:`~arbi_serve.realtime.duplex_tick_pump
    .DuplexTickPump` (barge-in AND-gate, one audio-out/TTS tick, event
    synthesis), and publish the resulting events to the connection's
    outbound queue.

Everything model-specific goes through
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter`;
this module never imports NemotronVoiceChat.

=== Zero cost when no duplex session exists ===

The engine carries a single ``eng.duplex_lane`` slot, ``None`` until the
first connection registers and set back to ``None`` when the last one
leaves. Both hooks are guarded by one ``is not None`` test in
``run_forever``'s loop body, so an ordinary deployment pays one predicted-
not-taken attribute read per iteration and nothing else. The scheduler's
own duplex gate is likewise fast-outed by a monotone
``_duplex_ever_admitted`` latch (``scheduler_admission.py``), which
actually REMOVES two per-step ``running``/``waiting`` scans that gate used
to run unconditionally for every deployment.

=== Threading ===

``run_forever`` runs on the dedicated ``arbi-engine-loop`` thread
(``ARBI_ENGINE_OWN_THREAD``, default ON) while WS handlers run on the
HTTP loop, so both queues cross a thread boundary:

  * inbound (HTTP -> engine): a :class:`collections.deque`, appended by
    the WS handler and ``popleft``-ed by the engine thread. ``deque``'s
    append/popleft are atomic under the GIL — the same lock-free
    single-producer/single-consumer discipline
    :class:`~arbi_serve.engine.loop_bridge.LoopBridge`'s own intake queue
    uses.
  * outbound (engine -> HTTP): a bounded ``deque`` plus ONE
    :meth:`~arbi_serve.engine.loop_bridge.LoopBridge.notify_http` hop per
    TICK to set the consumer's ``asyncio.Event`` — deliberately mirroring
    :class:`~arbi_serve.engine.output_bus.OutputBus`'s "one hop per step,
    never per item" invariant, scoped to a connection rather than routed
    through the global ``OutputBatchMsg`` pipeline (whose ``TokenOut``/
    ``AudioOut(codes)`` structs do not fit a raw waveform tensor, and
    whose applier sits on the hot path for all ordinary traffic). The
    crossing itself happens inside the marshaler seam, never here: this
    module contains no ``call_soon_threadsafe`` of its own, which is what
    ``tests/test_transport_boundary.py``'s engine<->HTTP invariant
    requires of every module outside ``output_bus``/``loop_bridge``/
    ``engine_thread``/``engine/proc``.

The connection registry itself is guarded by a plain lock and snapshotted
to a list inside each hook, so a register/unregister racing a tick can
never mutate the sequence being iterated.

=== Module layout (this module re-exports the public API) ===

This module holds :class:`DuplexLane` (the registry + the two hooks)
only. Its two collaborators live in sibling modules and are re-exported
here so every existing import site (``from
arbi_serve.realtime.duplex_lane import ...``) keeps working unchanged:

  - :mod:`arbi_serve.realtime._duplex_connection` —
    :class:`~arbi_serve.realtime._duplex_connection.DuplexConnection`,
    one live connection's lane-side state.
  - :mod:`arbi_serve.realtime._duplex_admission` — the connection
    lifecycle free functions (:func:`attach_duplex_connection`,
    :func:`publish_duplex_connection`, :func:`close_duplex_connection`,
    :func:`warn_if_still_scheduled`) plus the engine-thread marshaling
    primitive they and :class:`DuplexLane` share.

PRE-SCHEDULE hook — see :meth:`_prepare_tick_inner`.

Both hooks are wrapped so that NOTHING escaping the lane can reach
``run_forever``. A hook is called from the engine's own run loop,
outside its per-step try/except, so an exception here does not
"fail one request" — it kills the engine thread and takes down
serving for every request on the box. Observed live:
a single missing method on the notification path propagated out of
``complete_tick`` and terminated the run loop. The lane may lose a
duplex connection; it must never lose the engine.

Close every registered connection whose Request the ENGINE has
already finished.

Whatever finished it — ``run_forever``'s own except branch (which
finishes EVERY row in the failing slate), the timeout sweep, an
explicit cancel, or ``cache.max_context`` in
``engine/run_step/terminal.py`` — the Request is already torn down
at the engine/scheduler level (pages freed, out of
``running``/``waiting``). Only this lane's bookkeeping is left, and
until it is released the connection's consumer has no way to learn
that no further tick can ever arrive.

Called before :meth:`_prepare_tick_inner`'s due gate on purpose —
see the call site. Runs on the engine thread once per step in a
duplex-serving deployment, over a dict that holds one entry per
live WS connection; the hot path for every non-duplex deployment
is unchanged (``eng.duplex_lane`` is ``None``, so neither hook is
called at all).

Give every live connection this tick's audio.

Called from ``run_forever`` immediately before
``eng.scheduler.schedule(_now=now)`` — with the SAME ``now``, so
the readiness question this method asks (``duplex_tick_due``) and
the one ``schedule()`` asks cannot disagree across a deadline
crossing.

A no-op unless the duplex lane is actually due this step; ordinary
steps in a duplex-serving deployment therefore still cost only the
one predicate call.

Put a prepared-but-never-stepped tick's frame back at the FRONT
of the inbound queue, oldest-first order intact.

``_next_frame`` ``popleft``s the frame in the pre-schedule hook, so
a tick the scheduler then declines (no page room) or that never ran
(empty slate, a step that raised) would otherwise drop 80 ms of the
client's audio on the floor — silently, since nothing downstream
can tell a dropped frame from one the client never sent. This
method is what makes ``_complete_tick_inner``'s own "no frame is
consumed" contract true rather than merely intended. A SILENCE
frame is never re-queued (``_next_frame`` leaves
``prepared_frame_embed`` ``None`` for one): it was synthesized to
fill a gap, and re-queuing it would push real audio that arrived
meanwhile behind a filler frame.

POST-STEP hook — see :meth:`_complete_tick_inner`. Wrapped for
the same reason :meth:`prepare_tick` is: a lane bug must never
reach ``run_forever``.

``eng`` is unused since the retained logits heads moved onto the
request (design doc §7.42). It stays in the signature deliberately:
this is one of the two hooks ``run_forever`` calls, they are
documented and called as a symmetric pair (``prepare_tick(eng,
now)`` genuinely needs the engine), and narrowing the contract of a
run-loop hook to save one unused parameter would be a worse trade
than keeping the pair legible.

Run the second half of every prepared tick.

Called from ``run_forever`` after the step for ``slate`` has run —
and also on the paths where no step ran at all (empty slate, or a
step that raised), where it degrades to a clean rollback: the
prepared overrides are cleared and the frame this tick dequeued is
put BACK on the inbound queue, so the same audio is retried on the
next duplex tick instead of being silently dropped.

``slate`` decides which prepared connection actually ran. Its
per-row

One-shot soft KV/state-cap warning (design doc §3.5's 300 s
recommendation) — an OTel counter + one log line, the FIRST tick a
genuine (unbounded) duplex connection's tick count crosses its
configured soft threshold. Purely advisory: never closes the
connection, never repeats, and is a no-op for a bounded
(turn-shaped) session, whose ``max_ticks`` is a short test/turn
safety valve, not this connection-lifetime cap.

Drop a connection whose Request is STILL LIVE at the moment this
fires — an adapter exception mid-tick, a bounded session's own
natural EOS, or the Runs from inside exception handlers (and the ordinary
cap-hit path) on the ENGINE THREAD, so it must never raise: an
exception escaping here propagates out of the hook and kills
``run_forever`` itself, turning one bad connection into a total
engine outage (observed live).

Beyond the lane's own bookkeeping, this ALSO cancels the Request
through the same ``cancel_by_id``/``Scheduler.remove`` path any
other request's teardown uses — freeing its KV/Mamba pages and
removing it from ``Scheduler.running`` for good. Without this, a
connection the LANE decides to stop (as opposed to one the engine
already finished — see :meth:`_close_already_finished`) would
leave its Request orphaned: still ``is_duplex_frame``/``DECODING``,
re-selected by every future GPU/page leak. This was a real, confirmed gap
this task closed: before this fix, cap-hit and adapter-exception
paths alike only marked the lane-local ``conn.closed`` flag.

Close ``conn``'s CURRENT bounded turn WITHOUT tearing the
connection/request down (design doc §7.24) — the persistent-state
counterpart of :meth:`_end_connection`, for a
:attr:`~DuplexConnection.turn_rearmable` connection whose KV/Mamba
state must survive into the next turn.

The Request stays exactly as live as it was the instant before
this call: still admitted, still ``DECODING``, its pages
untouched. All this does is stop the pre-schedule hook from
arming ``conn`` for any further tick (:attr:`DuplexConnection
.awaiting_turn`) until :meth:`DuplexConnection.begin_turn` stages
a re-arm — the same "no override, no frame, cost is one
``is not None``-shaped check" idle state a prefill-phase
connection already sits in between its own hook visits.

Never raises, for the same reason :meth:`_end_connection` does
not: this can run from inside ``_complete_tick_inner``'s own
per-connection loop.

Drop a connection whose Request the ENGINE ITSELF already
finished — an engine-level step error's own finish-every-slate-row
branch, the timeout sweep, an explicit cancel, or the ordinary
``cache.max_context`` finish check in ``engine/run_step/
terminal.py`` (a real, already-live backstop against a duplex
request's unbounded KV growth, just a much looser one than the
lane's own §3.5 cap). Its pages are already freed and it is
already out of ``Scheduler.running``/``waiting`` — calling
``cancel_by_id`` again here would re-run the generic finish
machinery against an already-FINISHED Request, which is not
proven idempotent for every side effect (detok teardown, LoRA ref
release, the terminal output emission) the way a genuine
double-cancel of a still-live request is documented to be. Only
the LANE's own bookkeeping needs releasing. Never raises, for the
same reason :meth:`_end_connection` never does.

Model-agnostic duplex-mode adapter seam.

Mirrors :mod:`arbi_serve.multimodal.registry`'s ``MediaBinding``/
``resolve_mm_bindings`` pattern (a genuine model-agnostic seam already
shared across Step-Audio-2 and NemotronVoiceChat) and
``SUPPORTS_EMBED_OVERRIDE``'s plain-declared-capability-flag convention
(``arbi_serve/runtime/forward_exec.py``), applied to duplex mode: the
per-tick driver loop (:mod:`arbi_serve.realtime.duplex_driver`) never
imports a specific model's fusion method, token ids, or session class —
it only calls through :class:`DuplexModelAdapter`, resolved generically
via :func:`resolve_duplex_adapter`.

NemotronLabs-VoiceChat-11B is the FIRST model to implement this protocol
(``arbi_serve.realtime.nemotron_voicechat_duplex_adapter``,
registered below), not the only one duplex mode is architecturally
allowed to ever support — a future model (or a duplex-capable
Step-Audio-2, if that is ever built) declares
``SUPPORTS_DUPLEX_MODE = True`` on its model class and registers its own
factory in :data:`_ADAPTER_FACTORIES`; nothing in
:mod:`arbi_serve.realtime.duplex_driver` needs to change.

=== Why an adapter object, not more injected callables on the driver itself ===

``run_turn`` already established the "inject the model-specific step
function" pattern for ONE seam (``SttStepFn`` — see that module's own
"Why the STT step is an INJECTED callback" docstring section). Duplex
mode needs FOUR model-specific things, not one: how to build a fresh
session object, how to run one real STT tick, which token ids mean
turn-open/turn-close, and (optionally — a model needn't have a
TTS/audio-out side at all, e.g. a hypothetical ASR-only duplex model)
how to drive one audio-out tick. Bundling these into one small,
explicitly-named adapter object (rather than four more keyword
parameters on :func:`~arbi_serve.realtime.duplex_driver
.drive_duplex_or_turn_session`) keeps that function's own signature
stable as more models implement the protocol, and keeps one model's
whole adapter reviewable/testable as a unit — the same reason
``MediaBinding`` bundles ``encode``/``num_placeholder_tokens``/
``uses_mrope`` instead of the runner taking four separate modality
parameters.

Structural contract the generic driver needs from a per-connection
duplex session object — deliberately minimal (two fields + one
method) so a model's OWN session class (e.g.
:class:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession`, which carries plenty of model-
specific state beyond these) satisfies it for free, with no base
class to inherit from.

VAD-driven barge-in AND-gate, see
:mod:`~arbi_serve.realtime.duplex_driver`'s own docstring) is the
escape-hatch call the generic driver forces when the AND-gate fires —
named to match the reference's own ``s2s_state.py`` API shape (design
doc §3.2.1), and already present on
:class:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession` (idempotent — safe to call regardless
of the current value, matching that method's own docstring).

Structural contract for one STT tick's result — the fields the
generic driver itself reads. Mirrors
:class:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.DuplexTickResult`'s
own shape (that class already satisfies this Protocol unmodified;
this Protocol exists so the driver's own signature does not import
that NemotronVoiceChat-specific module just to name a type).

One model's concrete implementation of the duplex-mode protocol —
everything :func:`~arbi_serve.realtime.duplex_driver
.drive_duplex_or_turn_session` needs that is NOT the same across
every model.

Args:
    turn_open_token_id: the text-channel token id that means "the
        agent just started a turn" (NemotronVoiceChat calls this its
        Site-3 BOS id — see ``resolve_turn_bos_id``; a different
        model's equivalent concept, whatever it is, lands here).
    turn_close_token_id: the turn-close counterpart.
    new_session: builds a FRESH, this-model-shaped session object
        (satisfying :class:`DuplexSessionLike`) — called once per new
        connection when the caller doesn't already have one.
    step_stt: ``async (eng, session, audio_frame_embeds, *, now) ->
        DuplexStepResultLike`` — drives ONE real STT decode tick.
        For NemotronVoiceChat this thinly wraps
        :func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.step_duplex_tick`
        with this model's own ``text_bos_id``/``text_eos_id`` and
        tokenizer already bound.
    create_audio_out_driver: ``(model, stt_step_fn, session, *,
        max_ticks) -> Iterator[TurnEvent]`` or ``None`` — builds this
        model's own per-tick audio-out generator (for NemotronVoiceChat,
        a thin wrapper around
        :func:`~arbi_serve.runtime.nemotron_voicechat_turn.run_turn`
        with this connection's ``audio_prompt_latent``/vocab/etc.
        already bound, taking the driver's own ``stt_step_fn`` mailbox
        closure as ``run_turn``'s injected callback — see that
        module's own docstring). ``None`` means this model has no
        audio-out/TTS side at all (a valid, real case this protocol
        must support — e.g. an ASR-only duplex model); the driver
        then yields text-only events and synthesizes its own
        :class:`~arbi_serve.runtime.nemotron_voicechat_turn.TurnDone`.
    prepare_stt_tick: ``(eng, session, frame_embed) -> None`` — the
        FIRST half of one STT tick: park this tick's fused input
        embedding on the session's request, ready for whoever
        schedules the step. Required only by the PUSH caller
        (:mod:`arbi_serve.realtime.duplex_lane`), where
        ``run_forever``'s own loop — not the duplex code — issues the
        step, so ``step_stt``'s "I own the whole tick" shape does not
        apply (design doc §7.8/§7.9). ``None`` means this model has no
        push-driven path yet; the lane then refuses to register a
        connection for it, loudly.
    finish_stt_tick: ``(session) -> DuplexStepResultLike`` — the SECOND
        half: read the completed step's per-head logits for this
        session's OWN row, run the model's own per-tick bookkeeping
        (feedback tokens, turn-open/turn-close switch), and advance the
        frame counter. Paired with ``prepare_stt_tick``; both or
        neither. Takes the session alone because the retained heads are
        per-REQUEST state reachable from it
        (``Request.retain_logits_heads``/``logits_heads``, design doc
        §7.42) — the lane used to hand this a batch-wide tensor plus a
        slate row index, which made every caller re-derive a row the
        scheduler is free to reorder.
    silence_frame_embed: ``() -> frame_embed`` — one frame of silence,
        in the SAME already-encoded shape ``prepare_stt_tick`` takes.
        The lane feeds this on a tick where the client's audio has not
        arrived yet, rather than skipping the tick (design doc §2.2:
        skipping would desync the frame-index-based turn-taking
        bookkeeping from the shared grid). ``None`` means the lane
        falls back to a zero-valued frame of the same shape as the
        last real one.
    encode_audio_chunk: ``(session, pcm_bytes) ->
        Sequence[(frame_embed, semantic_blank | RAW PCM16 (at
        :attr:`duplex_input_sample_rate`) from a live WS
        ``input_audio_buffer.append`` stream into zero or more
        already-encoded frame embeddings, in the SAME shape
        ``prepare_stt_tick`` takes — the counterpart, for LIVE
        streamed audio, of the precomputed-whole-utterance
        ``binding.encode(mm_feats, device)`` call every existing
        duplex test uses. Each frame is paired with its OPTIONAL
        semantic speech/silence verdict (design doc §7.32 — for
        NemotronVoiceChat, the checkpoint's own RNNT endpoint
        detector's blank/non-blank decode of the same frame);
        ``None`` per frame when the model computes no such verdict.
        Stateful per ``session`` (an incremental encoder needs to
        remember how much audio it has already turned into frames,
        plus its per-layer caches); ``None`` for the whole field
        means this model has no live-streaming audio-in encoder, so
        :class:`~arbi_serve.realtime.duplex_audio_in.DuplexAudioIn`
        (the WS-side caller) fails loud rather than silently dropping
        or mis-encoding audio.
    duplex_input_sample_rate: the PCM sample rate ``encode_audio_chunk``
        expects (NemotronVoiceChat's perception preprocessor is fixed
        at 16 kHz — see
        ``arbi_serve.multimodal.preprocess.nemotron_voicechat
        .SAMPLE_RATE``); the WS-side caller resamples a client's own
        ``input_audio_sample_rate`` down/up to this before calling.
        ``None`` when ``encode_audio_chunk`` is also ``None``.

Resolve ``eng.model``'s :class:`DuplexModelAdapter`, generically.

Raises ``RuntimeError`` (fails loud, mirrors
``forward_exec.resolve_embed_override``'s own
``SUPPORTS_EMBED_OVERRIDE`` gate) if the served model does not declare
``SUPPORTS_DUPLEX_MODE = True``, or declares it but has no registered
factory (a model author's bug — the flag promises this function will
find something). ``adapter_kwargs`` are forwarded verbatim to the
resolved factory (e.g. NemotronVoiceChat's own factory needs
``audio_prompt_latent``/``vocab``/etc. — see
:func:`arbi_serve.realtime.nemotron_voicechat_duplex_adapter
.build_nemotron_voicechat_duplex_adapter`); this function itself
stays agnostic to what any particular model's factory needs.

Everything that happens on one duplex tick AFTER the STT decode step.

This is the per-connection,
model-agnostic "second half" of a duplex tick, factored out of
:func:`~arbi_serve.realtime.duplex_driver.drive_duplex_or_turn_session`'s
loop body so the two DIFFERENT things that can drive a tick share ONE
implementation instead of carrying a copy each:

  * **PULL** — :mod:`arbi_serve.realtime.duplex_driver`, where an
    external caller explicitly drives ``adapter.step_stt`` per tick. This
    is proven correct in isolation (CPU tests + a live-GPU turn-parity
    test) but is NOT the production shape: see §7.8's blocking finding.
  * **PUSH** — :mod:`arbi_serve.realtime.duplex_lane`, where
    ``run_forever``'s own always-on loop issues the STT step (the
    scheduler's duplex-lane gate already builds and runs the slate by
    itself) and the lane's post-step hook calls this pump with the
    result.

Both paths hand this object a :class:`DuplexStepResultLike` and this
tick's raw PCM (if any) and get back the ordered
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TurnEvent` list for
that tick. Nothing here knows which model is being served — every
model-specific action goes through
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter`.

=== What one ``post_stt`` call does, in order ===

1. **VAD x Site-3 barge-in FIRST so
   a forced interrupt is visible to everything downstream in this same
   tick. When it fires, the session's ``agent_idle`` flag flips AND the
   model's own text-channel feedback is force-closed
   (``session.force_eos_feedback``, design doc §7.26) — the flag alone
   would leave the model's recurrent conditioning still following
   whatever it was saying the instant before the interrupt. The EOU-gated
   BOS admission tracker (design doc §7.30/§7.32) is then fed — from the
   frame's own RNNT endpoint-detector verdict
   (``session.note_semantic_frame``) on a connection that carries one,
   else from the SAME tick's VAD events
   (:func:`~arbi_serve.realtime.duplex_barge_in.vad_tick_events` →
   ``session.note_vad_events``) — state that :func:`~arbi_serve.runtime
   .nemotron_voicechat_duplex_step.finish_duplex_tick` reads on a LATER
   tick, mirroring how ``force_eos_feedback`` above feeds a later tick's
   fusion rather than this one's.

   The gate has a SECOND trigger on a semantic-EOU connection (design doc
   §7.57): ``session.semantic_barge_in_fires``, read straight after that
   tick's verdict is folded in. The VAD edge is edge-triggered on
   ``speech_started``, which a user who never stops talking cannot
   produce, so it is blind to a turn opened mid-utterance; the sustained-
   speech half is the reference's own ``RNNT_BARGE_IN_FRAMES`` backstop
   for exactly that case. Both triggers close through one
   ``_close_on_barge_in`` so they cannot drift apart, and each names its
   own cause in the close reason.
2. **``agent_was_active`` / termination bookkeeping** — the bounded
   (turn-shaped) session's stop condition, in two shapes: the ordinary
   one (a full active-then-idle cycle once the real input audio runs
   out) and the SILENT turn (the audio ran out and the agent never once
   became active — ``silent_grace_ticks``, see :meth:`DuplexTickPump
   .post_stt`). A genuine duplex session is unbounded and only ever
   stops at ``max_ticks`` / connection close.
3. **Incremental detokenization** of this tick's text token, into the
   one-slot mailbox the audio-out generator's injected ``stt_step_fn``
   reads back.
4. **One audio-out (TTS) tick** — two ``next()`` calls on the model's own
   audio-out generator, which yields a ``TextDelta`` then an
   ``AudioChunk`` (plus a trailing ``TurnDone`` on an EOS tick). A model
   with no audio-out side (``adapter.create_audio_out_driver is None``)
   gets the same event shape synthesized instead, so callers never branch
   on whether the served model has a TTS half.
5. ``call_id`` and prepended
   to this tick's events as a
   :class:`~arbi_serve.realtime.nemotron_voicechat_turn.FunctionCallReady`,
   ahead of the tick's own deltas — see :meth:`DuplexTickPump
   ._function_call_events`. The engine-side gate that PADs the text
   channel for the duration is already applied by the time this pump
   sees the tick, so nothing here needs to know about it.
6. **Re-assert a forced interrupt**, if one fired this tick — a model's
   audio-out driver may run its own turn-open/turn-close switch off the
   same token and re-open the turn behind the gate's back. See
   :meth:`DuplexTickPump.post_stt`'s own comment; design doc §7.9.8.

The one-slot mailbox is the same bridge the pull driver always used: the
audio-out generator is SYNC and pulls its STT result through an injected
callback, while the STT step is ASYNC (pull) or already-completed (push).
The generator stores the exact callable object handed to it at creation
time, so the only way to change what it sees per tick is to mutate a cell
that same closure reads — hence one dict created once per connection.

Move an

Every tick's ``AudioChunk`` comes straight off
``NemotronAudioCodec.decode`` and is therefore a **device** tensor, and
every consumer of it wants bytes:
:func:`~arbi_serve.realtime.duplex_ws_bridge._waveform_to_pcm16` ends
in ``.cpu().numpy().tobytes()``. Where that device->host copy happens
is not a detail. Both threads use the legacy default stream
(``cuda_stream=0x0``), so a ``.cpu()`` issued from the WS/HTTP loop is
ordered behind everything the engine thread has enqueued since — i.e.
the consumer blocks for up to a whole tick's GPU work, on the loop that
is simultaneously supposed to be receiving mic audio, running the
perception encoder and writing this connection's frames out. Measured
on an idle RTX 4090 with one tick (~70 ms) of engine work in flight:
**62.9 ms** blocked per chunk, against **0.083 ms** for the same call on
an already-host tensor.

Doing it HERE costs the engine thread a sync on work it just issued and
needs finished anyway (the codec decode, ~3 ms, already inside the
frame budget) plus a 3.5 KB copy, and leaves the consumer with pure CPU
work. The turn-based path reached the same conclusion by a different
route — ``stream_nemotron_voicechat_turn`` wraps the identical call in
``asyncio.to_thread`` rather than run it on its loop — which is the
off-loop-or-bust discipline the duplex bridge was missing.

Any non-``AudioChunk`` event, and any chunk already on the host, is
returned unchanged, so the pull driver and every model without a real
device tensor pay one ``isinstance`` per event.

One connection's post-STT tick machinery — see module docstring.

Args:
    adapter: this connection's resolved
        :class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter`.
    model: the served model, handed verbatim to
        ``adapter.create_audio_out_driver``.
    tokenizer: used only for ``incremental_decode`` of the text
        channel.
    session: this connection's
        :class:`~arbi_serve.realtime.duplex_model_adapter.DuplexSessionLike`.
    max_ticks: safety valve, and the ``max_tokens`` handed to the
        audio-out generator so its own step counter reaches the same
        bound at the same tick this pump does.
    bounded: ``True`` for a turn-shaped session that must terminate
        on its own (one full active-then-idle cycle after the real
        input audio runs out); ``False`` for a genuine duplex
        connection, which never ends on its own.
    num_real_audio_frames: where the caller's real input audio ends
        and silence padding begins — read by the ``bounded``
        termination check only.
    vad: optional :class:`~arbi_serve.realtime.vad.VadDetector` for
        the barge-in AND-gate. ``None`` leaves the gate fully inert.
    semantic_eou: ``True`` routes the EOU admission tracker to the
        per-frame semantic verdicts riding this connection's frames
        (``session.note_semantic_frame`` — the served checkpoint's
        own RNNT endpoint detector, design doc §7.32; a tick without
        a verdict counts as blank/silence, which is what a
        no-client-audio tick is). ``turn_detection.backend``, and every
        model without the bundled RNNT head. The barge-in AND-gate
        reads the acoustic VAD either way.
    silent_grace_ticks: how many ticks a BOUNDED session keeps
        stepping past the end of its real input audio while the agent
        has never once become active, before the turn is declared
        over with no response. ``None`` (the default) disables the
        check entirely, which is bit-identical to the behaviour
        before it existed. See :meth:`post_stt`'s "the silent turn".
        Ignored for an unbounded (genuine duplex) connection, which
        has no turn to end.

Mirror this BOUNDED turn's ``num_real_audio_frames`` onto the
session, where the EOU-gated BOS admission check can read it
(design doc §7.35).

A bounded pump's turn boundaries are owned by a client, so its
admission rule is "this turn's committed audio has been consumed"
rather than the VAD's confirmed-speech-then-silence — see
:meth:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession.eou_admits_bos`. Published from
HERE, off the same assignment that sets
:attr:`num_real_audio_frames`, so the pump's own
``no_more_real_audio`` check and the admission gate can never
disagree about where this turn's audio ends.

A genuine (unbounded) duplex pump publishes nothing, leaving the
session's target ``None`` and its admission rule untouched.
``getattr``: a minimal session stand-in in a CPU test predates
this hook and simply keeps the pre-existing rule.

Pull exactly one tick's worth of events out of the audio-out
generator: a ``TextDelta`` then an ``AudioChunk``, plus a trailing
``TurnDone`` when ``expect_eos``. Always exactly two (or three)
``next()`` calls per tick, no more, no less — including the very
first tick, which also runs the generator's one-time priming code
before its first yield.

Run everything after this tick's STT decode step; return this
tick's events in order.

``pcm_frame``: this tick's raw PCM16 bytes for the barge-in VAD
(the SAME frame the STT side consumed as an already-encoded
embedding), or ``None`` when the caller has no PCM for this tick
(e.g. a bounded turn's padded silence tail) — the gate then
simply does not run this tick.

``semantic_blank``: this tick's RNNT endpoint-detector verdict
(design doc §7.32), riding the frame the STT side consumed —
``True`` = blank/no transcribable speech, ``False`` = the head
recognized real content, ``None`` = the tick carried no verdict
(a synthesized silence frame, or a frame pushed without one).
Read only when :attr:`semantic_eou` is set; a ``None`` then
counts as blank.

**The silent turn.** A bounded (turn-shaped) session's ordinary
stop condition requires :attr:`agent_was_active` — one full
active-then-idle cycle. A turn where the model simply never
speaks therefore has no stop condition at all and runs to
:attr:`max_ticks`, which for a repointed ``response.create`` is
the difference between "one prompt response completion" and
"nothing for 400 s" (design doc §7.9.15's protocol-semantics gap
(a)). :attr:`silent_grace_ticks` closes it: once the real input
audio is exhausted, a bounded turn that has still never seen the
agent active gets that many more ticks to start, and then ends —
as a genuine, ``status="completed"`` empty turn, not an error and
not a timeout. The grace is counted rather than assumed because
the model legitimately answers a few frames after the audio ends;
``0`` means "end immediately", ``None`` disables the check.

This tick's detected tool calls, as
:class:`~arbi_serve.realtime.nemotron_voicechat_turn.FunctionCallReady`
events with freshly minted

``getattr``, not a bare read: ``function_calls`` is an OPTIONAL
part of the tick-result contract
(:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexStepResultLike`
does not require it), so a model whose duplex path has no function
channel — and every existing test double — stays on the untouched
path and pays one attribute lookup.

``call_id`` is minted HERE rather than on the engine thread for the
same reason
:func:`~arbi_serve.realtime.nemotron_voicechat_turn.stream_nemotron_voicechat_turn`
mints it at its own WS-facing layer: it is a protocol identifier the
client will echo back on ``conversation.item.create``, not model
state, and the same ``call_{uuid4}`` shape keeps the two paths'
wire output indistinguishable to a client.

Emitted BEFORE this tick's ``TextDelta``/``AudioChunk``: the client
must be able to start executing the tool before it sees the
(now-PADded) deltas of the very tick that triggered it.

Re-arm this pump for another bounded turn on the SAME
connection/session — the duplex analogue of building a fresh
``EngineSttStep`` for turn N+1, without discarding this
connection's persistent KV/Mamba state (design doc §7.24's turn
re-arm primitive).

Resets every field :meth:`post_stt` uses to decide when THIS turn
ends — :attr:`agent_was_active`, :attr:`silent_tail_ticks`,
:attr:`silent_turn`, :attr:`finished` — plus the transcript
accumulator (``_generated_token_ids``/``_decode_state``), so the
new turn's own ``TextDelta`` stream starts clean instead of
continuing the previous turn's. Deliberately does NOT touch
:attr:`tick_count`, and leaves :attr:`max_ticks` alone unless the
caller passes one: both remain the CONNECTION's own lifetime
safety valve (design doc §3.5), which keeps bounding a repointed
multi-turn session exactly as it already bounds a genuine duplex
one, across every turn it re-arms through.

``num_real_audio_frames`` is an ABSOLUTE ``session.frame_seq``
target — the value at which THIS turn's real input audio ends —
not a turn-relative count: :attr:`session`'s own frame counter is
connection-lifetime and never resets (see
``NemotronVoiceChatDuplexSession``'s own docstring on why), and
:meth:`post_stt`'s ``no_more_real_audio`` check already reads it
that way.

Closes and drops the audio-out generator (:attr:`_gen`) rather
than reusing it: the generator (e.g. ``run_turn``) is turn-scoped
in its own step counter and yields exactly one terminal event, so
a second turn needs a second generator instance. Sound because
the actual TTS feedback (``prev_tts_code``/``tts_past_kv``) lives
on :attr:`session`, not on the generator's own locals — the new
instance primes from it on its own first tick, so nothing about
the model's conditioning is lost across the swap.

Only valid for a ``bounded`` pump: a genuine duplex connection has
no turn boundary to re-arm at all.

Terminal events for a pump that ran out of ``max_ticks``
without an internal EOS.

With an audio-out side, asks the generator itself for its own
``TurnDone`` rather than fabricating one, so ``num_steps`` always
comes from the SAME source. ``max_tokens=max_ticks`` was passed to
it at creation, so its own step counter reaches the SAME bound at
the SAME tick this pump does — the ``next()`` below observes its
own ``while`` condition already false and gets straight to
``TurnDone``, never re-entering the loop body / calling
``stt_step_fn`` again. Note the generator's own ``finish_reason``
string is whatever that model's own generator hardcodes (e.g.
NemotronVoiceChat's ``run_turn`` always yields ``"max_tokens"``,
matching its turn-based-mode heritage) — this pump does not
rewrite it, keeping the audio-out seam untouched; callers that
need to distinguish a genuine duplex

Without an audio-out side this pump synthesizes the event itself,
so ``bounded`` decides the reason directly: ``"max_tokens"`` for a
bounded (turn-shaped) session's own safety valve, unchanged;
``"max_duration_reached"`` for a genuine (unbounded) duplex
connection, whose ONLY way to stop on ``max_ticks`` at all is this
cap (design doc §3.5) — it never finishes on its own (§7 step 2's
``ignore_eos=True``/``max_tokens=None``).

THIRD tier / §6
decision point 4): the consumer side that turns a live
:class:`~arbi_serve.realtime.duplex_lane.DuplexConnection`'s outbound
``TurnEvent`` queue into real ``/v1/realtime`` wire events sent to a WS
client.

=== What already existed before this module, and what was missing ===

PRODUCER half all the way to
a queryable per-connection queue: ``run_forever``'s own hooks
(:mod:`arbi_serve.realtime.duplex_lane`) already turn every duplex tick
into real :class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta`/
``AudioChunk``/``TurnDone`` objects and publish them onto
``DuplexConnection.outbound`` with one consumer wakeup per tick
(``DuplexConnection.drain_outbound``/``.wakeup``). Nothing before this
module ever called ``drain_outbound`` outside a test — there was no
consumer turning those objects into wire-format JSON and no synthesis of
the ``agent_idle``

=== The wire-framing choice (design doc §6 decision point 4(a)) ===

Duplex mode has no discrete ``response.create``/``response.done`` request
shape — the assistant can start speaking without being asked, and a
genuinely unbounded connection has no terminal point for one giant
response object. Two options were on the table; this module implements
the RECOMMENDED one: reuse the existing OpenAI-Realtime ``response.*``
wire vocabulary (:mod:`arbi_serve.realtime.events`'s builders, UNCHANGED)
with a SYNTHETIC response lifecycle keyed off ``agent_idle``:

  * ``agent_idle`` observed ``True -> False`` (Site-3 BOS) synthesizes a
    fresh ``response.created`` / ``response.output_item.added`` /
    ``response.content_part.added`` — exactly
    :meth:`~arbi_serve.realtime.session.RealtimeSession._run_response`'s
    own opening sequence, new ids minted the same way (``ev.new_id``).
  * ``agent_idle`` observed ``False -> True`` (Site-3 EOS) synthesizes the
    matching close (``response.audio.done`` / ``...audio_transcript.done``
    / ``...content_part.done`` / ``...output_item.done`` /
    ``response.done``) — exactly
    :meth:`~arbi_serve.realtime.session.RealtimeSession._finalize_response`'s
    own closing sequence.
  * While no response is open (``agent_idle`` still ``True``), STT/TTS
    ``TurnEvent``s are NOT translated onto the wire at all — there is no
    ``response_id`` for a client to attach them to under this framing, and
    the alternative (always-on session-scoped events with no
    ``response_id``) was the option NOT taken. A duplex connection that
    never leaves ``agent_idle=True`` (the model never spoke) therefore
    produces zero response-lifecycle traffic, which is correct: nothing
    happened worth reporting as a response.

This means an existing OpenAI-Realtime SDK client speaks to a duplex
session exactly as it would a turn-based one — many short-lived responses
instead of one long one — with NO new wire-protocol surface
(:mod:`arbi_serve.realtime.events`/``schemas.py`` are untouched by this
module).

=== ``max_ticks`` safety-valve close ===

A genuine duplex connection (``bounded=False``) never produces
``TurnDone`` via the model's own EOS — that path requires
``agent_idle`` already ``True``, which the transition synthesis above
already closes on its own. The one case that can still leave a response
open going into a ``TurnDone`` is the ``max_ticks`` safety valve
(:meth:`~arbi_serve.realtime.duplex_tick_pump.DuplexTickPump.finish`),
which fires independent of ``agent_idle`` — this module treats that as a
forced ``status="cancelled"`` close, mirroring
``RealtimeSession._finalize_response``'s own cancelled-response shape.

=== Threading / cadence ===

This module is entirely WS-side (the HTTP loop), never engine-thread
code — it contains no ``call_soon_threadsafe`` and imports nothing engine-
internal, so it is exempt from (and does not need to satisfy)
``tests/test_transport_boundary.py``'s engine<->HTTP crossing invariant;
the ONE crossing per tick already happened inside
``DuplexConnection._publish``/``LoopBridge.notify_http`` before any event
reaches here. :func:`pump_duplex_connection` drains once per wakeup — in
steady state that is once per completed tick
(``batch_cfg.duplex_frame_interval_s``, 80 ms default), matching the
level-per-tick model the transition synthesis above assumes; see
:func:`translate_duplex_tick`'s own docstring for the documented,
bounded-under-backpressure imprecision if the consumer ever falls behind.

Owns the synthetic ``response.*`` lifecycle for ONE duplex
connection — the ``agent_idle``-keyed half of design doc §6 decision
point 4(a). One instance per connection, held by
:func:`pump_duplex_connection`'s own loop (or a caller driving the
translation manually, e.g. for tests).

Every method is a PURE, synchronous ``-> list[dict]`` producer (no
``send``/awaiting inside this class) so the translation logic is
testable without an event loop; :func:`pump_duplex_connection` is the
thin async layer that awaits ``send`` for each returned event, in
order.

``[B, T]`` (or ``[T]``) float waveform in ``[-1, 1]`` -> raw
little-endian int16 PCM bytes. Deliberate small duplicate of
``nemotron_voicechat_turn.py``'s own private helper of the same name
(rather than importing across module boundaries) -- both are ~5 lines
and this module has no other reason to depend on that one; keeps this
module's own import list minimal (torch stays a function-local import,
not a module-top one -- see that module's own "Torch-free at import"
precedent, ``server/routes/openai_realtime.py``'s docstring).

Pure translation of one drained batch of ``TurnEvent``s + the
``agent_idle`` level observed AFTER that batch into wire events.

Synthesizes ``response.created``/``response.done`` off a ``True<
->False`` level flip (design doc §6 decision point 4(a); see module
docstring), then translates every ``TextDelta``/``AudioChunk`` in
``events`` through ``bridge`` (in order), then closes on a
``False->True`` flip, then forces a close if any ``TurnDone`` in
``events`` fired the ``max_ticks`` safety valve while a response is
(still) open.

``ensure_response`` — **exactly-one-response semantics for a BOUNDED
(turn-shaped) session** (design doc §7.9.15's protocol gap (a)). The
``agent_idle``-keyed framing above reports what the model DID: a turn
where it never spoke opens no response at all, which is right for a
genuine duplex connection (nothing happened worth reporting as a
response) and wrong for a turn a client explicitly asked for, which
owes exactly one ``response.*`` lifecycle whatever the model chose to
do. With this set, a terminal ``TurnDone`` that finds NO response
open synthesizes a complete, EMPTY one (``response.created`` …
``response.done``) instead of nothing. ``False`` (the default, and
every genuine duplex connection) leaves this function bit-identical
to before the flag existed.

**Known limitation, deliberately not silently resolved**: ``events``
may span MORE than one tick if the caller fell behind between drains
(``DuplexConnection.outbound`` is a bounded, coalescing queue by
design -- see that class's own docstring). This function reads
``agent_idle`` ONCE, at the level observed after every event in
``events``, so a flip-then-flip-back that completes ENTIRELY within
one coalesced batch is not observed as two separate responses -- it
is invisible on the wire rather than reported. In steady state (one
drain per completed tick, the common case: ``DuplexConnection
._publish`` wakes the consumer every tick) this is exact, not an
approximation. This mirrors ``DuplexConnection.outbound``'s own
documented lossy-under-backpressure discipline (drops the OLDEST
event, not the newest) rather than introducing a new kind of loss.

Drain ``conn``'s outbound queue until it closes, translating every
tick's events onto the wire via Runs on the WS
handler's own loop (never the engine thread); intended to be driven as
a background ``asyncio.Task`` alongside the caller's own client-event
read loop, exactly the shape ``RealtimeSession._start_response``
already uses for its (turn-based) response task.

``conn`` must have been built with a real ``consumer_loop``/``wakeup``
(``attach_duplex_connection(..., consumer_loop=asyncio.get_running_loop(),
wakeup=asyncio.Event())``) — this is the WS-side half of that contract.

``on_function_calls``: called with every
:class:`~arbi_serve.realtime.nemotron_voicechat_turn.FunctionCallReady`
in ONE drained batch, as a list, Both properties are
load-bearing:

  * **as a batch**, because one ``<TOOLCALL>[...]</TOOLCALL>`` block
    can name several calls and the model expects ONE
    ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>`` answering all of them —
    the same grouping ``stream_nemotron_voicechat_turn`` does;
  * **before sending**, so the handler has registered its pending
    futures by the time the client can possibly see the call and
    answer it. A client that replies inside one event-loop iteration
    would otherwise hit "no pending function call with call_id".

It must not block (it runs inline on this loop); the reference
implementation, :meth:`~arbi_serve.realtime.session.RealtimeSession
._on_duplex_function_calls`, registers the futures synchronously and
spawns the waiting/resume work as a task.

A raise anywhere in one iteration's body — the drain, the
``on_function_calls`` handover, the translation, a ``send`` — is
logged (throttled, ``_FAILURE_LOG_EVERY``) and the loop continues with
the next tick, rather than ending the task. Ending it is what used to
happen, and it is the worst available outcome: the engine goes on
ticking, the bounded outbound deque goes on dropping its oldest
events, the client hears nothing, and an ``asyncio.Task`` that dies
unobserved says nothing until GC. ``was_agent_idle`` is only advanced
on a tick that translated successfully, so a failed tick does not
consume the BOS/EOS level flip it never emitted.

Self-terminates once ``conn.closed`` is observed true, closing any
still-open synthetic response first (mirrors
``RealtimeSession._finalize_response``'s cancelled/failed shapes, keyed
off ``conn.finished_reason``). A connection closed EXTERNALLY between
ticks (``close_duplex_connection`` -> ``DuplexLane.unregister`` sets
``conn.closed`` but does not itself wake this loop — only a completed
or failed TICK does) is still noticed within one poll interval (module
docstring's ``_POLL_INTERVAL_S``) rather than hanging forever; a caller
that wants tighter, zero-latency teardown should also cancel this
coroutine's task directly when it tears the connection down (the same
pattern ``RealtimeSession.aclose`` uses for its own response task).

The wire->client stage, rebuilt if the client changed its
requested ``output_audio_sample_rate`` since the last delta.

A rate change necessarily discards the resampler's carried tail —
it is a different polyphase kernel — but that is the truthful
cost of the client asking for a different rate mid-stream, it
happens only on an explicit ``session.update``, and the
alternative (freezing the rate at connection-open) silently
ignores the request. The codec->wire stage above is unaffected:
both of its rates are model constants.

BOS: mint a new ``response_id``/``item_id`` and emit the
opening sequence — see :meth:`~arbi_serve.realtime.session
.RealtimeSession._run_response`'s own opening block, which this
mirrors exactly (id minting via ``ev.new_id``, same three
events, same order).

Idempotent: a no-op ``[]`` when a response is already open. ``agent_idle`` has to open a response to carry its
``response.function_call_arguments.done`` (see
:func:`translate_duplex_tick`), and the agent then very often
starts speaking a tick or two later, which is a second, ordinary
BOS open of the SAME response. Reporting that as an error (the
previous ``assert``) would kill the consumer loop over a
legitimate sequencing; returning ``[]`` folds the call and the
speech it triggered into one response, which is also what a
client wants to see.

One ``TextDelta``/``AudioChunk`` -> its wire event(s) (beta +
GA-vocabulary sibling, matching ``_run_response``'s own dual-emit).
A caller must not call this while :attr:`response_active` is
``False`` — :func:`translate_duplex_tick` enforces that; direct
callers get an empty list plus a debug log rather than a raise,
since a TurnEvent arriving on an idle tick is a real (if
unexpected) possibility this bridge should degrade on, not crash
on (see design doc §7.9.8's own "the lane may lose a connection;
it must never lose more" discipline, extended here to the WS
side: a wiring bug here must not kill the whole consumer loop).

EOS (or a forced close): emit the closing sequence — mirrors
:meth:`~arbi_serve.realtime.session.RealtimeSession._finalize_response`'s
own closing block. A no-op (``[]``) if no response is open, so
callers can call this unconditionally at connection teardown.

OpenAI Realtime server-event builders.

Every ``server -> client`` event the ``/v1/realtime`` WebSocket emits is
constructed here so the wire shape lives in one place and the session state
machine stays readable. Each builder constructs the typed model from
``schemas.py`` and returns ``.model_dump()`` — a plain ``dict`` ready for
``websocket.send_json``. Because the emit path and the published OpenAPI
schema share those models, the two cannot drift.

Ids are minted with the OpenAI ``event_`` / ``item_`` / ``resp_`` / ``msg_``
prefixes. Only the subset of the protocol this server implements is modelled;
the shapes track the OpenAI Realtime beta so unmodified OpenAI-Realtime SDK
clients parse them. Fields the model cannot populate (e.g. a user-audio
transcript when input transcription is disabled) are emitted as ``null``.

NemotronLabs-VoiceChat-11B's concrete
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter` — the
FIRST implementation of the generic duplex-mode protocol (see that
module's own docstring for why the protocol exists and why this is not
assumed to be the only implementation ever).

Everything model-specific the generic driver
(:mod:`arbi_serve.realtime.duplex_driver`) would otherwise need to know
about NemotronVoiceChat lives HERE and only here: the real per-tick STT
stepper (:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.step_duplex_tick`,
reused unmodified), this checkpoint's own turn-open/turn-close token ids
(:func:`~arbi_serve.runtime.nemotron_voicechat_stt_step.resolve_turn_bos_id`/
:func:`~arbi_serve.runtime.nemotron_voicechat_stt_step.resolve_turn_eos_ids`),
its own session class
(:class:`~arbi_serve.realtime.nemotron_voicechat_duplex_session.NemotronVoiceChatDuplexSession`),
and its own audio-out driver (:func:`~arbi_serve.runtime.nemotron_voicechat_turn.run_turn`,
reused unmodified, bound to this connection's own voice-conditioning/vocab
kwargs and used as the "audio-out generator" half of
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexModelAdapter`).

=== ``encode_audio_chunk``: live

WS
``input_audio_buffer.append`` PCM stream is now turned into duplex-lane
frame embeddings incrementally, one per 80 ms, with no whole-utterance
re-encode anywhere.

``binding.encode`` over a trailing raw-PCM window on
every chunk — and was RIGHT to: that window's analytically safe size
compounds across all 24 Conformer layers into tens of seconds of audio
re-encoded per tick. What it named as the real requirement ("a genuine
incremental/cached-state streaming Conformer forward, per-layer
conv/attention state carried tick to tick, INTENDED mode
rather than new capability grafted on: this checkpoint's perception
tower is NVIDIA's cache-aware streaming Fast Conformer
(``nvidia/nemotron-speech-streaming-en-0.6b``) pinned to its 80 ms,
zero-lookahead ``att_context_size=[70, 0]`` streaming mode — the same
80 ms the duplex lane's own ``duplex_frame_interval_s`` tick already
runs at, one encoder frame per tick exactly. See
:mod:`arbi_serve.models.audio.nemotron_voicechat_perception`'s module
docstring for the architecture, the cache shapes, and the measured
agreement with the offline path.

Two pieces of per-connection state make this work, both created lazily
on the first chunk and living on the session (so a connection's audio
history is never shared or replayed across connections):

  * :class:`~arbi_serve.multimodal.preprocess.nemotron_voicechat.MelStreamer`
    — raw PCM16 -> log-mel frames, chunk-boundary invariant, so the WS
    client's arbitrary ``input_audio_buffer.append`` sizes produce
    exactly the frames a single buffered call would have.
  * :class:`~arbi_serve.models.audio.nemotron_voicechat_perception
    .PerceptionStreamState` — the encoder's per-layer attention/conv
    caches.

Chunk sizes need not line up with the 80 ms grid in either direction: a
chunk smaller than one frame's worth of audio returns no frames (the
lane's ``silence_frame_embed`` covers that tick), and a burst carrying
several frames' worth is encoded in ONE forward rather than one per
frame.

The checkpoint's own 1024-entry ``rnnt_tokenizer/vocab.json``
SentencePiece piece list, for rendering the RNNT endpoint detector's
emitted label ids as text (design doc §7.48) — ``None`` when the
served directory does not bundle it (e.g. a quantized export that
only copied the model tensors), in which case the heard-window log
falls back to raw ids, which is still a usable diagnostic.

Build this connection's NemotronVoiceChat duplex adapter.

``audio_prompt_latent``/``vocab``/``subword_id_to_char_ids``/TTS
sampling knobs are this connection's own values (mirrors
:func:`~arbi_serve.runtime.nemotron_voicechat_turn.run_turn`'s own
parameter list — see that function's docstring for what each one
means) — bound here, once, into the returned adapter's
``create_audio_out_driver`` closure, so the generic driver loop never
needs to know these are NemotronVoiceChat-specific TTS parameters at
all.

One function-channel token's visible text — the SAME lookup
``EngineSttStep.decode_function_token`` performs for the
turn-based path's own detector (``tokenizer.decode([id],
skip_special_tokens=True)``), minus the re-sampling step: the
duplex tick has already committed this channel's greedy token, so
re-running a ``Sampler`` over the same logits here would only risk
disagreeing with the token that actually entered the recurrence.

Live WS PCM16 (16 kHz mono) -> zero or more duplex frames.

Returns ``(frame_embed, semantic_blank)`` pairs — each frame in
the SAME ``(1, hidden)`` shape ``prepare_stt_tick``/
``silence_frame_embed`` use, so the lane cannot tell a
live-encoded frame from a pre-encoded one, paired with the
frame's RNNT endpoint-detector verdict (design doc §7.32): the
bundled ``rnnt_head`` greedily decodes each frame's
PRE-projection encoder output (the reference's ``asr_emb``, a
second output of the SAME ``encode_stream`` forward) and a frame
is blank exactly when the checkpoint's own ASR head finds no
transcribable speech in it. ``None`` verdicts (head absent — a
model built without audio) leave the EOU tracker on its acoustic
fallback.

Admission path for a persistent NemotronVoiceChat duplex-lane ``Request``.

ONE
scheduler-owned, persistent ``Request`` a duplex WS connection lives on
for its whole life (the ``Request`` the session will eventually hold), and admits it
into :class:`~arbi_serve.scheduler.scheduler.Scheduler` via the scheduler's
own, unmodified admission entrypoint.

Explicitly OUT of scope here (later steps, per the design doc's own
ordering):

* THIS module, ``Scheduler`` does not branch on
  ``Request.is_duplex_frame``/``SamplingParams.priority == "duplex"``
  ANYWHERE. A request admitted here is scheduled exactly like an
  ordinary ``"interactive"`` decode row — see
  ``tests/test_nemotron_voicechat_duplex_admission.py``'s
  ``test_duplex_request_is_scheduled_as_ordinary_decode_today`` for the
  proof and its implications.
* Wiring ``fuse_stt_step_embeds`` -> * This module's
  :func:`new_duplex_request` takes the (already-tokenized) seed prompt
  ids as a plain argument instead, so it stays engine/model-free.

Why the request needs a non-empty ``prompt_token_ids`` seed
-------------------------------------------------------------
A literally EMPTY ``prompt_token_ids`` combined with an empty
``output_token_ids`` and ``state=DECODING`` is a real, already-documented
landmine: ``tests/test_chunk_prefill_admission.py::
test_build_batch_index_error_on_emptied_prompt_regression`` pins that
``ModelRunner._build_batch``'s decode-row tail-token lookup
(``req.output_token_ids[-1] if req.output_token_ids else
req.prompt_token_ids[-1]``) raises ``IndexError`` on exactly that shape —
and that test's own docstring notes the scheduler "guarantees a request
whose slate row says ``is_prefill=False`` carries either a non-empty
prompt or a non-empty output", i.e. this shape is supposed to be
UNREACHABLE via normal admission. A duplex request built with a truly
empty prompt (the literal reading of "no prefill-then-decode transition")
would defeat that guarantee the moment it's real-forwarded even though
nothing in the CPU-only scheduler bookkeeping would catch it. An empty
prompt is also a zero-width prefill row, which the slate builder cannot
size. :func:`new_duplex_request` refuses it outright.

Why the seed prompt This module used to build the request with ``prompt_consumed ==
len(prompt_token_ids)``, so ``is_prefill`` was ``False`` from birth and
no prefill row was ever built for the seed. The justification given was
that this reused ``Scheduler.schedule``'s "full cached match"
zero-remaining-prompt fast path. That path is correct ONLY because a
radix full-hit means the prompt's KV ALREADY EXISTS in the cache,
computed by an earlier request. A duplex request is built with
``cache_enabled = False`` and has no cached prefix, so the premise never
held: the seed's KV was simply absent. The first tick took the decode
branch and fed ONLY the last seed token, anchored at ``total_length - 1``
— measured on a 5-token seed against the real ``FlatPageTable`` +
``_gather_slate``, the row's ids were ``[50]``, its position ``4`` and
``seq_lens`` ``1``, i.e. the kernel attended over ONE KV entry while the
position claimed index 4. Nothing ever wrote KV for seed positions 0-3,
so a live session's system prompt had no effect on generation beyond its
final token and a RoPE offset.

The seed is therefore left UNCONSUMED (``prompt_consumed == 0``) and
prefilled for real, exactly like any other request's prompt: the duplex
lane's slate builder (``Scheduler._build_duplex_slate``) emits chunked
prefill rows until the seed is consumed, then the request settles into
its steady-state one-row-per-frame decode cadence. See design doc
§7.9.13.

The refusal message for a duplex admission, or ``None`` to allow it.

The single predicate behind every duplex refusal site, so the lane
(:func:`~arbi_serve.realtime.duplex_lane.attach_duplex_connection`)
and the scheduler admission (:func:`admit_duplex_request`) can never
answer differently for the same engine.

Reads the world size two ways because the two call sites know
different things. ``eng`` (when the caller has one) goes through
:func:`~arbi_serve.engine.run_step.engine_is_multi_rank`, which also
catches an engine carrying a worker bridge. The engine-free form falls
back to the process-wide :class:`ParallelConfig`, which every rank of
a real launch shares — so the verdict is RANK-SYMMETRIC: two ranks
asked the same question return the same answer, and a refusal can
never split the group by firing on one rank only. Pure attribute
reads; no collectives, no CUDA.

Raise when duplex is not admissible on this engine's rank layout.

``site`` prefixes the message with the entry point the caller refused
at, so the traceback names which duplex handoff was attempted.

Render + tokenize + EMBED a duplex connection's one-time seed
prompt exactly the way NVIDIA's reference serves this checkpoint's
system prompt (design doc §7.31).

Three reference behaviors, all verified against the extracted

1. **Token stream** — ``[text_bos] + encode(prompt) + [text_eos]``,
   repeated :data:`SEED_PROMPT_REPEAT_N` times (matching NeMo's
   ``collate_system_prompt`` bracketing and the **Ack stripping** — the template's ``<TOOL_ACK_MESSAGES>`` block is
   removed before tokenizing (it is server-side data, see
   :data:`_TOOL_ACK_MESSAGES_RE`).
3. **Channel-fused embeddings** — every prompt position enters the
   model as the SAME per-channel weighted sum a steady-state frame
   tick uses (``w_user*embed(tok) + w_text*embed(pad) +
   w_func*embed(pad)``; the reference's
   ``_prepare_system_prompt_embeddings`` / the NeMo wrapper's
   fusion-module call, with the BOS position's extra addend being
   ``embed(pad)`` too — ``_get_bos_embedding`` embeds ``pad_id`` — so
   the sum is uniform across positions). ``embed(tok)`` seed the model used to receive was nearly
   orthogonal (cos ≈ 0.30) to its trained prompt-input distribution.

Returns ``(prompt_token_ids, seed_embed_override)`` — hand both to
:func:`new_duplex_request`. The override's rows are consumed
front-to-back across the seed's chunked-prefill steps by
``forward_exec.resolve_embed_override`` (see
``Request.pending_seed_embed_override``'s docstring).

Build a persistent duplex-lane ``Request``, not yet admitted.

``prompt_token_ids`` — the connection's one-time seed context
(design doc §7 step 2: "no prompt to prefill beyond the
system-prompt prefix"). MUST be non-empty (see the module docstring
for why); in production this is the tokenized system-prompt-prefix
rendered once at connection open. It is left UNCONSUMED
(``prompt_consumed == 0``), so ``is_prefill`` is ``True`` and the
seed gets a REAL prefill — its The prefill happens on duplex ticks, served by
``Scheduler._build_duplex_slate``, so it never shares a step with
ordinary traffic.

``sampling`` — an optional caller-supplied base
:class:`SamplingParams` (e.g. to set ``temperature``/``seed``/etc.
for the duplex session's sampling). The duplex-lane invariants below
are forced onto it regardless of what the caller passed (mirrors how
the submission layer force-rewrites ``mtp_k`` at admission — these
are not caller-tunable knobs for this request class):

* ``priority = "duplex"`` — marks it as duplex-lane traffic
  (``SamplingParams.is_duplex_frame`` / ``Request.is_duplex_frame``,
  the ``is_batch_priority``-style field this task introduces;
  design doc §6 decision point 3 / §7 step 3).
* ``ignore_eos = True``, ``max_tokens = None`` — the request never
  finishes on a stop token or a length cap (design doc §0's "no state
  meaning alive indefinitely" gap, §7 step 2). KV/state growth is
  governed by the frame-cap policy (design doc §3.5), not by
  sampling limits — enforcing that cap is a later step (§7 step 8),
  out of scope here.
* ``cache_enabled = False`` — a duplex session's seed context is
  per-connection and has nothing to share with (or usefully match
  against) other requests' prompts; opting out keeps it out of the
  shared radix tree entirely (``SamplingParams.cache_enabled``
  docstring, ``request.py``).

``SamplingParams.duplex_frame_interval_s`` is deliberately NOT forced
here — unlike the invariants above it IS a caller-tunable knob for
this request class. Left unset (``None``) the session paces on the
engine-wide ``batch_cfg.duplex_frame_interval_s`` 80 ms grid, which is
what a live real-time conversation wants. A caller that knows its
session is bounded and has no real-time constraint passes ``0.0`` to
run it unpaced, at engine speed. See design doc §7.9.14, and note the
exclusivity consequence recorded there: an unpaced session holds the
duplex lane against ordinary traffic for as long as it runs.

Returns a ``Request`` in ``state=WAITING`` — see
:func:`admit_duplex_request` for why that is the correct,
reuse-the-existing-admission-path way in.

Admit ``req`` into ``scheduler`` via the scheduler's OWN normal
admission entrypoint (:meth:`Scheduler.add`) — reused unmodified, not
hand-rolled (per this task's explicit scope).

``Scheduler.add``/``_add_inner`` hard-requires ``req.state ==
WAITING`` (``scheduler_admission.py``: "not WAITING (got ...)" raise),
which is exactly what :func:`new_duplex_request` builds. So: call
``scheduler.add(req)`` here, identical to how every other request is
admitted.

From there the request is served ONLY by the duplex lane
(``Scheduler._build_duplex_slate``), never by the ordinary admission
loop — it is invisible to ``_next_waiting_candidate``. Its first
duplex tick(s) carry chunked PREFILL rows for the seed prompt
(``WAITING -> PREFILLING``), and the tick that consumes the last seed
token promotes it to ``DECODING`` (``commit_state``) for the
steady-state one-row-per-audio-frame cadence. The duplex lane holds
the connection's audio unarmed for the whole prefill phase, so no
frame is consumed by a tick that carries no frame — see
``DuplexLane._prepare_tick_inner``.

Raises ``ValueError`` (fails loud, mirrors ``Scheduler.add``'s own
contract) if ``req`` was not built by :func:`new_duplex_request` (or
an equivalent caller that set ``priority="duplex"``/
``state=WAITING``) — this helper is specifically for the duplex
admission shape, not a generic ``scheduler.add`` wrapper.

Raises ``RuntimeError`` on a MULTI-RANK engine
(:func:`duplex_multi_rank_refusal` — see that function for the parts
of a duplex tick that have no rank-symmetric form). This is THE
admission boundary for a duplex row: every duplex driver reaches the
scheduler through here, so a refusal here holds for the WS lane, the
pull driver, and any offline harness alike — the lane's own
``attach_duplex_connection`` refusal fires earlier on the WS path but
covers only that path.

Connection-lifetime duplex session state for NemotronLabs-VoiceChat-11B.

DATA
SHAPE half of "promote ``_FrameLockstepState`` +
``prev_tts_code``/``tts_past_kv`` to a connection-lifetime object" — no
engine wiring, no scheduler changes, no live-GPU work. Admission (building
the persistent ``Request``), the duplex-lane scheduling gate, and the
``fuse_stt_step_embeds`` → ``pending_embed_override`` wiring are later,
separate steps (design doc §7 steps 2-4) and are explicitly OUT of scope
here.

:class:`NemotronVoiceChatDuplexSession` holds the SAME kind of per-step
bookkeeping two turn-scoped pieces of code hold today, but shaped to live
for a whole WS connection instead of being rebuilt/discarded every turn:

  * The STT-side frame-lockstep feedback —
    :class:`~arbi_serve.runtime.nemotron_voicechat_stt_step._FrameLockstepState`
    (``nemotron_voicechat_stt_step.py:181-240``): ``prev_text_token_id``/
    ``prev_function_token_id``, seeded with the tokenizer's text-PAD/blank
    id and updated via :meth:`observe` after each step's own sampling so
    the NEXT step's ``fuse_stt_step_embeds`` call conditions on what the
    model itself just emitted. :meth:`consume_frame` mirrors that class's
    method of the same name, but — because a duplex connection has no
    fixed ``num_frames`` turn boundary — it never raises on exhaustion;
    the KV/state growth cap (design doc §3.5) is enforced elsewhere
    (design doc §7 step 8), not by this bookkeeping class.
  * The TTS-side ``cond_on_prev_audio_tokens`` feedback — ``run_turn``'s
    ``prev_code``/``past_kv`` locals
    (``arbi_serve/runtime/nemotron_voicechat_turn.py:465-474, 519-522``):
    promoted here to ``prev_tts_code``/``tts_past_kv`` fields, written via
    :meth:`observe_tts_frame` instead of a loop-local reassignment.

Both pieces are reused NEARLY VERBATIM in their update logic — only their
lifetime changes (connection-scoped: never reset except at connection
close or an explicit barge-in hard reset per design doc §3.2, instead of
turn-scoped: rebuilt fresh every turn). TTS-side gap, generalized: neither
``prev_text_token_id``/``prev_function_token_id`` nor ``prev_tts_code``
may ever be silently reset at a simulated "turn" boundary once a session
is duplex — there IS no turn boundary at this layer (design doc §3.3).

``agent_idle`` (+ its lock) mirrors the reference's
``S2SStreamingState._agent_idle_lock``/``agent_idle`` flag
(``s2s_state.py:93-100, 215-237`` per the design doc's citation) — the
flag, a thread-safe setter, and idempotent ``mark_agent_idle``/
BOS/EOS transition TABLE that
decides WHEN to call them — read every tick off the STT backbone's sampled
text-channel token id — is implemented in
:func:`arbi_serve.runtime.nemotron_voicechat_turn.run_turn` (design doc
§3.2.1's own "Where this hooks into the per-tick loop"), not in this
module.

One instance per WS connection running NemotronVoiceChat in duplex
mode (design doc §1.1's field table) — a sibling of
:class:`~arbi_serve.realtime.session.RealtimeSession`, not a subclass
(design doc §1's own framing: the two have almost disjoint state
shapes).

``eq=False``: several fields are tensors (``prev_tts_code``,
``pending_frame_embeds``); dataclass-generated ``__eq__`` would compare
them with ``==``, which is elementwise on a ``Tensor`` and raises on
``bool()`` — object identity is the only sound equality for a mutable,
per-connection singleton like this anyway.

Whether :attr:`frame_budget_used` is still under
:attr:`frame_cap` — a query only; unlike
``_FrameLockstepState.frames_remaining`` (which gates a hard raise
in :meth:`consume_frame`), THIS class never refuses to advance on
its own. Cap enforcement (closing the session cleanly, design doc
§3.5) is a later step's job — a caller that cares checks this
first.

Return this tick's frame index (the value of :attr:`frame_seq`
before advancing), then advance :attr:`frame_seq` and
:attr:`frame_budget_used`.

Connection-scoped counterpart of
``_FrameLockstepState.consume_frame``
(``nemotron_voicechat_stt_step.py:205-217``): that method raises
once ``num_frames`` frames have been consumed because a turn has a
fixed, known-in-advance audio length. A duplex connection has no
such fixed length — audio keeps arriving indefinitely — so this
method never raises; :attr:`frame_budget_remaining` is the
(advisory, caller-checked) analog of the turn-scoped
``frames_remaining`` gate.

Record this step's own sampled tokens as the NEXT fused step's
feedback — same contract as ``_FrameLockstepState.observe``
(``nemotron_voicechat_stt_step.py:219-223``): starts at
``text_pad_id`` ("blank"/listening) and is overwritten every call,
so a real emitted token feeds forward until the model goes back to
listening (both are correct — see that method's own docstring).

:attr:`prev_text_token_id` is never reset once set except by an
explicit new :class:`NemotronVoiceChatDuplexSession` instance: in
duplex mode every 80 ms tick's fused embedding always conditions on
the immediately-preceding tick's own emitted text token, listening
or speaking, forever (design doc §3.3). :attr:`prev_function_token_id`
follows the same rule with one exception —
:meth:`end_function_call` also resets it, since a resolved tool
round trip is the function channel's own boundary (see that
method's docstring).

Also reconciles the request's own committed token history — see
:meth:`_commit_text_feedback`.

Set :attr:`prev_text_token_id` AND mirror it into the request's
``output_token_ids`` — the ONE place a duplex tick's decided text
token is recorded (design doc §7.42).

Why this exists. ``finish_duplex_tick`` starts from the token the
generic sampler already committed to ``req.output_token_ids`` and
then runs six sequential rewrite gates over it (the FC-pause/ack
forcing, the speak gate, the duplicate-BOS guard, the EOU veto, the
EOU force-open, and the babble watchdog), plus the barge-in
override in :meth:`force_eos_feedback`. Every one of those routes
its answer here. Before this method existed, none of them wrote it
back, so a duplex request carried TWO divergent token histories:
``output_token_ids`` holding the RAW sampler output, and
``prev_text_token_id`` holding the REWRITTEN one that actually
conditions the next tick's fusion and is what the agent actually
said.

That divergence was benign only by accident — ``ignore_eos=True``
meant no stop logic read the raw stream, and ``run_step/output.py``
suppresses the output bus for duplex rows so nothing detokenized
it. Any generic feature that reads a duplex request's token history
(logprobs, stop strings, prefix caching, an operator dump, a trace
exporter) would have read the wrong stream and been silently wrong.

The invariant this establishes: **for a duplex request,
``output_token_ids[-1]`` is the text token that will condition the
next tick's fused embedding** — i.e. it always equals
:attr:`prev_text_token_id`. That is exactly what the reference does
on its own barge-in path, which overwrites its OWN token stream
(``generated_text_tokens[0, t] = eos_id``) rather than keeping a
second, truer history beside it.

KV-cache safety. This rewrites the last committed id IN PLACE; it
never appends or truncates, so ``len(output_token_ids)`` — and
therefore ``total_length``, the page accounting and every KV
position — is untouched. The id itself does not describe any
computed KV either: a duplex row's slot-0 embedding comes from
``pending_embed_override`` (the audio-frame fusion), not from
``embed_tokens(output_token_ids[-1])``, so this makes the recorded
id AGREE with what actually conditions the next position instead of
contradicting it.

Skipped when there is no request (a bare session in a CPU test, or
one not yet admitted) or when nothing has been committed yet — an
empty ``output_token_ids`` means this step sampled nothing for this
row (mid-prefill chunk), and appending there WOULD move a KV
position.

Append a REAL spoken text token to the request's penalty history
(design doc §7.53).

Called from the one place a duplex tick's text token is decided, so
the penalty history contains exactly the tokens the agent actually
SAID — in the same order, and only once each.

PAD/BOS/EOS are skipped, which is the whole special-token carve-out:
a token absent from the history is never penalized, so the turn
machinery (which re-emits PAD on ~95% of ticks and needs BOS/EOS to
stay freely reachable) cannot be touched by any penalty value. The
reference samples those three from a plain argmax for the same
reason (``infer/utils.py:191-194``).

``getattr``: minimal request stand-ins in CPU tests predate the
field, and a request that never opted in keeps the generic
prompt+output history.

Overwrite :attr:`prev_text_token_id` only, leaving
:attr:`prev_function_token_id` untouched — so the NEXT tick's
fusion conditions on ``eos_token_id`` instead of whatever THIS
tick's step actually sampled.

The barge-in AND-gate (design doc §7.26,
:func:`~arbi_serve.realtime.duplex_barge_in.vad_barge_in_fires`)
calls this alongside :meth:`mark_agent_idle`: setting the
client-visible flag alone leaves the model's OWN recurrent
conditioning on the text channel untouched, still following
whatever it actually said the instant before the interrupt — so
without this, a barged-in turn can keep regenerating from a
"mid-utterance" internal state indefinitely, immune to the flag a
caller just flipped. Mirrors the reference's own barge-in branch,
which overwrites the model's OWN token stream
(``generated_text_tokens[0, t] = eos_id``), not just an external
state flag — and, like that branch, forces only the text channel;
the function channel has no barge-in analog.

Routes through :meth:`_commit_text_feedback`, so the forced EOS
lands in the request's committed history too — the reference
overwrites its own token stream here for exactly this reason, and
leaving the two to disagree is the divergence that method exists to
close.

Hold the agent silent: forces this tick's TEXT token to
:attr:`text_pad_id` through the same three consumers
:attr:`fc_in_progress` forces it through
(:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`)
until :meth:`open_speak_gate` releases it. A caller that owns turn
boundaries (a repointed bounded session) calls this before feeding
a turn's audio, so the model can listen without being free to
answer until asked — the turn-based contract, reproduced on the
duplex path.

Whether a tool call is currently pending for this connection —
the reference's ``_fc_in_progress``.

While ``True``,
:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
replaces the tick's sampled TEXT token with :attr:`text_pad_id`
before :meth:`observe`, the Site-3 BOS/EOS switch and the TTS
subword ever see it, so the agent's text/audio stream goes silent
for the duration of the call. The frame grid, the audio-in side and
the function channel itself keep running — this is a mute, not a
pause of the recurrence (design doc §3.2.1: "keep the frame grid
advancing and force the agent TEXT channel to PAD").

Derived from :attr:`fc_state` rather than being its own field so
there is exactly ONE thing to reset: clearing ``fc_state`` reopens
the gate, and a ``fc_state`` that exists only because the detector
has been scanning is correctly NOT "in progress".

Reopen the text channel after a tool call resolves, times out,
or is abandoned — idempotent, and safe to call from the WS thread
while the engine thread is mid-tick (a plain attribute store; the
engine side only ever READS the flag).

Deliberately does NOT drop :attr:`fc_state` itself: the detector's
scan buffer is connection-lifetime state (a duplex connection has
no turn boundary to rebuild it at) and must survive a completed
call so a LATER ``<TOOLCALL>`` block is still detected. The
``ToolResultBridge`` timeout hatch's harder ``fc_state = None``
reset is the deliberate exception — an abandoned call is exactly
when starting the scan clean is right.

Any ack phrase still unspoken IS dropped (design doc §7.9.18): the
gate's contract is "while a call is pending", and the tick that
reopens it is the tick the tool response is injected on, so
continuing to force ack tokens would be forcing them over the
model's own resumed speech. A phrase is a handful of frames and a
real call is measured in tens, so this truncation is reachable
only for a tool that answers almost instantly — where there was
nothing to fill in the first place.

Also resets :attr:`prev_function_token_id` back to
:attr:`text_pad_id` (design doc §7.69). A resolved round trip is
the function channel's own boundary — the call is answered and
there is nothing further to say about it — so its own feedback
returns to blank exactly like a fresh connection's does, rather
than continuing to condition every later tick's fused embedding
(at :data:`~arbi_serve.models.nemotron_voicechat
.DUPLEX_FUNCTION_CHANNEL_WEIGHT`, the single highest-weighted
input any tick receives) on whatever the channel's own
unconstrained continuation settles into once a call closes with
no terminator ever entering its own token stream. Unconditional,
not nested under the ``fc_state is not None`` check below: the
``ToolResultBridge`` timeout hatch clears ``fc_state`` to
``None`` before this method ever runs on that path, and the
function channel needs the same reset whether a round trip
resolves cleanly or times out.

Record that a ``<TOOL_RESPONSE>`` just entered the running
sequence, so the agent may open a turn to speak it — design doc
§7.63.

Called from the WS thread by
:meth:`~arbi_serve.realtime._session_bounded_duplex._BoundedDuplexMixin
._resolve_duplex_function_calls` on the same tick it pushes the
response tokens. A plain attribute store, matching
:meth:`end_function_call`'s own thread contract: the engine side
only ever READS the flag.

Why this is needed at all. :meth:`eou_admits_bos` admits a
model-sampled BOS only on confirmed user speech since the agent's
last turn opened, and :meth:`mark_agent_active` spends that
evidence the moment the ACK turn opens. A tool round trip
therefore ends with the agent holding an answer, the user silent
by construction (they asked once and are waiting), and the gate
refusing every BOS the model raises to deliver it — measured live
as seven consecutive refusals in the ~1.5 s after the response was
injected, then permanent silence for the rest of the connection.
SERVER knows the agent has to say.

Arms the VETO half and STANDS DOWN the force half — the delivering
turn belongs to the model's own turn-taking head. This section
always intended the first half ("deliberately arms the VETO half
only, not :meth:`eou_forces_bos`", on the reasoning that the veto
exemption cannot invent a turn and the model demonstrably does
propose one here), but arming the veto is not enough on its own:
:meth:`eou_forces_bos` is suppressed during a tool call ONLY by
``fc_in_progress``, while :attr:`eou_silent_ticks` keeps
accumulating unconditionally through the whole round trip. It is
therefore already over threshold the instant the pause lifts, and
fires on that frame — one frame BEFORE the injection tick that
actually puts the response into the sequence
(``push_context_tokens``: "the NEXT duplex tick becomes an
injection tick"). Measured live, 3 captures out of 3: a turn
opened with the model unable to have read the answer yet, 0 real
text tokens, EOS after 3 ticks, and the exemption spent on it.

So the force half is stood down for as long as the veto half is
armed, which is what makes "the model proposes the delivering BOS"
true rather than merely intended. :meth:`tool_result_owed_now`
bounds that stand-down to one :attr:`eou_silence_ticks` window —
the force-open's own — so the worst case is a turn forced one
window later than today, by which point the injection tick has
certainly landed and the model has something to say.

Queue an ``ack_messages`` phrase to be SPOKEN through the pause
gate — the duplex equivalent of the turn-based path's
``_maybe_speak_ack`` (design doc §7.9.18).

One id is forced per 80 ms frame, in place of the PAD the gate
would otherwise force, which is exactly the rate (and the
mechanism) ``_synthesize_ack_waveform_frames`` drives its own
standalone TTS loop at: this checkpoint's STT->TTS coupling is a
bare id passthrough, so a fully-known phrase's own token ids stand
in for per-tick STT samples with no re-tokenize round trip.

Ignored unless a call is actually pending: the queue exists only
for the duration of a gate, and a phrase arriving after the tool
already answered has nothing left to fill.

Called on the ENGINE thread (the duplex lane's pre-schedule hook
drains its own WS-side staging deque into here), for the same
reason ``pending_context_token_ids`` is written there.

Thread the TTS side's ``cond_on_prev_audio_tokens`` feedback
forward onto this connection-lifetime object — the same update
``run_turn`` performs on its own turn-scoped locals
(``prev_code = new_code`` /
``past_kv = model.generate_tts_frame(...)``'s return,
``nemotron_voicechat_turn.py:498-522``), except never reset at a
turn boundary (design doc §3.2). ``prev_tts_code``/``tts_past_kv``
now survive across simulated turn boundaries exactly the way
:meth:`observe` already makes ``prev_text_token_id``/
``prev_function_token_id`` survive on the STT input side.

Thread-safe flip of :attr:`agent_idle`, guarded by
:attr:`agent_idle_lock` — mirrors
``S2SStreamingState._agent_idle_lock``. The BOS/EOS transition
table that decides WHEN to call this with which value is design
doc §7 step 5's own follow-up (undesigned as of this task, §6
decision point 5) — this method only makes the flag itself safe to
mutate from whatever caller eventually implements that state
machine.

Name what is about to end this turn, for the ledger
(design doc §7.50).

Separate from :meth:`mark_agent_idle` rather than an argument to
it because that method is part of
:class:`~arbi_serve.realtime.duplex_model_adapter.DuplexSessionLike`'s
zero-argument structural contract, which generic call sites (the
pump's barge-in, a client cancel, a tool-call timeout) reach
through ``getattr`` and which every minimal session stand-in in
the CPU tests implements. Consumed and cleared by the next
:meth:`mark_agent_idle`.

Idempotent convenience wrapper over ``set_agent_idle(True)`` —
matches the reference's ``s2s_state.py`` API shape (design doc
§3.2.1's "What arbi-serve concretely needs to build" recommends
adding these so escape-hatch/Site-3 call sites read clearly).
Safe to call regardless of the current value.

Also resets the babble watchdog's per-turn counters (design doc
§7.23): every caller of this method — the Site-3 switch, barge-in,
``response.cancel``, a tool-call timeout — is ending a turn, and
the watchdog's accounting is scoped to one turn's speech.

Also emits the turn's audio ledger (design doc §7.50) BEFORE those
resets, since it reports the very counters they clear.
``close_reason`` names what ended the turn for that line; a caller
that cannot pass it (the zero-argument structural contract) sets
it through :meth:`note_close_reason` instead.

Also resets the EOU admission tracker (design doc §7.30/§7.32):
:attr:`eou_user_spoke`/:attr:`eou_nonblank_total` — a turn just
ended, so the NEXT candidate BOS needs its own fresh
confirmed-speech signal, not whatever was left over from before
this one opened. Mirrors the reference's own reset-on-EOS
(``model.py:1147``, which clears ``rnnt_user_speaking`` AND
``rnnt_nonblank_total``).

Idempotent convenience wrapper over ``set_agent_idle(False)`` —
the BOS-side counterpart of :meth:`mark_agent_idle`. Also resets
the babble watchdog's counters (design doc §7.23), so a new turn
starts its own accounting from zero rather than inheriting
whatever the previous turn left behind.

Also stamps :attr:`agent_has_opened_once` and resets the EOU
admission tracker (design doc §7.30/§7.32), mirroring the
reference's own reset-on-BOS (``model.py:1137``, which clears
``rnnt_user_speaking`` AND ``rnnt_nonblank_total``): the
confirmed-speech signal that justified opening THIS turn is
spent, and the connection has now used its one ungated "may open
with nothing confirmed yet" turn — see :meth:`eou_admits_bos`.

Record that THIS tick's synthesized waveform was replaced with
digital silence before it reached the client (design doc §7.50).

Called by :mod:`~arbi_serve.runtime.nemotron_voicechat_turn`'s
PAD-tail substitution, which is the only place a generated frame
is muted. Observability only — nothing decides on this counter.

Emit ONE line accounting for the audio a just-ended turn
produced — design doc §7.50's standing instrument for the
"the audio cut off but the text kept going" class of report.

A turn's text and its audio are generated in per-tick lockstep,
so the only ways the client can end up with less speech than the
text promised are: the turn ended while the TTS still had queued
subwords to render (``close_reason`` says which mechanism ended
it), or frames were generated and then MUTED
(:attr:`agent_muted_frames`). Both are reported here, next to the
text-token count they have to be read against, so the next report
is a log read rather than another live investigation.

A no-op unless the turn actually produced frames, which also
keeps the idempotent second :meth:`mark_agent_idle` call the
Site-3 switch makes on the same tick from logging twice.

Record one tick of continued agent speech for the TTS
ratio-cap watchdog (design doc §7.23) and report whether the cap
has just tripped.

Mirrors the reference's ``_apply_turn_taking_from_rnnt_states``
TTS CAP branch: :attr:`agent_talking_frames` counts every tick the
agent is speaking; :attr:`turn_text_tokens` counts only the ones
that carried genuinely NEW text content (``is_real_text`` — the
caller excludes BOS/EOS/PAD). The cap fires once both
:attr:`tts_ratio_min_tokens` real tokens have been seen AND the
agent has spoken :attr:`tts_ratio_cap` audio frames for each one —
i.e. the agent is producing audio much faster than new content,
the shape a stuck/degenerate turn has whether or not the model
ever emits a real EOS token on its own.

Called once per tick, only on a tick where the agent is speaking
and this tick's own token is neither BOS nor EOS
(:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
owns that gating). The BOS tick itself resets both counters to
zero via :meth:`mark_agent_active` rather than counting as the
turn's first frame — a one-frame difference against a cap
measured in dozens of frames, well inside its own granularity.

Whether an injected tool answer is still owed AND still inside
the window it may bend turn-taking for — design doc §7.64.2.

The single read point for :attr:`tool_result_owed`:
:meth:`eou_admits_bos` exempts a BOS while this holds, and
:meth:`eou_forces_bos` stands down while it does, so the delivering
turn is the model's own to propose.

**Why it is bounded.** Standing the force-open down is safe only
because it cannot last: an unbounded stand-down would mean a model
that never proposes a BOS — §7.37.2's cold always-PAD fixed point,
the exact case the force half exists for — leaves the agent mute
for the rest of the connection, which is worse than the bug this
fixes. The bound is :attr:`eou_silence_ticks`, the force-open's OWN
window, rather than a new constant: the guarantee it buys is
therefore "a turn this defers is forced at most one force-open
window later than it would have been", and by that point the
injection tick has certainly landed. Measured live, the model's own
BOS arrived 7, 7 and 26 ticks after injection across three
captures, all well inside the 40-tick default.

``eou_silence_ticks <= 0`` disables the EOU machinery outright, in
which case :meth:`eou_admits_bos` already admits everything and
:meth:`eou_forces_bos` already forces nothing; this returning
``True`` there changes neither.

Whether this turn's own text channel has said it is finished —
the ADDITIVE CLOSE-half of the turn-taking mechanism (design doc
§7.44), and the exact mirror of :meth:`eou_forces_bos` on the other
end of a turn.

Read by
:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
on every tick of an open turn, immediately after
:meth:`note_babble_frame` has updated the counters for THIS tick.
``True`` forces this tick's ``text_token_id`` to the checkpoint's
turn-EOS id, so the turn closes through the ordinary Site-3 switch
and nothing downstream can tell it from a model-emitted EOS —
exactly the way :attr:`tts_ratio_cap`'s watchdog already closes a
babbling one.

**Why this has to exist.** §7.37.2 measured that under the
checkpoint's own greedy default the text channel keeps BOS a
persistent runner-up to PAD, so nothing could ever OPEN a turn the
model did not argmax; §7.37 built the force-open that supplies it.
The closing side has the identical shape and was never built.
Measured live (§7.44), on the tick after an answer's last word the
text channel goes to PAD and stays there, with **EOS at rank 1 on
every subsequent tick** and a stable ~10-13 logit margin behind
PAD — the model is telling us it is done and greedy argmax can
never act on it. The observed tails were 26.9 s, 30.5 s, 37.8 s and
47.9 s, one of which never ended inside the capture at all.

The turn's only other close conditions cannot bound that:

  * the model's own EOS — the thing just measured not to arrive;
  * :attr:`tts_ratio_cap` — a RATIO, so its trip point is
    For the measured 51-token turn that is 816 frames of
    dead air (65 s); for the 75-token one, 96 s. It is a
    degenerate-babble backstop, not a turn boundary, and §7.39.5
    was right that it "correctly declines" — it simply cannot see
    this;
  * barge-in — requires the USER to interrupt, and a turn that
    will not close is precisely a turn during which the user
    cannot be heard: :attr:`eou_user_spoke` is only ever set while
    :attr:`agent_idle` (:meth:`note_vad_events`,
    :meth:`note_semantic_frame`), so an unclosed turn makes the
    connection deaf and no later signal can rescue it.

Deliberately NOT gated on :attr:`tts_ratio_min_tokens`, unlike the
two ratio mechanisms. Those need the floor because their
denominator is ``turn_text_tokens`` and a near-zero denominator
trips them on noise; this counter has no denominator. The floor
would also exempt the single worst case — a turn that opens and
produces NO real text at all (a force-open the model had nothing
for) would never reach the floor and so could never be closed by
anything. A run of :attr:`agent_pad_tail_close_ticks` consecutive
PAD ticks means the turn is over whether or not a word was said.

**And why the PAD run alone is not sufficient** (design doc
§7.50). §7.44 read the first PAD after an answer's last word as
the start of dead air. That is true of the TEXT channel and false
of the AUDIO one: this checkpoint's text channel emits a subword
on EVERY tick of an answer (§7.44.1's zero intra-answer gaps),
i.e. at 12.5 subwords/s, while its TTS renders that same answer at
roughly a THIRD of that rate — measured live at **2.92 audio
frames per real text token** (§7.50.2: an 89-token answer whose
text channel finished at 7.1 s and whose audio was still
mid-sentence until 20.8 s). The PAD tail is therefore not dead
air; it is where most of the answer is actually SPOKEN, and a
fixed 40-tick (3.2 s) run ends the turn — and with it the audio —
with the model roughly half way through saying it. That is the
"text runs ahead of the audio, which cuts off" report.

So the run is necessary but not sufficient: the turn also has to
have spent its audio-rendering budget, which is exactly what
:meth:`pad_tail_should_mute` already computes from the
reference's own ``S2S_TTS_PAD_TAIL_RATIO``. Past that point the
client is being sent substituted silence anyway, so closing costs
nothing; before it, closing costs the rest of the sentence.
Measured against the same live capture, the reference's boundary
— PAD-conditioned hallucination
resumes, at every turn length. §7.56 corrected which counters
that ratio reads; this predicate is unchanged and simply
inherits the corrected boundary.

The budget clause is SKIPPED when it does not apply — the mute
disabled outright, or fewer than :attr:`tts_ratio_min_tokens`
real tokens, which is the near-zero denominator the ratio
mechanisms cannot speak about. That keeps §7.44's worst case
exactly as it was: a turn that opens and says nothing still
closes on the bare consecutive-PAD run.

Whether a MODEL-SAMPLED

§7.50 established the rule this applies: TTS, so a close taken on a
text-channel signal discards whatever the renderer still owes.
§7.50 enforced that on ONE of the three close paths — its own
PAD-tail close (:meth:`agent_pad_tail_ends_turn`) — because that
was the path its captures exercised, and §7.44 had measured the
model's own EOS as "the thing just measured not to arrive".

In a
10-turn session two turns closed on ``the text channel's own The one turn in that capture that closed on the
budget-guarded Same session, same
checkpoint, same question as a truncated turn — the close path is
the only variable.

So the guard belongs on the quantity, not on one path to it. A
model EOS arriving while the turn still owes real audio is
DEFERRED (rewritten to PAD by the caller), which does not lose the
close: the turn flows into its PAD tail exactly as a turn whose
model never emitted EOS already does, and
:meth:`agent_pad_tail_ends_turn` — already budget-aware — ends it
the moment the budget IS spent. That is the path the complete
246-frame turn took.

BARGE-IN is deliberately NOT routed through this. An interrupt is
the user asking the agent to stop talking; making it wait out the
rest of the sentence would defeat its whole purpose. NAMED-not-fixed question of whether arbi-serve's one-VAD-edge
trigger is too eager next to the reference's 3.2 s one.

Skipped exactly where :meth:`agent_pad_tail_ends_turn` skips it —
mute disabled, or fewer than :attr:`tts_ratio_min_tokens` real
tokens — so a turn with no denominator still closes on its own EOS
immediately, as it always did.

Whether THIS tick's TTS output should be replaced with genuine
silence rather than the backbone's own PAD-conditioned generation
(design doc §7.31.6/§3.2's flagged unknown (4),
``S2S_TTS_PAD_TAIL_RATIO``).

Reads :attr:`agent_pad_tail_ticks`/:attr:`turn_text_tokens` —
:meth:`note_babble_frame` has already updated both for THIS tick
by the time a caller reaches this method
(:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
calls it before :mod:`~arbi_serve.runtime.nemotron_voicechat_turn`'s
own per-tick TTS-generation call runs), so the ratio this computes
already includes this tick's own frame/token count.

**The numerator is the CONSECUTIVE PAD RUN, not the turn's total
frame count** (design doc §7.56). The reference's condition is
``tts_in_turn_pads > S2S_TTS_PAD_TAIL_RATIO * tts_in_turn_content``
(``model.py:658-665``), and its two counters are not symmetric:
``tts_in_turn_pads`` is reset to 0 by every content token while
``tts_in_turn_content`` accumulates across the whole turn — i.e.
exactly this class's :attr:`agent_pad_tail_ticks` and
:attr:`turn_text_tokens`. This method used to read
:attr:`agent_talking_frames`, which counts content ticks AND pad
ticks, so

That gap is the whole "the last sentence is abruptly cut" report.
Measured natural rendering on this checkpoint is **2.92 audio
frames per real text token** (§7.50.2 run C), i.e. The reference's **2.7%** headroom, so any turn rendering even
marginally slower than the average (a multi-sentence answer pays
inter-sentence pauses in frames and nothing in tokens) had real,
still-speaking audio replaced with silence.

No :attr:`tts_ratio_min_tokens` floor, also matching the
reference, which has none here. The floor was this method's own
addition ("a thinking pause before the first word is not a PAD
tail"), and with the corrected numerator it is not merely
unnecessary but harmful: a turn with zero content tokens is
precisely the one whose every frame is PAD-conditioned, and the
floor exempted it from muting entirely — a force-opened turn the
model had nothing to say for broadcast its full
:attr:`agent_pad_tail_close_ticks` run of hallucination at full
volume. With ``content == 0`` the reference mutes from its second
PAD tick, and a real word arriving later resets the run to 0 and
unmutes on the same tick.

One tick tighter than the reference, unavoidably: the reference
decides the substitution from the PREVIOUS frame's counters and
updates them afterwards (``model.py:679``'s own note), while
``note_babble_frame`` has already counted THIS tick by the time
anything calls this. So the boundary lands at ``4C + 1`` total
frames where the reference's is ``4C + 2``. One 80 ms frame,
named rather than papered over.

The floor still guards the two CLOSE predicates
(:meth:`agent_pad_tail_ends_turn`, :meth:`eos_closes_turn_now`),
which compute it themselves — muting a frame and ending a turn
are different decisions and only the second one needs a
denominator it can trust.

How many CONSECUTIVE blank ticks clear the BOU evidence
accumulator (:attr:`eou_nonblank_total`) — design doc §7.46.

The reference's own constant is :data:`_EOU_NOISE_RESET_TICKS`
(``RNNT_NOISE_RESET_FRAMES=10``, 800 ms), and porting it verbatim
made the semantic feed unable to confirm BOU for a normal, short
utterance. Measured on live production traffic (§7.46.2) against the
bundled RNNT head: it argmaxes non-blank on only **1.9% of all
ticks and 11-20% of ticks inside genuine continuous speech** —
it is blank BETWEEN subwords by construction, since one 80 ms
frame carries at most one greedy label step's worth of evidence.
Over 3969 traced ticks :attr:`eou_consecutive_speech` reached 3
exactly three times and 6 never, so the consecutive-run half of
the BOU test is effectively dead and confirmation rests entirely
on this cumulative counter — which an 800 ms gap zeroed, and
800 ms gaps occur INSIDE every utterance. A user speaking a
normal one-second sentence therefore accumulated 0-2 non-blank
ticks, never reached the 3-or-6 threshold, and was never heard:
:attr:`eou_user_spoke` stayed ``False`` for the life of the
connection, so :meth:`eou_forces_bos` could never fire and
:meth:`eou_admits_bos` refused every BOS the backbone proposed —
a connection that is live, ticking and completely deaf, with no
log line to say so.

So the window is the connection's OWN definition of "this
speaker has stopped", :attr:`eou_silence_ticks` (default 40
ticks = 3.2 s) — the same duration :meth:`eou_admits_bos` and
:meth:`eou_forces_bos` already use to decide an utterance ended.
A gap shorter than that is BY THIS CLASS'S OWN DEFINITION still
inside the utterance, and evidence that the utterance is
happening must survive it.

The reference constant is kept as a FLOOR, so a deployment that
lowers or disables the EOU gate (``eou_silence_ticks <= 0``)
still gets the reference's own reset rather than none at all.

"Sparse non-speech blips must never accumulate into a fake BOU"
is enforced by this reset window itself plus the head's own
content gate (it decodes non-speech to blank) — the separate
low-density threshold raise the reference layered on top is
retired for this feed; see :meth:`note_semantic_frame` for the
measured reason (design doc §7.48).

Update the EOU admission tracker from this tick's RNNT
endpoint-detector verdict (design doc §7.32) — called once per
tick from :meth:`~arbi_serve.realtime.duplex_tick_pump
.DuplexTickPump.post_stt` on a connection whose frames carry the
checkpoint's own per-frame blank/non-blank decode
(:meth:`~arbi_serve.models.audio.nemotron_voicechat_rnnt
.NemotronVoiceChatRnntHead.step_frame`). A tick that consumed no
client audio (a jitter-gap silence frame) counts as blank: no
audio is silence, which is also what the reference's own
silence-template frames decode to.

Ports the reference's per-frame counter update
(``_rnnt_eou_decode_frame``'s tail) plus the BOU/noise-reset
branches of ``_apply_turn_taking_from_rnnt_states``:

  * the rolling speech-density EMA advances only while the agent
    is idle (the reference's ``if not rnnt_agent_speaking``
    gate);
  * a non-blank tick zeroes :attr:`eou_silent_ticks` and advances
    :attr:`eou_consecutive_speech`/:attr:`eou_nonblank_total`; a
    blank tick zeroes the consecutive run and accumulates
    silence;
  * the noise reset: :meth:`_bou_noise_reset_ticks` of continuous
    silence with no BOU yet confirmed clears
    :attr:`eou_nonblank_total`, so sparse non-speech blips can
    never accumulate into a fake BOU across a long quiet gap —
    sized off :attr:`eou_silence_ticks` rather than the
    reference's own 800 ms constant, because the RNNT is blank
    between subwords and an 800 ms window zeroed the evidence
    mid-utterance (design doc §7.46);
  * BOU: :attr:`eou_user_spoke` is set once
    :attr:`eou_consecutive_speech` OR :attr:`eou_nonblank_total`
    reaches the effective minimum — relaxed before the
    connection's first turn (``rnnt_first_turn`` ≡ ``not
    agent_has_opened_once``), raised while the rolling density is
    low-but-nonzero (sparse/noisy audio), and gated on
    :attr:`agent_idle` exactly like the reference's own BOU
    condition (speech while the agent talks is barge-in
    territory, a different gate).

This is the semantic half of what makes the admission signal
content-aware: the RNNT head emits non-blank only for frames it
can actually transcribe, so breathing/mouth noise/background
sound — real acoustic energy that a VAD (neural included) can
classify as speech — never confirms BOU here.

**``is_blank=None`` is UNKNOWN, not blank** (design doc §7.48). A
synthesized jitter-gap silence frame carries no verdict because it
carries no microphone audio at all — the engine ran ahead of the
client's real-time delivery, and the mic frames on either side of
the gap are ADJACENT in microphone time. Live-measured tick rates
run ~10% above the 80 ms grid (§7.48.2), i.e. roughly one
synthesized frame per second, so treating these as blank was
resetting :attr:`eou_consecutive_speech` about once a second
INSIDE genuine utterances — which is why §7.46.2 measured the
consecutive-run half of the BOU test reaching 3 only three times
in 3969 ticks. An unknown tick therefore advances
:attr:`eou_silent_ticks` (so a client that stops streaming
entirely still reaches EOU/force-open) and runs the noise-reset
check, but neither builds nor destroys speech evidence: no
density update, no consecutive-run reset, no cumulative count.

**The low-density threshold raise is retired for this feed**
(design doc §7.48). The reference's guard raises the confirmation
bar to :data:`_EOU_BOU_LOW_DENSITY_MIN_TICKS` while the rolling
density sits in ``(0, _EOU_DENSITY_THRESHOLD)`` — designed for
sparse blips in mostly-quiet ambient audio. Against THIS head it
is unsatisfiable by real speech: the measured in-speech duty
cycle is 11-20% of ticks (§7.46.2), i.e. genuine continuous
speech itself sits INSIDE the guard's "sparse" band, so no
density threshold in that band can admit an ordinary utterance
while refusing the blip pattern the guard exists for — the two
are the same signal at this duty cycle. Worse, the EMA never
returns to exactly ``0.0`` once any speech has ever been seen
(§7.46.5's named latent deadlock), so after the first utterance
the raised bar applied for the rest of the connection's life:
6 non-blank ticks is ~2.3 s of CONTINUOUS speech at this duty
cycle, and a normal 1.5-3 s conversational sentence could never
confirm BOU again — measured live as the minutes-long
"deaf while the agent idles" stretch of §7.48.2. The noise
defense this feed actually has is the head itself: it emits
non-blank only for frames it can transcribe (the module's own
§7.32 contract), which is a content gate no acoustic density
proxy approximates. The EMA is still maintained (and snapped to
a true zero once its residue decays below one tick's
contribution aged past the connection's own end-of-utterance
duration) as connection-state observability; nothing decides on
it any more.

Advance :attr:`eou_barge_in_ticks` — how long the user has been
talking THROUGH the agent (design doc §7.57).

Runs on every semantic tick, before the ``agent_idle`` guard,
because this counter exists only for ticks the guard excludes.
The run advances whenever the last non-blank verdict is within
:data:`_EOU_BARGE_IN_GAP_TICKS`, so the blank gaps that separate
subwords on this head do not break it, and it is zeroed the
moment the user is quiet for longer than that.

Deliberately NOT ``eou_consecutive_speech``: that resets on every
blank tick, and genuine continuous speech on this head is 63-82%
blank, so a consecutive-non-blank run of any useful length is
unreachable by real speech (design doc §7.57.2).

Whether the user has now talked through the agent for long
enough to count as an interruption (design doc §7.57).

The content-aware half of the barge-in gate, and the one that
catches a user who never stopped talking:
:func:`~arbi_serve.realtime.duplex_barge_in.vad_barge_in_fires`
needs a ``speech_started`` EDGE, which a continuing speaker never
produces, so a turn opened mid-utterance is otherwise immune to
the interruption it is talking over.

``barge_in_speech_ticks <= 0`` disables this half entirely,
leaving the acoustic edge alone — the same escape hatch
:attr:`eou_silence_ticks` uses.

Bin this tick's frame energy for the input-SNR estimate (design
doc §7.59).

Takes the audio and NOTHING else, deliberately. The first version
of this estimator took the RNNT head's per-frame ``is_blank``
verdict and split energy on it, and that is measurably wrong: the
head is blank for ~70% of frames DURING continuous speech
(§7.58.2 — blank means "no token this frame", not "no speech"), so
it put most of the speech in the background bucket and reported
5.2 dB for a clean 30 dB connection. Keeping the verdict out of
the signature is what stops that being reintroduced.

A frame with no PCM at all (a synthesized jitter-gap tick) is
skipped: it is an absence of evidence, and binning it as a very
quiet frame would drag the noise floor down and flatter the room.

This connection's loud-to-quiet frame-energy ratio in dB, or
``None`` until enough frames have arrived to mean anything.

Not a calibrated acoustic SNR: it is the 90th percentile of frame
energy over the 10th, which is the quantity that tracks this
checkpoint's accuracy (§7.59) and needs no microphone calibration,
no second detector and no per-frame classification to compute.

Say once, out loud, that this connection's audio is too noisy for
the checkpoint to transcribe reliably (design doc §7.59).

The live report this exists for was *"it's almost like the hearing
is quite faint acoustically"*, and the server had no way to agree or
disagree: LEVEL is not the variable (this checkpoint is
level-invariant — 30 dB of attenuation moves word error by nothing,
because log-mel turns a gain into an offset and the encoder's own
``layer_norm`` absorbs it), but the speech-to-background RATIO very
much is. Without this line a noisy room and a serving bug produce
the same transcript and the same silence in the log.

Say once, out loud, that this connection has stopped hearing
anything — and say which side of the wire went quiet (design doc
§7.54).

A live session ran 3275 consecutive silent ticks (262 s) with the
agent idle and the backbone proposing BOS on 40 of them, and the
only trace was a stream of INFO refusals whose numbers climbed
forever. From the tester's chair the conversation simply stopped;
from the log nobody could say whether the microphone had stopped
or the server had stopped hearing it, because a tick fed a
synthesized silence frame and a tick fed real audio the ASR head
decoded to blank are both just ``eou_silent_ticks += 1``.

:attr:`eou_silent_real_ticks`/:attr:`eou_silent_synth_ticks` carry
that split, so this fires ONE warning per silent stretch — latched
until the next non-blank tick or turn boundary clears it — naming
the dominant cause. It is deliberately a WARNING: an agent that
cannot hear its user is not a routine event, and §7.49.3's whole
finding was that this class of failure was invisible.

Only while :attr:`agent_idle`: a connection listening to its own
turn is not deaf, it is talking.

Below this, the rolling-density EMA is a stale residue, not a
live ambient-audio measurement: the contribution one non-blank
tick retains after :attr:`eou_silence_ticks` ticks of decay — the
connection's own "this speaker has stopped" duration, i.e. the
same boundary every other EOU decision in this class uses.
Floored at the equivalent for the reference's own 10-tick noise
window so a deployment that disables the EOU gate keeps a finite
floor.

Record the RNNT endpoint detector's own emitted subword ids for
one encoded frame — the server-side "what did I actually hear"
evidence (design doc §7.48). Called at encode/push time by the
adapter's ``encode_audio_chunk`` (WS worker thread); drained by
:meth:`take_asr_heard_ids` at each turn open (engine thread). The
head already produces these ids on every frame and they were being
discarded — which left a context-blind live session with no way to
show whether the failure was the audio the client delivered or the
server's own understanding of it.

Drain and return every heard label id since the last call —
one turn-boundary window of the server's own ASR evidence. Reader
side of :meth:`note_asr_labels`'s producer/consumer pair; popleft
per element keeps the deque single-consumer-safe against the WS
thread's concurrent appends.

Update :attr:`eou_user_spoke`/:attr:`eou_silent_ticks` from this
tick's already-fed VAD events — the RNNT semantic feed (:meth:`note_semantic_frame`): a bounded/
turn-shaped session fed pre-encoded frames, a client that
explicitly selected an acoustic ``turn_detection.backend``, or a
served model without the bundled RNNT head. Called once per tick
from :meth:`~arbi_serve.realtime.duplex_tick_pump.DuplexTickPump
.post_stt` with the SAME event list the barge-in AND-gate consumed
(:func:`~arbi_serve.realtime.duplex_barge_in.vad_tick_events`).
Never feeds the VAD itself.

Rebuilds a private in-speech mirror from the VAD's own discrete
``speech_started``/``speech_stopped`` edges (the VAD reports
transitions, not a per-tick level), then updates the two admission
signals for every tick, mirroring the reference's own
``rnnt_user_speaking``/``rnnt_silent_frames`` bookkeeping:

  * while voiced: :attr:`eou_silent_ticks` resets to zero, and
    :attr:`eou_user_spoke` is set — but only while
    :attr:`agent_idle`, matching the reference's own BOU condition
    (``model.py:1159-1169``, gated ``and not
    seq_state.rnnt_agent_speaking``) — speech arriving while the
    agent is talking is barge-in territory (a different gate), not
    a signal that the user has something new to say once the
    agent finishes;
  * while silent: :attr:`eou_silent_ticks` accumulates
    unconditionally (mirrors the reference's own unconditional
    ``rnnt_silent_frames += 1``) — :meth:`eou_admits_bos` is what
    decides whether the accumulated count matters, not this
    method.

Whether a candidate new-turn BOS should be honored right now —
the confirmed-speech-then-confirmed-silence admission check design
doc §7.30 builds from the reference's own EOU branch
(``model.py:1171-1188``), adapted onto arbi-serve's VAD. Read by
:func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
for every text-channel BOS the agent's own sampler (or a forced
EOU-style open — arbi-serve has none; every candidate BOS here is
model-native) produces while :attr:`agent_idle`.

**A client-owned-turn (bounded) connection uses a different rule
entirely** (:attr:`client_turn_audio_target` set — design doc
§7.35): admit exactly once this turn's committed real audio has
been consumed (:attr:`frame_seq` reached the target). Neither
half of the VAD rule below is meaningful there. The
first-turn exemption is wrong because the client has already told
us WHICH utterance it wants answered, so an open before that
audio is consumed answers a question the model has not heard —
the canned-greeting bug (§7.35). The confirmed-speech half is
worse than wrong: a bounded connection is built with ``vad=None``
(:meth:`~arbi_serve.realtime._session_bounded_duplex._BoundedDuplexMixin
._ensure_bounded_duplex_conn`, which deliberately does not share
the session's turn-boundary VAD), so
:meth:`note_vad_events` only ever receives an EMPTY event list,
:attr:`vad_user_spoke` can never become ``True``, and every BOS
after the first would be refused forever — one answer per WS
connection, then permanent silence.

BOS
sampling at any turn — which is true of the reference's veto
(``model.py:1134`` honors a sampled BOS unconditionally) but was
the wrong conclusion to draw, for two reasons:

* **it is inconsistent with this class's own force half.**
  :meth:`eou_forces_bos` already refuses to exempt turn one, in
  those words and for that reason. A connection whose agent has
  never spoken had a veto that waved a BOS through and a force
  that would not raise one — the two halves of one mechanism
  disagreeing about whether turn one is special;
* **what it actually produced was one fixed sentence, forever.**
  The text channel is decoded A deterministic decoder on a fixed prefix
  MUST emit the same tokens, so every session on the deployment
  opened with the same canned greeting, word for word. That is
  not a defect in the greeting; it is what a content-free
  proactive turn can only ever be under greedy decoding.

The fix is to stop taking the turn, not to randomize it. Waiting for the user costs nothing — the user's
first utterance is answered by :meth:`eou_forces_bos` on the same
confirmed-speech-then-silence rule that already serves every later
turn, which is exactly what the reference does (its own turn-one
special case is a RELAXED BOU threshold,
``RNNT_BOU_MIN_FRAMES_FIRST_TURN``, mirrored here at
:meth:`note_semantic_frame` — never a bypass).

Every candidate BOS therefore requires real confirmed user speech
since the agent's last turn opened (:attr:`eou_user_spoke`). The
confirmation signal is the checkpoint's OWN RNNT endpoint detector
whenever the connection carries it (:meth:`note_semantic_frame`,
design doc §7.32) — the SAME ``rnnt_user_speaking`` ground truth
the reference's own turn-taking reads, on the same bundled
weights — degrading to the acoustic-VAD rebuild
(:meth:`note_vad_events`) only where that head is unavailable.
``eou_silence_ticks <= 0`` disables this gate entirely (always
``True``), the same escape hatch :attr:`tts_ratio_cap` uses.

**This gate deliberately does That clause could not
prevent a single turn — it could only DELAY one — because it is
the exact predicate :meth:`eou_forces_bos` opens the turn on:

  * while it was ``False``, the force-open's own condition was
    ``False`` too, so nothing opened either way;
  * at the first tick it turned ``True``, the force-open fired on
    that same tick whether or not the model proposed a BOS.

So a turn the veto "refused" opened anyway a few ticks later, with
the ONE difference that it opened on the EOU timer's tick instead
of the tick the model's own turn-taking head chose — which made
that head, the thing this checkpoint is trained around, entirely
inert on the duplex path. Measured over 72 h of production
traffic: 689 refusals with speech already confirmed, every one of
them a turn that opened anyway, delayed by a median ~2 s; and all
408 refusals that genuinely had nothing to answer carried
``eou_user_spoke=False``, i.e. the protection §7.27/§7.30 built
this gate for lives entirely in the clause that REMAINS.

Restoring the asymmetry also restores the reference's own shape,
which §7.30.1 recorded and then went beyond: "the reference never
needed a hard block on its LLM's own BOS sampling because it never
built one — the EOU mechanism is purely additive." The
arbi-serve-specific hardening is kept at exactly the scope its own
live evidence (§7.27.1: the model reopening with the user silent
throughout) supports.

How much confirmed silence a MODEL-SAMPLED BOS must stand on —
        design doc §7.60.

        Not zero, and not :attr:`eou_silence_ticks`. Both extremes were
        tried and both are wrong, for reasons this project measured:

          * ``True`` the force-open fires anyway. It
            made the checkpoint's own turn-taking head inert on the duplex
            path and delayed 689 measured turns by a median ~2 s.
          * **Zero, what §7.51 replaced it with.** It leaves nothing at
            all between the model's sampler and an open turn, and the
            model samples BOS *mid-utterance*. Reproduced live with a
            perfectly-paced client (zero synthesized frames): a turn
            opened 16.2 s into a 19.3 s three-sentence question on
            ``eou_silent_ticks=7/40``, answered the fragment it had heard,
            and was killed 3.2 s later by the user still finishing their
            own sentence. §7.51's argument assumed the model's BOS and the
            EOU timer agree about WHEN; mid-utterance they do not.

        The quantity that separates them is measured, not chosen: this
        head's blank runs *inside* genuine continuous speech reach at most
        12 ticks across three files (7 while the audio is continuously
        voiced), and genuine silence is 100% blank. So any threshold above
        ~12 means "the user has actually stopped", while staying below 40
        leaves the model's own head choosing the moment inside that
        window — which is the whole of what §7.51 was protecting.

Capped by :attr:`eou_silence_ticks` so a
        connection that shortens or disables the EOU gate can never end up
        with a STRICTER admission rule than force-open rule.

Whether this tick must OPEN a turn the model did not sample on
        its own — the reference's EOU force-open branch
        (``_apply_turn_taking_from_rnnt_states``, which writes
        ``generated_text_tokens[0, t] = bos_id`` once confirmed speech is
        followed by ``RNNT_EOS_SILENCE_FRAMES`` of confirmed silence while
        the agent is idle). Read by
        :func:`~arbi_serve.runtime.nemotron_voicechat_duplex_step.finish_duplex_tick`
        on every tick whose text channel produced something other than
        BOS.

        This is the ADDITIVE half of the same EOU mechanism
        :meth:`eou_admits_bos` vetoes with, and it is what makes the pair
        complete: the veto alone can only ever remove a turn-open the
        model proposed, so a text channel that never proposes one leaves
        the agent permanently silent after its first turn. The duplex text
        channel is sampled greedily, and the checkpoint's own logits keep
        BOS as a persistent runner-up to PAD rather than an outright
        winner, so "never proposes one" is the ordinary steady state, not
        an edge case.

        **Two connection shapes, two turn-boundary signals — and the same
        additive half applies to both** (design doc §7.44). What differs is
        only WHICH signal says "the user's turn has ended and the agent
        owes an answer":

          * a GENUINE DUPLEX connection has no client to say so, so the
            signal is the reference's own EOU condition — confirmed user
            speech (:attr:`eou_user_spoke`) followed by
            :attr:`eou_silence_ticks` of confirmed silence. The veto's
            first-turn EXEMPTION is deliberately absent here: a connection
            whose agent has never spoken must still earn its open from
            real confirmed speech, matching the reference, whose
            force-open path requires BOU at every turn and only RELAXES
            the frame threshold for turn one. An unprompted greeting
            therefore stays the model's own choice to make; a user who
            speaks first is answered by this like any other turn;
          * a CLIENT-OWNED-TURN (bounded) connection has an EXPLICIT
            signal and needs no inference at all: the client committed an
            utterance and asked for a response (:attr:`speak_gate_closed`
            cleared by ``response.create`` / ``server_vad``'s auto-respond),
            and the model has now heard every frame of it
            (:attr:`frame_seq` reached :attr:`client_turn_audio_target` —
            the SAME rule :meth:`eou_admits_bos` vetoes on, §7.35). That
            IS the turn boundary; there is nothing left to wait for. The
            VAD/EOU rule is not merely unnecessary here, it is unusable:
            a bounded connection is built with ``vad=None``, so
            :meth:`note_vad_events` only ever sees an empty list and
            :attr:`eou_user_spoke` can never become ``True``.

``response.create`` owned those
        boundaries and the model would sample BOS naturally once its
        committed audio was consumed. NOTHING opened the turn. An ordinary ``/v1/realtime`` client could
        not get a single question answered; the turn ran mute until
        ``silent_grace_ticks`` (300 s) expired and the bridge synthesized
        an empty response. Owning a turn's boundaries is not the same as
        being able to make the model speak, and only the client half of
        that was ever true.

        Self-limiting on both paths, by two different latches that each
        fall out of an existing transition rather than being bolted on:

          * duplex — :meth:`mark_agent_active` clears
            :attr:`eou_user_spoke` as the forced turn opens, so the
            condition cannot hold again until the user speaks and falls
            silent afresh;
          * bounded — :meth:`mark_agent_active` sets
            :attr:`client_turn_agent_opened`, cleared only by the next
            :meth:`set_client_turn_audio_target` (i.e. the next
            ``begin_turn``), so exactly one turn is ever forced per
            committed utterance. See that attribute's own docstring for
            the tick this protects.

        The ``eou_silence_ticks <= 0`` escape hatch disables the DUPLEX
        half only, exactly as it disables the corresponding half of
        :meth:`eou_admits_bos`: it is a knob on the VAD/EOU silence rule,
        and the bounded path does not use that rule. A deployment that
        turns the EOU machinery off still gets answers from its
        turn-based clients.

Declare that a CLIENT owns this connection's turn boundaries,
and that the turn now starting ends its real input audio at
``frame_seq_target`` (an absolute :attr:`frame_seq` value) —
design doc §7.35.

Called by
:meth:`~arbi_serve.realtime.duplex_tick_pump.DuplexTickPump.begin_turn`
(and its constructor) for a BOUNDED pump only, off the very same
value that method installs as its own ``num_real_audio_frames``,
on the engine thread. A genuine duplex connection never calls
this, so :attr:`client_turn_audio_target` stays ``None`` there and
:meth:`eou_admits_bos` keeps its VAD/EOU rule bit-identically.

Also re-arms the bounded force-open by clearing
:attr:`client_turn_agent_opened`: every call to this method starts
a NEW client-owned turn (the pump's constructor, then one
``begin_turn`` per committed utterance), and each such turn is
owed exactly one opened response — see :meth:`eou_forces_bos`.

Realtime WS turn source for NemotronLabs-VoiceChat-11B — the SECOND
``turn_streamer`` :class:`~arbi_serve.realtime.session.RealtimeSession` can
be driven by, alongside :func:`arbi_serve.realtime.turn.stream_turn`
(Step-Audio-2's engine-scheduled turn).

:func:`stream_nemotron_voicechat_turn` is a NEW SOURCE feeding the SAME
SINK: it drives
:func:`arbi_serve.runtime.nemotron_voicechat_turn.run_turn`'s synchronous
generator (via ``asyncio.to_thread(next, gen)``, exactly the pattern that
module's own docstring documents for a WS-layer caller) and translates its
:class:`~arbi_serve.runtime.nemotron_voicechat_turn.TextDelta`/``AudioChunk``/``TurnDone``
events into :mod:`arbi_serve.realtime.turn`'s OWN
``TranscriptDelta``/``AudioDelta``/``TurnComplete`` — the exact dataclasses
:meth:`RealtimeSession._run_response` already knows how to turn into wire
events. The ordinary text+audio turn therefore needs NO new event types or
session-layer branches at all — see the module docstring of
``arbi_serve/realtime/turn.py`` for why those dataclasses already carry
everything a turn source needs to say.

Function-calling is the one genuinely new shape (see
:class:`FunctionCallReady`) — NemotronVoiceChat emits tool calls on a
SEPARATE function-logits channel (not interleaved into the text stream the
way a normal ``<tool_call>``-emitting model would), so detecting one can't
reuse ``ChatToolStream``'s text-stream design (see
:mod:`arbi_serve.engine.nemotron_function_channel`'s docstring for the
full reasoning). This generator samples + detokenizes that channel itself
(:meth:`~arbi_serve.runtime.nemotron_voicechat_stt_step.EngineSttStep.decode_function_token`),
feeds an incremental detector, and on a complete
``<TOOLCALL>[...]</TOOLCALL>`` block:

  1. mints a ``call_id`` per named call and yields one
     :class:`FunctionCallReady` per call (the session layer turns each
     into a ``response.function_call_arguments.done`` wire event and
     keeps streaming — no session-layer blocking);
  2. AWAITS every call's result via ``tool_bridge.wait_for_result(call_id)``
     — this is the "pause": the loop simply stops calling
     ``next(gen)`` until every pending future resolves, which happens
     when the client sends ``conversation.item.create``
     (``function_call_output``) and
     :meth:`~arbi_serve.realtime.session.RealtimeSession._on_conversation_item_create`
     calls :meth:`ToolResultBridge.resolve`;
  3. injects a ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>`` chunk (matching
     the reference chat template's own documented client-response
     format — see ``arbi_serve/chat_templates.py``'s
     ``NEMOTRON_VOICECHAT_CHAT_TEMPLATE``) into the running sequence
     via ``EngineSttStep.inject_tokens`` — this is the "resume": the
     next loop iteration calls ``next(gen)`` again, and ``run_turn``
     proceeds exactly as if nothing had paused, now conditioned on the
     injected tool response.

No change to ``run_turn``'s generator CONTRACT was needed to get this
pause/resume shape: because the async caller (this function) is the one
deciding when to call ``next()`` again, "pausing" is simply "don't call
next() yet" and "resuming" is "call it again" — the injected content goes
through the STT step function's OWN mutable state
(:meth:`EngineSttStep.inject_tokens`), not through ``run_turn``'s yield/send
protocol. ``run_turn`` itself is untouched (its own existing tests still
pass unmodified — see the module docstring there for the "why an injected
callback" design this leans on).

=== Barge-in cleanup ===

On cancellation (the session's ``_cancel_active_response`` cancels the
consuming ``asyncio.Task``, exactly as it does for ``stream_turn``), this
async generator's ``finally`` closes the underlying sync generator
(``gen.close()``) so its frame — and therefore its reference to the model
forward's intermediate tensors — is released immediately rather than
waiting for GC. ``run_turn`` itself holds no lock / context manager /
explicit resource (confirmed by reading its full body: no ``with``
blocks, no ``SpeechSlot``-style acquire — unlike ``stream_turn``, which
DOES need an explicit ``finally`` release for exactly that reason).
``EngineSttStep``'s ``MultiStatePool`` is a PRIVATE, turn-scoped pool (see
that class's own docstring) — not the shared engine pool — so it carries
no page-table entry needing an explicit free; abandoning the
``EngineSttStep`` instance (this generator's local variable, dropped when
the generator frame exits) reclaims its CUDA tensors via ordinary
refcounting once ``gen.close()`` runs. A pending ``tool_bridge`` future is
never left dangling either: cancelling this task while awaiting
``wait_for_result`` cancels that await directly, and
:meth:`ToolResultBridge.wait_for_result`'s own ``finally`` pops the
future out of the pending map regardless of how the await ends.

=== Tool-call ack UX ===

The real NIM reference's tool-definition format carries an optional
``ack_messages: list[str]`` per tool: "the model speaks one of these
while waiting for the tool result, keeping the conversation flowing
naturally" (``deploy.md#function-calling``). Two pieces make this real:

  1. **Selection + lookup** (:func:`_tool_ack_messages`): tools stay
     untyped OpenAI-format dicts throughout this layer (matching
     ``schemas.py``'s ``ConversationItemCreate`` convention), so lookup
     is a small ad-hoc dict walk, not a new Pydantic model. No selection
     rule is documented beyond "one of these" — :func:`_maybe_speak_ack`
     picks uniformly at random via :func:`random.choice`.
  2. **Speaking it** (:func:`_synthesize_ack_waveform_frames` /
     :func:`_maybe_speak_ack`): the ack phrase is synthesized DIRECTLY via
     the TTS pieces (:meth:`NemotronVoiceChatModel.prime_tts_turn` /
     ``generate_tts_frame``) at THIS realtime-integration layer, rather
     than by injecting the ack text into ``run_turn``'s own STT-decode
     token stream. Two reasons this is the cleaner integration point,
     not just the more convenient one:

       * This checkpoint's STT->TTS coupling is a bare id passthrough
         (see ``arbi_serve/runtime/nemotron_voicechat_turn.py``'s module
         docstring, "Lockstep coupling") — TTS synthesis for a FULLY
         KNOWN text never needed an STT decode step to begin with, so a
         standalone TTS-only loop over the ack text's own token ids is a
         faithful, not approximate, port of the same mechanism, and
         needs no change to ``run_turn``'s generator contract or its
         yield/send protocol.
       * ``run_turn`` (and ``EngineSttStep``) are the exact two runtime
         modules a concurrent frame-lockstep redesign is actively
         rewriting (see this task's coordination constraints) — routing
         ack playback through their internals would mean building on
         code known to be about to change out from under it. The
         standalone synthesis path touches neither: it opens its OWN
         ``past_key_values``/``prev_code`` thread (a fresh
         :meth:`~NemotronVoiceChatModel.prime_tts_turn` call), fully
         isolated from the paused ``run_turn`` generator's own KV-cache
         state, and reads (never writes) only the shared
         ``embed_subword`` char map that ``run_turn`` already set once
         at turn start.

  The synthesized ``TranscriptDelta``/``AudioDelta`` events are yielded
  from :func:`stream_nemotron_voicechat_turn` between the
  ``FunctionCallReady`` yields and the blocking
  ``tool_bridge.wait_for_result`` await — so they reach the client while
  it is (in parallel, off this loop) executing the tool — and the ack
  text is folded into the running ``transcript_parts`` accumulator so it
  is part of the turn's final transcript, same as any other spoken
  words. Best-effort throughout: no ``ack_messages`` configured, or a
  synthesis failure, silently yields nothing rather than failing the
  tool call.

Session-owned rendezvous for the pause/resume tool-call flow.

One instance per response that supports function calling (created by
:class:`~arbi_serve.realtime.session.RealtimeSession` before starting
the turn, held for that response's lifetime). The turn generator
calls :meth:`wait_for_result` and blocks; the session resolves it from
the client's ``conversation.item.create`` event via :meth:`resolve`.

Two-tier dispatch: does the SERVED model drive this module's turn
loop instead of ``stream_turn``'s ordinary engine-scheduled one?

Mirrors :func:`arbi_serve.multimodal.output.resolve_output_spec` /
:func:`arbi_serve.multimodal.registry.resolve_mm_bindings`: live mode
reads the loaded model's own class; process mode
(``ARBI_ENGINE_PROC``, no model object in the API child) resolves the
same fact weightlessly from the served model dir + the architecture
registry via
:func:`arbi_serve.models.uses_nemotron_voicechat_turn_loop_for_dir`.

Concatenate every system-role message's text content, in order.

``NEMOTRON_VOICECHAT_CHAT_TEMPLATE`` (``arbi_serve/chat_templates.py``)
renders ONLY its ``system_message`` variable (+ tool defs) — it never
loops over ``messages`` at all, since this checkpoint's trained prompt
shape has no textual turn representation for user/assistant content
(the reference tokenizes the system prompt once, then feeds audio
directly — see that template's own docstring). So the caller's
system-role message(s) (``RealtimeSession._turn_messages`` always
prepends one from ``self.instructions``) must be pulled out and
rendered via the ``system_message`` kwarg explicitly; any user/
assistant TEXT content is intentionally dropped here (their AUDIO
content is handled separately, appended as raw placeholder tokens —
see this module's prompt-construction code below).

``output`` (a plain string from ``conversation.item.create``) as a
JSON value fragment for ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>``: passed
through unchanged if it already parses as JSON (the reference's own
convention — a caller MAY JSON-serialize its result), else wrapped as a
JSON string.

``ack_messages`` for the tool named ``name``, from the session's raw
OpenAI-format tool-definition dicts (``RealtimeSession.tools`` — see the
real NIM reference's documented tool-definition format,
``ack_messages: list[str] | None``, and
:data:`arbi_serve.chat_templates.NEMOTRON_VOICECHAT_CHAT_TEMPLATE`,
which already strips this same field out of the LLM-visible tool JSON
and re-renders it into the prompt's own ``<TOOL_ACK_MESSAGES>`` block).

Tools stay untyped ``dict``s throughout this codebase's realtime layer
(see ``schemas.py``'s ``ConversationItemCreate`` docstring on why
``conversation.item.create``'s ``item`` is a plain dict too) — this is
the one lookup function that needs to reach into that shape, so the
ad-hoc parsing lives here rather than behind a new Pydantic model that
every other tool-consuming call site would then need to route through
for no benefit (the chat template above already does its own
equivalent ad-hoc ``tool.function if tool.function is defined else
tool`` unwrap for the same untyped-dict reason).

Standalone TTS-only synthesis for a FIXED, fully-known text (an
``ack_messages`` phrase) — the mechanism behind the "ack while waiting"
UX (module docstring's "Tool-call ack UX" section has the full design
rationale). Runs entirely on this thread; the caller wraps it in
``asyncio.to_thread``.

Why this can skip ``run_turn``'s STT-decode lockstep entirely: per
``arbi_serve/runtime/nemotron_voicechat_turn.py``'s module docstring
("Lockstep coupling"), this checkpoint's STT->TTS coupling is a bare id
passthrough — the STT backbone's OWN sampled token id is fed straight
into the TTS side as ``subword_ids``, no re-tokenize round trip. An ack
phrase's text is already fully known up front (it is not being
generated token-by-token), so ``token_ids`` (this text's OWN
tokenizer-encoded ids, standing in for what would otherwise be
per-step STT samples) can drive the exact same per-frame
``generate_tts_frame`` call directly, with no STT step at all.

Isolation from the PAUSED main turn: this builds its OWN
``past_key_values``/``prev_code`` thread via a fresh
:meth:`~NemotronVoiceChatModel.prime_tts_turn` call — entirely separate
Gemma3 KV-cache state from the one ``run_turn``'s generator is holding
(as a local variable in ITS OWN frame, inaccessible from here) while
parked on ``await tool_bridge.wait_for_result(...)``. The only model
state this reads (never writes) is
``model.tts_backbone.embed_subword``'s char map, which ``run_turn``
already set once at turn start, before any tool call could possibly be
mid-turn-detected — so there is nothing left to set up here. No
``arbi_serve/runtime/*`` code is imported or called by this function.

Ack-while-waiting filler speech: "the model speaks one of [the
called tool's ``ack_messages``] while waiting for the tool result,
keeping the conversation flowing naturally" (real NIM reference,
``deploy.md#function-calling``). See module docstring's "Tool-call ack
UX" section for why this synthesizes directly via the TTS pieces at
THIS (realtime WS integration) layer, rather than injecting the ack
text into ``run_turn``'s own STT-decode stream.

Selection: the reference names no rule beyond "one of these" — this
picks uniformly at random (:func:`random.choice`) among the first
called tool (in call order) that defines any ``ack_messages`` at all;
a tool-call batch where no called tool defines one yields nothing
(silent no-op, matching the field's optional/best-effort nature).
Synthesis failure is likewise best-effort: logged and swallowed, never
raised into the turn — a missing filler phrase must not fail the tool
call itself.

Submit one NemotronVoiceChat turn and yield transcript + audio (+
function-call) events. See module docstring for the full design.

``tool_bridge``: passed by :class:`~arbi_serve.realtime.session.RealtimeSession`
only when it detects this streamer supports function calling (see
``supports_tool_calls`` set on this function below — the same
attribute-flag dispatch pattern the session already uses to keep
``stream_turn``'s signature untouched). ``None`` means "no client-side
tool bridge for this response": a detected tool call is logged and
the turn continues without pausing (best-effort — never hangs a turn
that happens to emit ``<TOOLCALL>`` when the client never registered
for function calling).

``tools``: OpenAI-format tool definitions from
:attr:`~arbi_serve.realtime.session.RealtimeSession.tools` (set via
``session.update``), rendered into the ``<AVAILABLE_TOOLS>`` block by
:data:`arbi_serve.chat_templates.NEMOTRON_VOICECHAT_CHAT_TEMPLATE`.
``None``/empty renders no tool block — the model then never emits
``<TOOLCALL>``, matching a plain (no-tools) turn.

Create (or return) the pending future for ``call_id``, without
waiting on it.

The turn-based path never needs this — it awaits
:meth:`wait_for_result` immediately after yielding the call, on
the same coroutine. The duplex path does: its call is detected on
the ENGINE thread and surfaced through a queue, so the WS side
must be able to make the call resolvable BEFORE it sends
``response.function_call_arguments.done`` and hands the loop back
(design doc §7.9.16). Without that, a client answering inside one
event-loop iteration would race the ``create_task`` that is about
to await, and :meth:`resolve` would report an unknown ``call_id``.

Idempotent, and :meth:`wait_for_result` goes through it, so the
two entry points can never end up holding different futures for
one call. An ALREADY-RESOLVED future is returned as-is rather than
replaced — that is precisely the race this method exists for (the
client answered between registration and the await), and minting a
fresh one would throw the answer away and hang the call until its
timeout. Nothing goes stale: ``wait_for_result``'s ``finally``
pops the entry however the await ends, so a resolved future can
only ever be one nobody has read yet.

Await the client's result for ``call_id``, up to
:attr:`_timeout_s`.

On firing, this raises :class:`ToolCallTimeoutError`
(never hangs forever) and, if this bridge was constructed with a
``session``, calls ``session.mark_agent_idle()`` and clears
``session.fc_state`` — the reference production backend's own
broader "force idle on every FC-cycle-end, success or failure"
choice (design doc §3.2.1's "Flagged divergence" recommendation).

The duplex FC-pause PAD while a call is
pending, so the frame grid keeps advancing during this wait. That
makes the ``fc_state = None`` reset above the ESCAPE from a muted
stream rather than mere bookkeeping — clearing the state reopens
the gate (see
:attr:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession.fc_in_progress`, which derives
from it), so a client that never answers gets an agent that
starts speaking again instead of one that is silent forever.

Acoustic prosody / intonation markers for transcript segments.

Text alone loses *how* something was said — the rising pitch of a question, the
drop of a flat denial, a raised voice under cross-examination. Those live in the
audio, not the tokens. This module derives a compact, **deterministic** prosody
summary per utterance straight from the waveform (no model, no extra latency):

  * ``pitch_hz``  — median fundamental frequency over voiced frames (autocorrelation).
  * ``trend``     — ``"rising"`` / ``"falling"`` / ``"flat"`` from the f0 slope across
                    the utterance (rising often ⇒ a question; falling ⇒ finality).
  * ``energy``    — ``"loud"`` / ``"normal"`` / ``"soft"`` from RMS vs a speaking range.
  * ``marker``    — one glyph summarizing the above for a caption (▲ rising, ▼ falling,
                    ‼ loud, · soft, — flat).

This is *intonation*, not semantic sentiment: it reports measurable acoustics
honestly and never claims to read emotion. (A model pass could add that; this
path stays cheap so the live transcript is never slowed.)

Everything is wrapped so a failure returns ``None`` — prosody is a garnish, never
allowed to break transcription.

Typed wire schemas for the ``/v1/realtime`` WebSocket.

OpenAPI 3.x cannot describe a WebSocket, so the socket's message shapes are
invisible to the generated client. We follow the same trick the ARBI repo
uses for its ``/ws`` endpoint: model every message as a Pydantic model with a
``type`` ``Literal`` discriminator, then expose them all through a dummy HTTP
``GET /v1/realtime/schemas`` (see ``routes/openai_realtime.py``) whose
``response_model`` drags every model into ``components.schemas``. The TS/Python
clients then autogenerate types for what to *send* and what to *expect*.

These models are the SINGLE SOURCE OF TRUTH for the wire: ``events.py`` builds
each ``server -> client`` event by constructing the model here and calling
``.model_dump()``, so the published schema and the actual bytes cannot drift.

Only the subset of the OpenAI Realtime beta this server implements is modelled.
Composite sub-objects allow extra keys (``extra="allow"``) so the emit path
stays tolerant, and optional fields default to ``None`` — emitted as ``null`` —
matching OpenAI's convention that unpopulated fields are present-but-null.

The ``part`` object carried by ``response.content_part.{added,done}``
— the real NIM reference's GA-vocabulary content-part marker (see
:class:`ResponseContentPartAdded`). Minimal by design: the reference's
own examples show only ``{"type": "audio"}``, and ``_Wire``'s
``extra="allow"`` keeps room for a future ``text``/``transcript`` part
kind without a schema change here.

Signals a complete function call detected on NemotronVoiceChat's
function-logits channel. Field shape matches the real NIM reference's
documented wire format exactly (``call_id``/``name``/``arguments``, plus
the response-lifecycle ``response_id``/``item_id``/``output_index`` this
server's other ``response.*`` events already carry).

Client sends a function-call result back to the server
(NemotronVoiceChat function calling — see the real NIM reference's
``conversation.item.create`` / ``function_call_output`` shape). ``item``
is left as a plain dict (rather than a typed sub-model) since its only
currently-supported shape is
``{"type": "function_call_output", "call_id": ..., "output": ...}`` and
:meth:`arbi_serve.realtime.session.RealtimeSession._on_conversation_item_create`
validates those fields itself, the same way ``ResponseCreate``/``ResponseCancel``
carry no typed body beyond ``type``/``event_id``.

Container that pins every ``/v1/realtime`` message into the OpenAPI spec.

Returned (all-``None``) by ``GET /v1/realtime/schemas``; the body is
meaningless — only the referenced ``components.schemas`` matter, so the
client learns what to send and what format to expect.

Realtime session state machine (VAD-orchestrated, turn-based).

One :class:`RealtimeSession` per WebSocket connection. It owns the
session config, the input-audio buffer, the VAD (for ``server_vad`` turn
detection), the conversation history, and the lifecycle of the single
active response — and translates client events into the engine turn +
server events.

Decoupled from FastAPI: the constructor takes a ``send`` coroutine
(``event_dict -> awaitable``) and a ``turn_streamer`` async-iterator
factory. Production wires ``websocket.send_json`` +
:func:`arbi_serve.realtime.turn.stream_turn`; tests inject a collector +
a canned streamer, so the whole protocol (event ordering, VAD auto-commit,
barge-in cancel) is exercised on CPU with no engine.

Barge-in is turn-based, NOT full-duplex: the model does not attend to
overlapping audio. When VAD detects fresh user speech during an active
response, the session CANCELS that response (``eng.cancel`` via the turn
generator's ``finally``) and opens a new turn.

=== Module layout (mixins, this module re-exports the public API) ===

``RealtimeSession`` grew several genuinely distinct concerns bolted on
over many tasks; each now lives in its own mixin module and is combined
here by inheritance (never delegation — every ``self.`` attribute a mixin
method touches is still set on THIS class's own ``__init__``, so there is
no indirection cost and no need for delegating wrappers):

  - THIS module — the original turn-based WS dispatch: session config,
    the client-event dispatch table, the ``_on_*`` event handlers, commit/
    conversation-item plumbing, and connection teardown (``aclose``).
  - :mod:`arbi_serve.realtime._session_duplex_optin` —
    :class:`~arbi_serve.realtime._session_duplex_optin._DuplexOptInMixin`,
    the full duplex opt-in (design doc §5's THIRD tier — see
    that module's own docstring).
  - :mod:`arbi_serve.realtime._session_bounded_duplex` —
    :class:`~arbi_serve.realtime._session_bounded_duplex._BoundedDuplexMixin`,
    the bounded duplex backend (design doc §7.24) plus duplex
    function calling (shared with full duplex mode — see that
    module's own docstring).
  - :mod:`arbi_serve.realtime._session_response` —
    :class:`~arbi_serve.realtime._session_response._ResponseMixin`, the
    turn-based response lifecycle, speculative endpointing, and live
    voice cloning (``CLONE_VOICE``, re-exported here).
  - :mod:`arbi_serve.realtime._session_transcription` —
    :class:`~arbi_serve.realtime._session_transcription._TranscriptionMixin`,
    the off-hot-path input-audio transcription echo
    (``_asr_system_prompt``, re-exported here).

See each mixin module's own docstring for the design rationale of its
concern; this docstring covers only the turn-based dispatch backbone that
stays in this file.

Validate a ``session.update`` ``tool_choice`` into the engine's shape.

Accepts the same values ``/v1/chat/completions`` does — ``"auto"`` /
``"none"`` / ``"required"``, or ``{"type": "function", "function":
{"name": ...}}`` — and returns the value
:attr:`arbi_serve.engine.request.SamplingParams.tool_choice` expects.
``None`` resets the session to :data:`_DEFAULT_TOOL_CHOICE`.

Raises :class:`ValueError` on anything else so the caller can reject it
at config time with a param-tagged error, rather than silently serving a
connection under a policy the client did not ask for.

This session's resolved ``turn_detection`` knobs, for a detector
that will be fed PCM at ``sample_rate``, plus the acoustic backend
to build it with.

``sample_rate`` is a PARAMETER rather than always
:attr:`input_sample_rate` because the two detectors this session
builds are fed at two different rates: the turn-based detector
reads the client's own PCM off ``input_audio_buffer.append``,
while a genuine duplex connection's detector reads the frame
queue, whose PCM ``DuplexAudioIn`` has already resampled to the
served model's input rate. A detector told the wrong rate
resamples by the wrong factor and sees time-compressed,
formant-shifted audio — see :meth:`_rebuild_duplex_vad`.

Build (or rebuild) the detector a GENUINE duplex connection's
barge-in AND-gate and EOU admission tracker read — at the rate of
the audio it will actually be fed.

A duplex connection's detector is not ``self._vad`` and cannot be.
The turn-based path feeds the VAD the client's own PCM straight
off ``input_audio_buffer.append``, so
``VadConfig.sample_rate = input_sample_rate`` is right there. The
duplex path does not: ``DuplexAudioIn.push_pcm16`` resamples every
chunk to the served model's ``duplex_input_sample_rate`` before it
ever reaches the frame queue, and the frame queue is what feeds
this detector. Handing the

Whether a duplex connection opened NOW should feed its EOU
admission tracker from the served model's own endpoint detector
(design doc §7.32) rather than the acoustic VAD.

Resolution mirrors :func:`~arbi_serve.realtime.vad.build_vad`'s
own explicit-vs-default contract: an EXPLICIT
``turn_detection.backend: "semantic"`` requires the served model
to carry the detector (``rnnt_head``) and raises a clear error
when it does not (no silent downgrade); an explicit ACOUSTIC
backend (``silero``/``ten``/``energy``) means the client chose
that signal, so the semantic feed stays off; the DEFAULT
(``auto``/unset) uses the semantic feed whenever the model has
it — the strictly better signal costs nothing extra, since the
detector rides the same encoder forward the audio-in path
already runs.

Confirm-window (ms) for speculative endpointing; 0 disables it.

Only meaningful under server_vad: it trades a shorter VAD
``silence_duration_ms`` (fast provisional endpoint) for a hold window
during which a resumed turn is cancelled inaudibly.

Deliberately ``0`` whenever the bounded duplex backend is in use
(design doc §7.24's module-docstring "deliberately deferred"
section): the hold/replay mechanism is built on
:meth:`_run_response`'s own event buffering (:meth:`_emit_turn`),
which the bounded backend's direct
``pump_duplex_connection``-to-socket streaming has no equivalent
of. Falls back to ordinary immediate commit+respond — exactly
what a client that never set the knob already gets.

The temperature this session's TEXT-channel sampling should use,
for ANY backend (turn-based ``_run_response``, duplex, and
bounded-duplex all call this — the name predates the turn-based
call site, kept as-is to avoid a churny rename across 5 files +
``tests/test_duplex_text_temperature.py``).

Resolution order — most explicit wins, reference default last:

1. a client ``session.update {"temperature": ...}``
   (``self._temperature_client_set``);
2. an operator-explicit ``ARBI_REALTIME_AUDIO_TEMPERATURE`` (env or
   live override — :func:`~arbi_serve.runtime_flags.flag_explicitly_set`);
3. for a NemotronVoiceChat engine, the checkpoint's own reference
   production default (:data:`_NEMOTRON_VOICECHAT_DUPLEX_TEMPERATURE`,
   greedy — see that constant's docstring for the citations);
4. otherwise the cross-model flag value already in
   ``self.temperature``.

Client delivers a function-call result — the resume half of the
NemotronVoiceChat tool-call pause/resume flow (see
``arbi_serve.realtime.nemotron_voicechat_turn``'s module docstring).
Matches the real NIM reference's wire shape:
``{"type": "conversation.item.create", "item": {"type":
"function_call_output", "call_id": ..., "output": ...}}``.

Pipecat Smart Turn v3 — semantic end-of-turn classification on raw audio.

Smart Turn v3 (``pipecat-ai/smart-turn-v3``, BSD-2) answers a different
question from a speech/silence VAD. A VAD asks "is there speech in this
16 ms slice"; Smart Turn asks "given the last few seconds of waveform, is
the speaker DONE" — it reads the prosody of the pause (a trailing-off
cadence versus a mid-thought hesitation), so a thinking pause and a
finished sentence are separable even though both are silence.

The contract, taken from the published reference implementation
(``pipecat-ai/smart-turn``'s ``inference.py`` and pipecat's own
``LocalSmartTurnAnalyzerV3``):

  * **Input** — one fixed 8 s window of 16 kHz mono audio, converted to
    Whisper log-mel features: ``(80, 800)`` float32 under the ONNX input
    name ``input_features``. Audio longer than 8 s keeps its LAST 8 s;
    shorter audio is zero-padded ON THE LEFT, so the window is always
    right-aligned on the moment being judged.
  * **Waveform normalization** — Whisper's ``do_normalize`` zero-mean /
    unit-variance pass runs over the padded 8 s window, not over the raw
    segment.
  * **Output** — a single sigmoid probability that the turn is COMPLETE.
  * **Invocation** — once per candidate pause, not per audio hop. The
    reference pairs it with a cheap frame VAD that finds the pauses; this
    model is the second opinion on each one.

:func:`whisper_log_mel` is the feature front end, implemented on numpy
against Whisper's own definition (hann-windowed 400-point STFT at a
160-sample hop, 80 slaney-scale slaney-normalized mel filters over
0-8 kHz, log10, an 8-decade floor relative to the window max, then the
``(x + 4) / 4`` scaling). ``tests/test_smart_turn.py`` pins it against
``transformers.WhisperFeatureExtractor``, the upstream spec.

Weights are a single ONNX file, resolved by :func:`resolve_model_path`
from ``~/.cache/arbi-serve/smart-turn/`` and fetched from the Hugging
Face repo on first use. Prefetch it for an air-gapped or
cold-start-sensitive deploy with::

    python -m arbi_serve.realtime.smart_turn

Return a local path to the Smart Turn ONNX file, fetching it if needed.

An explicit ``path`` is used verbatim (and must exist). Otherwise the
file is looked up in :data:`CACHE_DIR` and, when absent and
``download`` is set, fetched from :data:`HF_REPO` into that cache.
Raises ``RuntimeError`` naming the exact file and URL when the weights
cannot be produced — a selected backend never degrades silently.

The ONNX session plus its feature front end.

``probability`` is the whole surface: 16 kHz float32 audio in (any
length — it is right-aligned onto the 8 s window), completion
probability out. The underlying ``onnxruntime`` session is safe to
call from several threads, so :func:`load` shares one instance across
every connection rather than paying a session per WebSocket.

Continuous, hours-safe streaming transcription / translation.

Step-Audio-2-mini is a *generative* model with a bounded audio+text context,
not a streaming CTC recognizer: you cannot feed it an hour and read tokens out
the side. The hours-safe pattern is **VAD-segmented streaming** — a
voice-activity detector cuts the continuous feed at natural pauses, and each
utterance becomes ONE bounded ASR call (well inside context). This is the same
shape as Whisper-streaming and every live-caption product.

Why it is hours-safe (the property that matters):

  * Memory is ``O(one utterance)``, not ``O(stream)``. Audio is retained only
    from the current utterance's start; once a segment is transcribed the bytes
    are dropped. A one-hour session holds at most one utterance of PCM at a
    time — exactly like the audio-encoder micro-batch bound, not a growing
    buffer that eventually OOMs.
  * A speaker who never pauses cannot defeat that: :data:`_MAX_SEGMENT_MS`
    force-cuts an in-progress utterance at a fixed ceiling (Whisper's 30 s
    window), emits it, and continues — so the retained buffer is bounded even
    with zero silence.

Timestamps are **segment-level and absolute**: each segment carries
``[start_ms, end_ms]`` measured from the start of the whole stream (the VAD is
never reset, so its offsets grow monotonically for the session's lifetime).
They are VAD-boundary timestamps, NOT model-emitted word timings — Step-Audio-2
does not produce those. See :class:`TranscriptSegment`.

Transport-independent by design: this core takes raw PCM in via :meth:`feed`
and yields :class:`TranscriptSegment`\s out. The realtime WebSocket session and
an HTTP/SSE endpoint both drive the SAME core — one windowing loop, two mouths.

Resource-exhaustion errors that must END the stream, not be swallowed.

An ordinary per-utterance failure (a decode glitch, a validation error) is
best-effort: skip that utterance and keep transcribing — one bad segment
must not kill an hours-long stream. But an OOM leaves the process in an
unreliable state, so it propagates to the transport, which ends the SSE with
an honest error frame. Mirrors ``_request_ctx._oom_types`` without importing
a route helper into the realtime core (layering).

One finalized transcript/translation segment.

``start_ms`` / ``end_ms`` are absolute offsets from the start of the stream,
at VAD-boundary granularity (NOT model word timings — Step-Audio-2 emits
none). ``language`` is the model's detected source language (from its
``<英语>``/``<中文>`` tag), or ``None`` if it emitted no tag. ``text`` is the
verbatim transcript, or the translation when the transcriber is in translate
mode. Non-speech segments are dropped and never yielded.

``is_partial`` marks an *interim* hypothesis: the growing text of an utterance
the model is still decoding (the live-caption reveal). A partial is always
followed by more partials or exactly one final (``is_partial=False``) segment
with the same ``start_ms``; a consumer replaces the current in-progress line
on each partial and commits it on the final. Partials are emitted only when
the transcriber is constructed with ``emit_partials=True``.

Drive continuous ASR/translation over an unbounded PCM16 stream.

Construct once per stream, then ``async for seg in st.feed(pcm): ...`` for
each incoming chunk and finally drain ``st.flush()`` at end-of-stream. Each
yields zero or more :class:`TranscriptSegment`\s as utterances finalize.

``translate_to`` is validated up front (raises
:class:`~arbi_serve.multimodal.asr_text.UnsupportedTranslationTarget` for a
target the model cannot serve) so a stream never starts only to silently
return the source language.

First word of ``raw``, lowercased and clipped; ``None`` if unusable.

The model usually answers with a single word (``"neutral"``) but can add
punctuation or a stray tag; keep only a leading alphabetic run so the field
is always a clean label or nothing.

Short verbatim ASR probe used only to gate translate mode: True unless
the model flags the segment as non-speech (``<非语音>``) or empty. Bounded
to a few tokens (we only need the leading tag / first word). Fails OPEN
(returns True) on any error so a probe glitch never drops real speech.

Run ONE bounded ASR/S2TT call over a single utterance's PCM.

Yields ``(text, language, is_final)``. When ``emit_partials`` is set, each
model delta yields the accumulated hypothesis so far with
``is_final=False`` (the live-caption reveal); otherwise only the single
final tuple is yielded. Exactly one ``is_final=True`` tuple always ends
the utterance. A per-utterance failure is swallowed (yields an empty
final so the stream continues); an OOM propagates to end the stream.

Second, paralinguistic pass: one lowercase English emotion word.

A separate bounded call over the SAME utterance audio with
:data:`_EMOTION_SYSTEM_PROMPT`. Best-effort: any failure (or an empty /
non-word answer) returns ``None`` so the transcript stream is never held
up or broken by the emotion probe. StepFun recommends a non-greedy
temperature for paralinguistic understanding, so this runs at 0.7 (the
transcript pass stays greedy at 0.0).

Pin the ambient CUDA device on a thread the engine did not start.

The realtime layer offloads its blocking model work with
``asyncio.to_thread`` — the audio-in encode
(:mod:`arbi_serve.realtime.duplex_audio_in`) and the turn-based TTS
drive (:mod:`arbi_serve.realtime.nemotron_voicechat_turn`) both do. Those
run on the event loop's default executor, and **the only thread that ever
calls** ``torch.cuda.set_device(eng.device)`` **is the engine thread**
(:mod:`arbi_serve.engine.build_phases_load`). A thread that never made
that call has an ambient current device of cuda:0.

Tensors are still built on the right device, so nothing looks wrong: the
kernels simply launch on GPU 0 holding pointers into GPU 1's address
space. The driver reports it as an Xid 31 ``MMU Fault ... FAULT_PDE
ACCESS_TYPE_VIRT_READ`` against the *ambient* GPU, the CUDA context is
poisoned, and the ENGINE thread dies at its next sync point with a
traceback naming whatever it happened to be compiling — never the
offloaded call that actually did it.

A single-GPU deployment cannot see any of this, because ambient cuda:0 IS
``eng.device``. It needs a host with more than one visible GPU and a
non-zero ``--device``, which is why it reached production. Design doc
§7.58.10 has the four-arm isolation that pinned it.

Make ``device`` the ambient CUDA device for the calling thread.

Returns a no-op for anything that is not a CUDA device, so a CPU test
and a CPU-only build run exactly the same code path as a GPU serve.

``torch`` is passed in rather than imported at module scope: several
callers keep ``torch`` out of their import-time dependencies and hold
it as a local import already.

Deliberately NOT memoized. ``torch.cuda.device`` records the device it
displaced on the INSTANCE (``prev_idx``) when entered, and
``asyncio.to_thread`` hands the same pooled executor thread to
different connections, so one shared guard entered twice concurrently
would restore the wrong device on exit. A fresh guard per call costs
nothing next to the model forward it wraps.

Engine-driven Realtime turn generator (the speech-out streaming core).

:func:`stream_turn` is the production ``turn_streamer`` the
:class:`~arbi_serve.realtime.session.RealtimeSession` drives: given the
rendered conversation (system + history + the committed user audio turn),
it submits ONE speech-output request through ``eng.asubmit`` and yields
:class:`TurnEvent`\s as the engine emits — interleaved transcript text
and streamed 24 kHz PCM16 audio, chunk-by-chunk, exactly like the SSE
chat path (:func:`arbi_serve.server.routes.openai_chat._chat_stream`)
but shaped for the WebSocket session instead of SSE frames.

Cancellation (barge-in / ``response.cancel``): the consuming task is
cancelled; the ``finally`` here cancels BY REQUEST ID
(``eng.cancel_by_id`` — P3 CancelMsg semantics) so the GPU work stops
within a token. A submit rejection (``SubmitRejected``, a ValueError)
propagates to the session task like any pre-P3 validation error. This module imports torch / the engine
multimodal stack only at call time so the session + route layers stay
importable on CPU/CI.

Render + tokenize + expand audio placeholders (+ optionally seed speech).

Mirrors the chat route: render the chat template, append the model's
``<tts_start>`` speech seed (only when ``seed_speech`` — omitted for a
TEXT-only turn such as ASR / translation), encode, then expand the audio
placeholder tokens against the preprocessed features.

Submit one speech-output turn and yield transcript + audio events.

``voice`` is a prepared voice name OR an ephemeral conditioning triple
(the live-voice-clone path: the incoming speaker's own
``{prompt_speech_tokens, spk_emb, prompt_mels}``) — passed straight through
to ``decoder.open_stream``.

Submit one TEXT-ONLY turn over input audio — the ASR / speech-translation
echo. No ``<tts_start>`` speech seed (so the model emits text, not codes) and
``priority="batch"`` so it runs opportunistically and never steals scheduling
from the interactive response. Yields :class:`TranscriptDelta`s + a final
:class:`TurnComplete`; no audio.

Server-side voice-activity detection for ``server_vad`` turn detection.

Step-Audio-2-mini ships no VAD on disk, so the ``server_vad`` turn
detector needs a server-side detector to decide when a user turn starts
and ends (and therefore when to auto-commit + when a barge-in fires).

Backends (pluggable behind :class:`VadDetector`, selected by
``turn_detection.backend`` / :func:`build_vad`'s ``prefer``):

  * :class:`SileroVad` — the **default and recommended backend**
    open-source VAD). Ships as a dependency of the audio/omni serving
    extra.
  * :class:`EnergyVad` — the **zero-dependency fallback**. numpy RMS + a
    silence hangover; always available, and what the default gracefully
    degrades to if ``silero-vad`` is somehow missing. Its ``threshold`` is
    an RMS gate (see :class:`VadConfig`), not a neural probability. Fine in
    quiet environments; poor at rejecting background noise.
  * :class:`TenVad` — an **opt-in backend** (TEN-framework ``ten-vad``).
    NOT a dependency; ``pip install ten-vad`` to try it. 16 kHz, 256-sample
    (16 ms) hop.
  * :class:`SmartTurnVad` — an **opt-in SEMANTIC backend** (pipecat's
    Smart Turn v3). Two-stage by construction, because that is the shape
    of the model: a cheap frame VAD finds the candidate pauses and Smart
    Turn v3 judges each one on the trailing waveform, so a mid-thought
    hesitation does not end the turn. Needs ``onnxruntime`` plus an ONNX
    weight file (fetched + cached on first use, see
    :mod:`arbi_serve.realtime.smart_turn`).

The default (server_vad with no explicit ``backend``, or ``prefer="auto"``)
resolves to **silero**; if ``silero-vad`` can't be imported it logs a
WARNING and falls back to energy so a bare / misconfigured install still
boots. An EXPLICIT ``backend: "silero"`` / ``"ten"`` / ``"smart_turn"``
with its package missing raises a clear pip-hint error instead (explicit
= hard fail). ``ten-vad`` and ``smart-turn`` weights are never
dependencies.

Every backend except :class:`SmartTurnVad` is a per-frame speech gate: it
classifies each fixed hop voiced/unvoiced and runs the shared
:class:`_FrameStateMachine` hysteresis over that. :class:`SmartTurnVad`
does not fit that shape and is deliberately not forced into it — it wraps
one of the others and vetoes its ``speech_stopped`` edges. See that
class's docstring.

Detector contract: ``feed(pcm_bytes) -> list[VadEvent]``. The caller
streams appended PCM16 chunks in; the detector returns zero or more
``("speech_started" | "speech_stopped", offset_ms)`` transitions, where
``offset_ms`` is measured from the start of the buffer the session has
fed since the last ``reset``. State is per-connection; ``reset`` clears
it at the start of each turn.

Resolved ``turn_detection`` knobs.

``threshold`` semantics depend on the backend:

  * **energy** — an RMS gate in ``[0, 1]`` on the normalized waveform
    (voiced speech is typically 0.02-0.2). This is why the energy
    default uses 0.02, NOT the OpenAI ``server_vad`` default of 0.5.
  * **ten / silero** — the neural speech PROBABILITY in ``[0, 1]``
    (native semantics; OpenAI's 0.5 default applies directly).

``prefix_padding_ms`` is subtracted from the reported speech-start
offset; ``silence_duration_ms`` is the trailing-silence hangover that
ends a turn (the barge-in / auto-commit trigger). Under
:class:`SmartTurnVad` that hangover becomes the *candidate* pause
length — the point at which the semantic model is asked whether the
turn is really over.

The ``smart_turn_*`` knobs are inert for every other backend.

Shared per-frame speech-start/stop hysteresis for every backend.

A backend classifies each fixed-duration frame voiced/unvoiced and
calls :meth:`_step`; the machine emits ``speech_started`` after
``start_frames`` voiced frames and ``speech_stopped`` after
``silence_duration_ms`` of continuous unvoiced audio following speech.
``_frame_dur_ms`` (set by the backend) is how much wall-clock each
frame advances — 20 ms for energy/silero framing, 16 ms for TEN's
256-sample hop.

RMS-energy VAD with a silence hangover — no external dependency.

Buffers incoming PCM16 into fixed ``frame_ms`` frames, marks each
frame voiced by RMS vs ``threshold`` (an RMS gate, see
:class:`VadConfig`), and runs the shared start/stop hysteresis.

Shared plumbing for the neural backends.

Both silero and TEN want 16 kHz mono and a FIXED per-inference window
(silero v5: 512 samples/32 ms; TEN: 256 samples/16 ms). Incoming PCM
at ``cfg.sample_rate`` is resampled to 16 kHz, buffered, and processed
in exact ``_hop``-sample windows here — so callers keep feeding
arbitrary-length chunks and each backend still gets its required window
size. Subclasses supply the model (``_load``) and the per-window
probability (``_prob``); the shared :class:`_FrameStateMachine`
hysteresis turns the gated probability into start/stop events.

TEN VAD backend — an opt-in alternative to :class:`SileroVad`.

Wraps TEN-framework's ``ten-vad`` (permissive HF weights). TEN runs at
16 kHz over a fixed **256-sample (16 ms) hop**, returning a per-hop
speech probability we gate with ``cfg.threshold`` (the neural
probability, default ~0.5). :class:`SileroVad` remains the default
recommended neural backend. Lazy import: a missing ``ten-vad``
package raises a clear pip-hint error only when selected.

silero-vad backend — the proven neural standard (opt-in).

silero v5 is the de-facto open-source VAD (small, fast on CPU, well
validated), so this is the **safe recommended neural backend** when an
operator wants better precision than the energy gate. It runs at
16 kHz over its required fixed **512-sample (32 ms) window**, treating
``cfg.threshold`` as the silero speech PROBABILITY (native semantics).
Lazy import: a missing ``silero-vad`` package raises a clear pip-hint
error only when this backend is selected.

Semantic end-of-turn detection — pipecat's Smart Turn v3, composed.

Smart Turn v3 is not a frame classifier and is deliberately not
wrapped as one. It reads a trailing window of waveform and answers
"is this speaker DONE" — a question that only has meaning at a pause,
and only against seconds of context, so :class:`_NeuralHopVad`'s
per-hop ``_prob`` shape does not apply to it. What it composes with is
exactly what this class builds: a cheap frame VAD
(``smart_turn_first_pass``, silero by default) proposes candidate
pauses, and Smart Turn v3 confirms or vetoes each one.

The composition, in the caller's terms:

  * ``speech_started`` passes straight through from the first pass.
  * the first pass's ``speech_stopped`` (after ``silence_duration_ms``
    of silence) is a CANDIDATE. The model scores the audio from the
    padded speech start through now; at or above
    ``smart_turn_threshold`` the stop is emitted, below it the stop is
    SWALLOWED and the turn stays open, so a hesitation mid-sentence
    does not commit a half-finished utterance.
  * a swallowed stop is bounded: once continuous silence measured from
    the start of that pause reaches ``smart_turn_max_silence_ms``, the
    turn ends anyway. A refused pause can never hold a turn open
    forever.
  * while a turn is held open, a resumed-speech ``speech_started`` is
    swallowed too — the caller was never told the turn ended, so it is
    the same turn continuing, not a new one.

Cost is one inference per candidate pause, not per hop, which is what
makes an 8 s-window model affordable on the audio tick path at all.

Construct a detector for the requested backend.

``prefer``:

  * ``"auto"`` / unknown (the DEFAULT, i.e. server_vad with no explicit
    ``backend``) -> :class:`SileroVad`. This is a *default*, not
    an explicit request, so if ``silero-vad`` is not installed we log a
    clear WARNING and **fall back to** :class:`EnergyVad` rather than
    failing to boot.
  * ``"silero"`` -> :class:`SileroVad`, the recommended neural backend;
    EXPLICIT request, so a missing package raises the clear
    ``pip install silero-vad`` error (no silent downgrade).
  * ``"ten"`` -> :class:`TenVad`, an opt-in neural backend; EXPLICIT,
    raises a clear ``pip install ten-vad`` error if absent.
  * ``"smart_turn"`` -> :class:`SmartTurnVad`, the opt-in SEMANTIC
    backend; EXPLICIT, raises a clear error naming
    ``pip install onnxruntime`` or the weight-prefetch command if
    either is absent. Its own first-pass VAD is resolved through this
    same function (``cfg.smart_turn_first_pass``, ``"auto"`` by
    default), so the frame stage degrades on the auto contract while
    the semantic stage the caller actually asked for never does.
  * ``"energy"`` -> :class:`EnergyVad`, the zero-dependency fallback.

Contract: an EXPLICIT neural request fails hard when its package is
missing (the caller wanted that detector); only the DEFAULT degrades
gracefully so a bare / misconfigured install still serves.

Per-model recipe registry.

A *recipe* is a YAML file declaring the known-good calibration knobs
and serving defaults for a model: cal corpus, target bit widths,
eval set, default max-context, supported weight quants. The CLI
looks up a recipe by:

  1. Reading ``config.json`` from the ``--model`` directory.
  2. Matching the directory's basename or any architecture name in
     ``architectures: [...]`` against every recipe's ``matches:`` /
     ``architectures:`` field.
  3. Returning the first match (filesystem ordering is the tiebreak; in
     practice users should not ship two recipes with overlapping match
     sets).

Lookup is best-effort: a model with no matching recipe is fine — the
CLI just falls back to its own defaults. The recipe NEVER overrides an
explicit CLI flag (see :func:`Recipe.merge_into_args`).

Recipe parsing + lookup.

Each YAML recipe declares::

    name: qwen3-4b
    architectures: [Qwen3ForCausalLM]
    matches: [Qwen3-4B, Qwen3-4B-Instruct-2507]
    calibration:
      corpus: c4
      num_runs: 3
      num_tokens: 4096
      target_bit_widths: [2, 4, 8]   # OR target_bpw: 4.0
      min_improvement_pct: 5.0
    weight_quants_supported: [bf16, fp16]
    default_max_context: 32768
    eval_set: docs/eval_sets/short_recall.txt

Matching algorithm at boot:

  1. ``matches:`` is checked against the model directory basename
     (case-insensitive substring match in either direction — the recipe
     lists the canonical model basename(s); a deployment may rename the
     directory locally and we still want to find the recipe).
  2. ``architectures:`` is checked against ``config.json``'s
     ``architectures: [...]`` array (exact, case-sensitive — HF
     conventions).
  3. First recipe to match wins. Filesystem listing is sorted so the
     order is deterministic.

A user can also pass ``--recipe <path-or-name>`` to bypass auto-detection.

Tiny YAML subset parser — top-level scalars + one-deep blocks +
flow-style lists. Block-style lists with leading dashes are also
accepted at indent 2.

Raises ``ValueError`` on unrecognized constructs so a user mistake
surfaces at boot rather than producing a silently-wrong recipe.

Load a recipe from a path or registry name.

``path_or_name`` is treated as a path if it contains a separator or
ends with ``.yaml`` / ``.yml``; otherwise looked up in the built-in
registry under ``arbi_serve/recipes/<name>.yaml``.

Raises :class:`FileNotFoundError` if the recipe is missing.

Fill in any *missing* CLI defaults from this recipe.

argparse ``Namespace`` attributes that already carry a non-default
value are NOT overridden — CLI flags always win. Each
recipe-touchable arg is parsed with ``default=None``, so a value
of ``None`` here means the CLI flag was not given and the recipe
may fill it in.

arbi-serve runtime helpers.

Four modules live here:

  - :mod:`arbi_serve.runtime.capture` — per-kind capture/replay impls.
    Each kind keeps its own typed pool (``CapturedGraphPool``,
    ``LayerCapturedGraphPool``, ``DrafterChainGraphPool``) —
    different graph classes
    with different replay signatures, so polymorphic unification
    never paid; ``engine.cudagraph_pools`` constructs each one
    directly.

  - :mod:`arbi_serve.runtime.model_runner` — :class:`ModelRunner`
    Protocol + :class:`EagerModelRunner` impl. The "execute one
    slate" seam — flat-tensor build, captured-graph or live forward,
    sample. The run loop owns orchestration; the runner owns
    execution.

  - :mod:`arbi_serve.runtime.multi_group_cache` —
    :class:`MultiGroupCapturedGraphCache`. Wraps the
    :class:`CapturedGraphPool` with an extra ``(model_path, params_hash)``
    key for multi-tenant deployments.

  - :mod:`arbi_serve.runtime.sleep_allocator` —
    :class:`SleepableTensorPool` for sub-1s same-process sleep / wake.

  - :mod:`arbi_serve.runtime.timing` — :class:`RequestTiming` (per-
    request observability with first-call cost separation). Lets
    benchmarks distinguish "real work" from one-shot warmup costs
    (CUDAGraph capture, JIT first-call, autotune lookup).

Shared helpers for the per-step batch materializer.

Small, side-effect-free helpers used by both the persistent-decode
fast path (:mod:`arbi_serve.runtime._batch_build_persistent`) and the
full-rebuild materialization path
(:mod:`arbi_serve.runtime._batch_build_materialize`): decode input-stream
context resolution, sliding-window caching, host-mirror construction, and
the incremental block-table patch. Every name here is re-exported from
:mod:`arbi_serve.runtime.batch_build`.

Return ``(stream_ctx, event)`` for the decode input-prep pipeline.

When ``ARBI_DECODE_INPUT_PIPELINE`` is ON (and CUDA is present), lazily
create the runner's dedicated input stream + reusable CUDA event,
``wait_stream`` the default (compute) stream onto it so the input-prep
ops order AFTER the prior default-stream work (this step's sample,
page-table writes), and return a ``torch.cuda.stream(input_stream)``
context + the event to record at the end of the block. The captured-
graph replay then ``wait_event``s that event on the default stream
before reading ``pb.*``.

When OFF / on CPU, returns ``(nullcontext(), None)`` — every input-prep
op runs on the default stream exactly as before (bit-identical).

Return the model's shared PAGED_KV ``sliding_window``, cached.

Every PAGED_KV :class:`LayerSpec` carries the same ``sliding_window``;
the value is resolved with a ``next(...)`` scan over ALL layer specs.
The model's ``layer_specs`` list and each spec's ``state_kind()`` /
``sliding_window`` are fixed for a loaded model, so the scan result
is loop-invariant across decode steps. Cache it on the runner keyed
on ``id(eng.model)`` so a hot-swap (new model object) re-resolves.

Returns the same value as scanning ``eng.model.layer_specs`` for
the first spec with ``state_kind() == StateKind.PAGED_KV`` and
reading its ``sliding_window`` — computed once per loaded model
instead of once per step.

Return the ``scratch.forward_arena`` named-pool ``use()`` context.

Routes the fresh-alloc per-step batch tensors through the
activation arena so they are attributed instead of leaking into the
default allocator. Falls back to :func:`contextlib.nullcontext` when
the engine has no named pools (CPU stub / fixtures).

Host twin of ``cu_seqlens_k`` = ``[0, cumsum(h_seq_lens[:B])]``.

The DEVICE ``cu_seqlens_k`` is produced by a device cumsum in the
build path (it has no H2D ``h_*`` copy source), so the TP transport
host mirror recomputes it on the host from ``h_seq_lens`` — an exact
integer cumsum, bit-identical to the device ``torch.cumsum``.

The cumsum is written IN PLACE into the persistent
``pb.h_cu_seqlens_k`` buffer (allocated once, sized ``max_num_seqs +
1``) and a ``[:B+1]`` view is returned: element 0 is the leading
zero (written once at buffer construction, never touched here —
same as the device ``cu_seqlens_k[0]``), and ``[1:B+1]`` is the
int32 cumsum. Reuse is safe because the mirror is consumed
SYNCHRONOUSLY within the step that builds it — the TP path pickles
the bytes into the shm ring (``ShmControlQueue.enqueue``) and the
TP=1 ``bypass_safe`` read completes on the CPU — both before the
next step overwrites the buffer.

Host-side ``cu_seqlens_q`` for FLA's ``cu_seqlens_cpu``, or ``None``.

THE single decision point for whether a :class:`ScheduledBatch` carries
:attr:`~arbi_serve.engine.batch.ScheduledBatch.cu_seqlens_q_cpu`. Serving
(:func:`~arbi_serve.runtime.batch_build.build_batch`) and every boot-time
synthetic-batch builder call THIS function, because the field is one the
compiled GDN layer GUARDS on: ``_resolve_cu_seqlens`` branches on its
None-ness inside the ``fullgraph=True`` region, so a boot that warms one
state while serving produces the other leaves the first live request to
JIT-compile the layer (``jit_compile_serving``). Two constructions that
merely agree today drift the moment either side is edited; one function
cannot.

``None`` on any boot that can capture a prefill forward. That is not a
capability check — the host tensor is always available — it is the
confinement the capture hazard requires: host values baked into a
recorded graph are frozen launch parameters, and a replay at other
lengths would produce silently wrong chunk indices. See
:attr:`arbi_serve.engine.batch.ScheduledBatch.cu_seqlens_q_cpu`.

``pb`` is the persistent :class:`PiecewiseBuffers` when the caller wrote
this batch's boundaries into the pinned ring (the twin is then that same
ring slice — no allocation, and the exact object serving hands the layer),
``None`` otherwise, in which case ``cu_q`` supplies the host values — as a
sequence, or as the int32 host tensor the caller already built to source
its own H2D (:func:`~arbi_serve.runtime.split_mixed._slice_sub_batch`),
which is then adopted rather than copied.

The non-``None`` host twin, with the config gate BYPASSED.

The gate in :func:`host_cu_seqlens_q_twin` reads ``cfg.prefill_capture``,
which the admin console can flip at runtime — so a boot warmup that only
covered the state the boot resolved would hand the post-flip traffic an
uncompiled branch. Boot warmup enumerates BOTH states through this pair;
nothing on the serving path should call this one directly.

Build the TP-transport :class:`HostBatchMirror` from the pinned ring.

Every core scheduling tensor was just written into the pinned-host
``pb.h_*`` buffers before the H2D copy; attaching these host views
lets rank-0 ``_serialize_batch_to_plan`` build the worker StepPlan
from host memory with no device→host copy / no per-step sync.

Patch ``pb.h_block_table`` in place for rows whose page chain changed.

Rewrites only the rows whose ``(request_id, page_count)`` differs
from the cached occupant (a decode row grows only at a
``block_size`` boundary), zeroing each rewritten row's tail to the
FULL buffer width so a later wider-``max_pages`` step reads clean
zeros. ``remap_epoch`` movement suppresses every skip this step so
the rows rebuild from the live chains. Drops the whole occupant
cache on first use or a piecewise-buffer rebuild. Returns ``True``
when any row's slice was rewritten this step.

A row that kept its occupant and only GREW writes the appended
columns alone. Its prefix already holds this request's pages (the
occupant key says so — the same premise the skip above rests on:
page content that moves under a live request bumps ``remap_epoch``)
and its tail is already zero (the previous rewrite of this row
zeroed to ``full_w``). The full-width rewrite is ``O(max_pages)``
per growing row while the growth itself is one page, and
``max_pages`` is ``ceil(context / block_size)`` — at a small block
size that is the dominant per-step host cost.

Shared verbatim by :func:`try_persistent_decode` (which consults the
return to decide whether to re-issue the block-table H2D) and the
persistent materialization path (which always copies), so the two
stay in lockstep.

Full-rebuild slate gather + tensor materialization for the batch builder.

The per-step full-rebuild path: the slate walk (:func:`_gather_slate`)
that allocates page slots and gathers the flat batch vectors, the
persistent-buffer (:func:`_materialize_persistent`) and fresh-alloc
(:func:`_materialize_fresh`) tensor materializers, the recurrent
slab-row resolution (:func:`_resolve_recurrent_state`), and the
host-mirror + occupant refresh (:func:`_attach_host_mirror_and_refresh_occupant`).
Every name here is re-exported from :mod:`arbi_serve.runtime.batch_build`.

Per-step flat batch vectors gathered from the slate walk.

The dual-state ``prep_writes`` / ``flat_*`` carry both the CPU-prep
numpy-slice form and the Python list-extend form because
``use_persistent`` is decided AFTER the slate walk; the
materialization path picks the right form.

``flat_ids`` / ``flat_positions`` cover only the flat spans no
``prep_writes`` entry owns, so their length is ``n_tokens`` minus the
prepared token count; :func:`_prep_list_runs` maps them back onto the
flat axis. ``flat_slots`` is always full-length.

Allocate page slots and gather the per-token flat batch vectors.

Walks the slate once, allocating slots and building ``flat_ids`` /
``flat_positions`` / ``flat_slots`` / ``seq_lens`` / ``cu_q`` plus
the per-row bookkeeping (``per_row_pages`` / ``per_row_is_prefill``)
and the deferred-write side lists (``prep_writes`` /
``gpu_overwrites`` / ``gpu_overwrite_host_tokens``). ``use_persistent``
is decided AFTER this walk, so the dual-state forms are carried and
the materialization path picks the right one.

``(dest_offset, src_offset, length)`` runs for the Python lists.

``prep_writes`` owns a set of disjoint, ascending spans of the flat
axis; every other span is owned by ``flat_ids`` / ``flat_positions``,
which carry those tokens contiguously and in the same order. Each run
says where a contiguous list segment starts on the flat axis, where it
starts in the list, and how long it is.

Persistent-buffer (cudagraph) materialization of the flat batch.

Writes the gathered vectors into the engine's pinned-host ``pb.h_*``
ring via zero-copy numpy views, issues one H2D ``copy_`` per field,
patches the block table incrementally, applies the decode-row GPU
sample overwrites, and computes ``cu_seqlens_k`` into the persistent
slice. Returns the prefix-sliced device tensors
``(input_ids, positions, slot_mapping, seq_lens_t, cu_q_t, bt, cu_k)``
whose ``data_ptr``s match the captured per-layer graph references.

Fresh-alloc fallback materialization (CPU stub / oversized / capture off).

Allocates fresh per-step device tensors routed through the
``scratch.forward_arena`` named pool so they are attributed instead of
leaking into the default allocator, builds the eager ``(B, max_pages)``
block table, applies the decode-row GPU sample overwrites, and
computes ``cu_seqlens_k``. Returns
``(input_ids, positions, slot_mapping, seq_lens_t, cu_q_t, bt, cu_k)``.

Resolve the recurrent (GDN / Mamba / ShortConv) slab-row mapping.

Returns ``(state_indices, state_indices_long, has_initial_state)`` —
all ``None`` on a pure-attention model. On a persistent prefill step
the row mapping AND the per-row precise ``has_initial_state`` mask
(``prompt_consumed > 0``) are written into the pinned
``pb.recurrent_*`` slices so the captured kernels read a stable
``data_ptr``, and the mask slice is ALSO returned for
``batch.has_initial_state`` — the same refreshed buffer serves both
the captured graphs and the eager builders (single derivation, no
extra H2D). Otherwise the mapping is materialized directly on
``batch.state_indices`` so the per-kind metadata builders honour it
without falling through to ``arange(B)`` (which would leak the prior
tenant's end-of-decode recurrent state), and the mask is ``None``
(builders derive ``seq_lens > 1``).

Attach the TP-transport host mirror + refresh the fast-path occupant.

On the persistent (cudagraph) hot path every core scheduling tensor
was just written into the pinned-host ``pb.h_*`` ring before the H2D
copy; attach those host views so rank-0 ``_serialize_batch_to_plan``
builds the worker StepPlan from host memory (no device→host copy / no
per-step sync). Then refresh the persistent-input-batch occupant
cache so the next step's :func:`try_persistent_decode` validates
against the composition this slow-path step materialized — recorded
only when the slate is all single-token rows (``N == B``); a mixed /
prefill composition invalidates it so the next fast attempt rebuilds.

Persistent-input-batch fast path for a pure-decode step.

:func:`try_persistent_decode` implements the per-step batch
materializer's decode fast path. On a steady-state decode step
(unchanged slate composition) it rewrites only the changed slices of
the persistent pinned-host ring buffer instead of rebuilding every
flat Python list, returning a :class:`ScheduledBatch` bit-identical to
the full rebuild. Re-exported from :mod:`arbi_serve.runtime.batch_build`.

Persistent-input-batch fast path for a pure-decode step.

On a steady-state decode step the slate is the same requests in
the same row order as the previous step, every row a single-token
(``n == 1``) decode with its previous sample held on GPU. The full
:func:`build_batch` slate loop rebuilds every flat Python list
(``flat_ids`` / ``flat_positions`` / ``flat_slots`` / ``seq_lens`` /
``cu_q``) and re-coerces them into the pinned-host ring buffer; for
a pure-decode step those vectors change only by a +1 increment on
``positions`` / ``seq_lens`` and a fresh per-row slot, while
``cu_seqlens_q`` is the constant ``[0, 1, ..., B]``. This path
rewrites just the changed slices and skips the unchanged H2D
copies (``cu_seqlens_q`` always; ``block_table`` when no row's page
chain grew this step), producing a :class:`ScheduledBatch` that is
BIT-IDENTICAL to the full rebuild.

Returns ``None`` (caller falls through to :func:`build_batch`) on
ANY condition the fast path does not cover: a non-persistent
engine (CPU stub / capture off / oversized), a prefill / multi-token
row, a row whose previous sample is not on GPU (first decode after
prefill, freshly-preempted, sync CPU path), any LoRA / multimodal /
MTP row, or a slate whose composition changed vs the previous step
(admit / evict / preempt / reorder). All those route to the full
rebuild, which refreshes the occupant cache so the NEXT step's fast
attempt validates correctly.

Recurrent (hybrid GDN / Mamba / ShortConv) decode pins the per-row
slab-row mapping on ``batch.state_indices`` (resolved via
``view.row_for(req_id)``), exactly as the slow path does. This is
load-bearing under tensor parallelism: rank 0 re-derives the rows
from the live slate, but the TP WORKERS have no request stream and
no recurrent pool — they consume ONLY the broadcast batch. When
``state_indices`` is None the worker's captured replay falls back to
``arange(B)`` (``captured_lookup.live_recurrent_state_indices_for_batch``),
reading the WRONG slab rows (row 0 is the permanent zero-sentinel,
never a live request's row), so its GDN/Mamba shard advances against
near-zero state while rank 0 advances the real state — the merged
cross-rank output is corrupt. Pinning the rows here makes the
broadcast carry them so both ranks read identical cells.

The slab rows are stable within a decode window (allocated once at
admission, freed at finish/preempt) and the eligibility gate above
already rejects any composition change (admit / evict / preempt /
reorder), so re-deriving from the unchanged slate each step is
bit-identical to the slow path's ``state_indices`` mapping.

StepPlan-side batch utilities for the per-step batch materializer.

The captured-prefill replay staging
(:class:`PrefillReplayStaging` / :func:`resolve_prefill_replay_staging`),
the batched decode-sample scatter (:func:`scatter_gpu_overwrites`), and
the :class:`StepPlan` ↔ :class:`ScheduledBatch` bridges
(:func:`populate_tensors` / :func:`reconstruct_verify_batch`). Every name
here is re-exported from :mod:`arbi_serve.runtime.batch_build`.

Persistent, reusable staging buffers for captured-prefill padding-replay.

The captured-prefill replay path pads the live request's flat
tensors up to the chosen ``bucket_N`` before
:meth:`CapturedGraph.replay`. ``replay`` ``copy_()``s every passed
kwarg into the graph's OWN persistent buffers (see
:meth:`arbi_serve.runtime.capture.decode.CapturedGraph.replay`,
the ``self.input_ids[:N].copy_(input_ids)`` block) and then launches
against those — it does NOT retain the passed ``data_ptr``. So the
padding tensors are pure copy SOURCES, transient per replay, and a
single reusable scratch set is trivially safe: replay snapshots their
content before the next call can touch them. The forward path is
single-threaded per engine step (no nested/concurrent prefill
replay), so one set per runner suffices.

Sized ONCE to the worst case so every reachable bucket fits a
contiguous prefix slice:

  * flat-token dims (``input_ids`` / ``positions`` / ``slot_mapping``)
    at ``max(prefill_cudagraph_buckets)`` — the largest bucket the
    lookup can pick.
  * ``block_table`` at ``(1, max_cap_pages)`` where ``max_cap_pages``
    is the captured prefill graphs' ``max_pages_in_table`` (full
    ``ceil(max_context / block_size)`` page width — capture_prefill
    bakes this).

Mirrors :class:`arbi_serve.spec_decode.verify_buffers.VerifyBuffers`
(the decode-verify reference pattern): allocate worst-case once into
a named pool, ``copy_()`` the live prefix in per step. Unlike
VerifyBuffers these buffers are NOT referenced by any captured graph
(replay copies out of them), so they need no ``data_ptr`` stability /
sleep-pool registration — they are plain scratch, dropped + realloced
on the next prefill replay after an ``eng.pool`` swap.

The ``cu_seqlens_q`` / ``cu_seqlens_k`` / ``seq_lens`` slots are tiny
fixed ``(2,)`` / ``(1,)`` int32 tensors (B=1 prefill); they replace
the per-replay ``torch.tensor([0, real_N])`` materializations.

Return the runner's persistent captured-prefill staging buffers,
allocating them once (per ``eng.pool`` generation) on first use.

Sized to the worst case so every reachable bucket / page width
fits a contiguous prefix slice: flat-token dims at
``max(prefill_cudagraph_buckets)`` and ``block_table`` at the
captured graphs' ``max_pages_in_table`` (passed as ``cap_pages``;
invariant across all prefill captures — ``capture_prefill`` bakes
the full ``ceil(max_context / block_size)`` page width). Dtypes /
device follow the live batch tensors (``like_*``) so the per-step
``copy_()`` is a pure D2D refresh with no dtype cast.

Cache is keyed on ``id(eng.pool)`` (mirrors the flush-hook
resolvers): a hot-swap that rebuilds the pool forces a fresh
allocation, so the staging
buffers never outlive the VA generation their content is paired
with. The allocation routes through the engine's
``graph_buffers_pool`` named pool — the same sink
:class:`VerifyBuffers` uses for engine-lifetime captured-graph
scratch — so it is attributed once at first-replay rather than
per-step in the activation arena. Falls back to the default heap
on CPU stub paths / fixtures without a live engine.

Write the previous-step GPU samples into their ``input_ids``
slots in one scatter instead of one D2D ``copy_`` per row.

Each entry is ``(slot, src)`` with ``src`` a 0-d device tensor
holding the row's previous sample. ``target`` is the persistent
(or fresh) ``input_ids`` buffer; ``slot`` indexes it absolutely.
The write is in place, preserving the buffer's ``data_ptr`` so it
is safe under cudagraph capture. Equivalent to
``target[slot].copy_(src)`` for every entry.

Return a new :class:`StepPlan` with the flat tensor fields
populated from a materialized :class:`ScheduledBatch`.

The mapping is direct:

  - ``plan.input_ids``  ← ``batch.input_ids``     (sum_T,)
  - ``plan.positions``  ← ``batch.positions``     (sum_T,)
  - ``plan.cu_seqlens`` ← ``batch.cu_seqlens_q``  (B+1,)
  - ``plan.request_ids`` ← per-row ``str(req.request_id)``

The plan is frozen, so we use :func:`dataclasses.replace` to
produce a fresh instance. ``slate`` is preserved as-is on the
new plan so any downstream consumer reading from ``plan.slate``
continues to work.

``cu_seqlens_k`` and ``slot_mapping`` are NOT lifted onto the
plan — they live on the per-:class:`StateKind` metadata the
backend's :class:`MetadataBuilder` consumes; the plan carries a
``state_meta`` dict for that.

Rebuild the verify-pass :class:`ScheduledBatch` from ``plan``.

Single source of truth for the verify ``StepPlan`` →
``ScheduledBatch`` mapping. The flat token tensors are the very
same staged :class:`VerifyBuffers` views the plan carries
(``input_ids`` / ``positions`` / ``cu_seqlens`` are aliased, not
copied); the paged-KV routing header comes from
``state_meta[StateKind.PAGED_KV]`` and the recurrent slab-row
mapping from the plan's recurrent ``state_meta`` entry (``None``
on pure-attention models). The remaining batch-shaped fields are
derived from the canonical plan payload without a stashed copy:
``mtp_meta`` from :attr:`StepPlan.mtp_step_plan`,
``lora_assignments`` / ``req_ids`` from :attr:`StepPlan.slate`,
and ``mtp_fill_enabled=True`` (the verify forward always fills
the bundled MTP-head KV).

C++ shim for the cuMem pluggable allocator.

Isolates the Python<->C++ language boundary used by
:mod:`arbi_serve.runtime.cumem_allocator`: this module holds the embedded
C++ source and builds it (once, process-cached) into a loadable extension
via ``torch.utils.cpp_extension.load_inline``. ``cumem_allocator`` imports
:func:`build_ext` from here.

The extension exports the two C-ABI symbols ``torch``'s
``CUDAPluggableAllocator`` expects:

  void* my_malloc(ssize_t size, int device, cudaStream_t stream);
  void  my_free  (void* ptr, ssize_t size, int device, cudaStream_t stream);

``my_malloc`` reserves a VA, creates+maps a physical handle, sets access,
then calls a Python callback with ``(ptr, handle, size, device)`` so the
Python side records the bookkeeping for sleep/wake. ``my_free`` calls a
Python callback that returns the handle, then unmaps + releases + frees the
VA.

Two extra entry points drive sleep/wake from Python:

  void unmap_and_release(uintptr_t ptr, ssize_t size, uintptr_t handle);
  uintptr_t create_and_map(int device, ssize_t size, uintptr_t ptr);

``unmap_and_release`` drops physical pages but keeps the VA;
``create_and_map`` binds fresh physical pages at the SAME VA and returns the
new handle.

Construction + live-flag / flush-hook lifecycle for the eager runner.

The :class:`EagerModelRunner` constructor (the per-step cache fields it
allocates), the live-refreshable runtime-flag snapshots, and the
recurrent-pool flush-hook resolvers. :class:`EagerModelRunner` mixes this
in; every attribute the constructor sets is read by the runner's forward /
execute methods and the sibling :class:`_RunnerOpsMixin`.

No-op stand-in for :meth:`MultiStatePool.flush_pending_zero_clears`.

Pure-attention models have no recurrent slabs, so the engine's
pool exposes no flush hook. The runner caches this sentinel on
first ``execute`` so the per-step path can call unconditionally,
skipping a per-step ``hasattr`` probe.

No-op stand-in for
:meth:`MultiStatePool.flush_pending_savepoint_resumes`.

Pure-attention models (no recurrent slabs) and hybrid-model
builds without a savepoint store wired both bind to this sentinel
so the per-step ``execute`` / ``forward`` paths skip the
``hasattr`` probe.

Re-read the live-applicable runtime flags into the cached attrs.

The runner snapshots several runtime flags into instance attributes at
construction so the hot path reads a plain bool instead of re-deriving
the env-overlaid :func:`runtime_flags` every step. A live config
override of a ``capture_affecting=False`` flag installs the new value
into the active-member flag overlay (``set_active_flag_overlay``) and
then calls THIS so the already-built runner picks the new value up
WITHOUT a member re-prepare / re-capture.

Only flags whose effect is purely a per-step routing decision — they
do NOT change any captured cudagraph's shape or the compiled forward —
are re-snapshotted here. ``split_mixed_decode_prefill`` /
``decode_pad_cudagraph`` pick between an already-captured graph and an
eager fallback per step; the persistent-buffer / async-output /
input-pipeline flags toggle host-side fast paths that produce the same
device values either way. All are safe to flip on the live member.

Snapshot the live-refreshable runtime flags into cached attrs.

Single source of truth for the per-step routing-only flags so the
construction-time snapshot (:meth:`__init__`) and the live
re-snapshot (:meth:`refresh_runtime_flag_snapshots`) cannot drift.
Only flags whose effect is purely a per-step routing decision —
they do NOT change any captured cudagraph's shape or the compiled
forward — belong here. (``time_build_batch`` is read once at
construction and is NOT live-refreshable, so it stays out.)

Resolve a recurrent-pool flush hook, keyed on ``id(eng.pool)``.

Shared body for :meth:`_resolve_flush_zero_clears` and
:meth:`_resolve_flush_savepoint_resumes`, which differ only in
the cached attr, the bound-method name, and the noop sentinel.
``eng.pool`` is None during construction;
the boot sequence wires it later, and sleep-mode / hot-swap may
swap it, so a changed pool id forces a fresh ``getattr``. Returns
``(hook, pool_id)``; the caller writes both back onto its cache
attrs. Returns ``noop`` when the pool is None or the model is
pure-attention (no ``method_name`` bound method).

Drop the cached bound methods of ``eng.pool`` so the pool can free.

The two flush hooks are cached as BOUND METHODS keyed on ``id(eng.pool)``
and are therefore strong references to the pool live at resolve time. A pool the engine has stopped using — the activation
profile's scratch pool, the pool a per-kind teardown just dropped —
stays reachable through them, and with it every recurrent and paged slab
it owns, until the NEXT step re-resolves the hook. That is a whole
boot's worth of resident slabs held past their owner, so the code that
drops a pool releases the hooks in the same breath. Idempotent, and the
next resolve re-binds from the live pool.

Thin execution delegators + sampling + verify/seed helpers for the
eager runner.

The one-line delegators onto the ``runtime.*`` free-function
collaborators (batch build, captured-graph lookup, forward exec,
multimodal merge, activation profiling), the sampler dispatch
(:meth:`_sample` / :meth:`_sample_async`), and the MTP verify/seed
forward helpers. :class:`EagerModelRunner` mixes this in; the methods
read the instance state the :class:`_LifecycleMixin` constructor sets.

Arm the fold emit's tables for a replayed prefill; the plan or ``None``.

The metadata builder's ``_attach_fold_split`` runs only on the eager
forward; a captured replay skips it. Same derivation
(:func:`~arbi_serve.cache._fold_emit_staging.publish_step_plan`), same
refusals, on ``gdn_fold_split_delivered``.

Delegate to :func:`forward_exec.run_model_forward`.

Publishes any host-staged rows the model prepared for this batch
first. This is the canonical ``model.forward`` call site, which is
why the publish lives here: a path that reaches the model another
way (the MTP seed forward does) would otherwise read whatever the
previous step left in the fixed staging buffer — same shapes,
another step's values, and nothing raised.

Single canonical ``model.forward`` call site (inference_mode +
activation-arena wrapping). ``execute`` / ``forward`` /
``run_step.step`` MTP-seed all route through this wrapper.

``reqs`` (optional): the batch's per-row request list, in slate
order — threaded through to :func:`forward_exec.resolve_embed_override`
for callers that want to honor ``Request.pending_embed_override``.
``None`` (the default; every call site not explicitly updated) is
a complete no-op — see that function's docstring.

``collapse_heads`` (default ``True``): threaded straight through to
:func:`forward_exec.run_model_forward`'s own parameter of the same
name — see its docstring. Every existing call site leaves this at
the default, so it is a complete no-op for them.

``(logits, hidden_last)`` for a seed-flavor forward, where
``hidden_last`` is the LAST-SLOT ``(B, H)`` hidden per row:
whole-forward prefill replay when a hidden-retaining rung matches,
else the raw eager ``_run_model_forward(return_hidden_state=True)``
plus the same last-slot gather (device-index, no host sync).

``owns_step_prep=False`` (run_step seed path): the caller drained
the recurrent flush queues and built metadata
(``_build_mtp_seed_batch``). ``owns_step_prep=True`` (SPMD seed
tick, where this call REPLACES ``forward``): drain the zero-clear +
savepoint queues here — the drains must precede any read of
recurrent state, replay included — and build metadata only on the
eager branch, mirroring ``forward``'s own preamble.

The rung resolves from batch facts + the captured pool only
(``lookup_captured_graph_for_prefill(require_hidden=True)``), so any
two ranks holding the same pool resolve the same branch. Fire/refuse
attribution: ``mtp_seed_prefill_replay`` here plus the lookup's own
``prefill_graph_lookup`` counter.

The six per-step replay tensors pulled straight off ``batch``.

``input_ids`` / ``positions`` / ``slot_mapping`` / ``seq_lens`` /
``cu_seqlens_k`` / ``block_table`` — the bundle the exact-shape
decode replay (:meth:`execute`) and the verify-forward replays
(:meth:`forward`) all pass to ``CapturedGraph.replay`` identically.
The caller adds ``recurrent_state_indices`` (and, on the verify
path, ``return_hidden_state``). The prefill replay does NOT use
this — it passes padded staging tensors plus ``cu_seqlens_q``.

StepPlan-aware verify-forward entry point.

Reconstructs the per-step :class:`ScheduledBatch` from the
plan's flat tensors + ``state_meta`` and delegates to
:meth:`forward`. The MTP verify path's public engine boundary is
always a :class:`StepPlan` — there is no staged
``ScheduledBatch`` payload riding along; the per-step routing
header (paged-KV tensors + the scalar attn fields + the
recurrent slab-row mapping) is transported in ``state_meta`` and
the per-:class:`StateKind` :class:`MetadataBuilder` chain runs
over the reconstructed batch exactly as the main path's builders
do.

Only verify-pass plans are accepted: a plan without the
``StateKind.PAGED_KV`` ``state_meta`` payload (e.g. an MTP-off
:meth:`StepPlan.from_slate` plan, whose tensors the runner's
``_build_batch`` owns) raises rather than silently no-op.

Delegate to :func:`batch_build.reconstruct_verify_batch`, then
attach the per-row LoRA state the verify forward needs.

The pure ``StepPlan`` → batch mapping is adapter-agnostic (it has
no engine handle), so it fills ``lora_assignments`` but leaves
``lora_state`` ``None``. Building the state needs the engine's
store + capture pool, which only the runner holds — so it happens
here, mirroring :func:`batch_build.build_batch`'s LoRA path.

Without this the verify/bonus forward would run with the adapter
UNAPPLIED: its logits would diverge from the LoRA'd main-model
decode, silently corrupting the committed bonus token and the
accept/reject decisions (a losslessness break for any
``lora_id != None`` request under MTP). Passing the capture pool
also keeps the captured verify graph correct — its replay reads
the pool's persistent ``data_ptr``s, which this ``populate()``
refills (the capture sweep seeds the same pool per ``lora_bucket``
in :mod:`runtime.capture.decode`).

Run one synthetic forward to measure the activation peak.

Thin delegator to :func:`profile_peak.profile_activation_peak`;
see that free function for the full contract (it runs once at
boot from ``engine.build.profile_and_size_kv_pool`` to derive
the auto-sized KV pool depth).

Delegate to :func:`batch_build.build_batch`.

The hot per-step flat-tensor materializer lives as a module-
level free function (taking this runner by reference) so the
hot loop carries no delegating boundary at the kernel level;
this one-line method preserves the ``eng.model_runner._build_batch``
call site + the persistent caches it reads on ``self``.

A model that derives host-resident rows from the request stream
(Qwen4-Exp PLE) prepares them HERE, against the slate that built
the batch, so every consumer of a slate-built batch has its rows —
including the MTP seed forward, which builds its batch through this
method and then reaches the model without going through
:meth:`execute`.

Sample without the host sync — return a deferred PendingSample.

fast paths as :meth:`_sample`, but:

  * NO ``.cpu().tolist()`` — sampling stays on GPU.
  * Every ready row's ``last_sampled_gpu`` is stamped from the
    GPU sample tensor so the NEXT step's ``input_ids`` slot is a
    D2D copy with no host round-trip (vLLM ``prev_sampled_token_ids``).
    This is REQUIRED here because the host ints aren't available
    this step — including the B=1 fast path, which the synchronous
    ``_sample`` left to read ``output_token_ids[-1]`` (not yet
    appended under deferred commit).
  * Issues the non-blocking D2H on the dedicated copy stream and
    records a CUDA event; returns a :class:`PendingSample` the
    run loop drains one tick later.

Returns ``None`` when the slate sampled nothing (all-dead /
all-mid-prefill); the caller treats that as "no pending output".

Per-step activation arena — bump-pointer allocator for inference transients.

Inference activations have *stack discipline* within a forward step:
tensors are allocated in forward order and the entire set is dropped at
end-of-step (no backward pass). The PyTorch caching allocator treats
each as a heap allocation — freeing, caching, and reusing per size
class — and *fragments* across layers when sizes vary. This fragments
even when the headline free counter looks healthy, producing an OOM
on a small allocation despite most of the pool nominally being free.

A stack arena gives O(1) alloc, zero fragmentation, and a predictable
peak. The pattern is:

  - reserve one big uint8 buffer up front (inside the named
    ``scratch.forward_arena`` MemPool so it sits in its own bucket
    and never fragments against KV / CUDAGraph / LoRA);
  - per allocation, return a tensor view whose storage is the arena
    buffer at the current bump-pointer offset (16-byte aligned);
  - rewind the offset to zero at every layer boundary and again at
    end-of-step. The buffer itself stays alive across steps; only the
    offset rewinds.

Critical invariant
------------------
**Every tensor the arena hands out must have its last reference
dropped BEFORE the next :meth:`reset`.** Otherwise a view outlives
its arena slot and the next allocation overwrites the bytes
underneath it. Two seams rewind: the eager per-layer dispatcher,
immediately before it allocates that layer's buffers, and the
:meth:`scope` context manager at the forward boundary
(``with arena.scope(): model.forward(...)``). The per-layer rewind is
what makes the arena's capacity a MAX over layers rather than a sum —
a bump allocator that only rewinds at end-of-step is not a stack.

Liveness contract
-----------------
``torch.Tensor.set_(storage, offset, size)`` re-points the existing
empty tensor at a slice of the arena's storage. The storage is the
arena's *untyped* storage; it is reference-counted by the arena's
own ``self._buf`` reference, so the per-allocation tensor's
contribution to its refcount is just an additional view. Dropping
the tensor does NOT free the storage — the arena keeps it alive for
the engine's lifetime.

This also makes the arena a GRAD-FREE construct. Autograd saves the
tensors a backward would need — the attention output is one of them
(``o_proj``'s weight gradient reads it) — so a forward run with grad
enabled keeps every arena buffer alive through the returned logits'
graph. The engine always forwards under ``torch.inference_mode``; a
caller that does not gets an :class:`ArenaEscape`, which is the right
answer rather than a silent overwrite.

Every allocation is tracked by a weak reference, and :meth:`reset`
checks that none of them survived. A survivor is a tensor whose bytes
the next allocation is about to overwrite while something still reads
it, so :meth:`reset` raises :class:`ArenaEscape` naming the survivors
rather than letting the corruption run silently.

The tracking is UNCONDITIONAL, not debug-gated: a guard that is off by
default leaves production on the side that corrupts silently, which is
the failure this arena is least able to survive. It is affordable
because the per-allocation record is a weak reference plus the raw
(shape, dtype, offset) triple; the human-readable description is
formatted in :meth:`reset`, and only for a tensor that actually
survived — so the healthy path never builds a string it will not read.

Eager-only
----------
:meth:`alloc` MUST run outside a Dynamo trace. ``Tensor.set_`` with
more than one argument is not traceable, and even a traceable
formulation would bake the bump-pointer offset into a guard and
recompile the region once per offset value. The engine therefore
allocates at the eager per-layer dispatch seam
(:meth:`~arbi_serve.models._layer_stack_model_mixin.LayerStackModelMixin.
_dispatch_layer_args`) and threads the resulting tensor into the
compiled block as an ordinary argument. :meth:`alloc` refuses a
traced call by name so a future caller that allocates from inside a
block gets that sentence instead of an opaque Dynamo error.

Stream awareness
----------------
The arena buffer is allocated on the current CUDA stream at construct
time; subsequent ``alloc()`` calls hand back tensors that share that
storage. Because async kernels read/write storage by raw pointer,
the per-tensor stream tag is irrelevant — what matters is that no
stream is reading from the arena buffer when :meth:`reset` is called
and the next allocation overwrites the same bytes. The same
discipline applies to PyTorch's caching allocator: the arena does
not weaken it, only inherits it. Callers running multi-stream
forward (we don't today) would need to ``torch.cuda.synchronize()``
or insert a barrier event before :meth:`reset`.

Opt-in
------
This module is OPT-IN behind ``--enable-activation-arena``. The
default-off behaviour is preserved by passing ``arena=None`` into
``model.forward``; every allocation site has a fallback to plain
``torch.empty``/``torch.zeros``.

Raised by :meth:`ActivationArena.alloc` when the requested
allocation would exceed the arena's capacity.

Carries the requested-end offset and capacity so callers /
operators can decide whether to bump the arena size or retry on
the heap path. The arena does NOT auto-resize — a resize during a
step would orphan every live view.

Raised by :meth:`ActivationArena.reset` when an arena tensor
outlived the region it was allocated for.

The arena rewinds its bump pointer on reset, so the next allocation
hands out the same bytes. A tensor still alive at that moment is a
tensor whose contents are about to change underneath its reader —
silently wrong output, not a crash. Every survivor is named
(shape / dtype / byte offset) so the escaping buffer can be found
from the message alone.

``survivors`` carries the same descriptions as a list.

Bump-pointer arena over one big uint8 buffer in a named pool.

Args:
    capacity_bytes: total arena buffer size. The buffer is allocated
        once at construct time and survives for the arena's lifetime.
    named_pool: the :class:`NamedMemPool` the arena buffer lives
        in. Use the bucket named ``scratch.forward_arena`` so the
        buffer sits in its own segment and never fragments against
        KV / CUDAGraph / LoRA pools.
    device: optional CUDA device override. Defaults to the named
        pool's device.
    debug: when True, :meth:`reset` logs the allocation count and
        bytes it is rewinding. Off by default — the escape check
        that catches the actual correctness failure
        (:class:`ArenaEscape`) is unconditional and does not need
        it.

Helper: pull the ``activation_arena`` named pool out of the
registry and construct an :class:`ActivationArena` against it.

The named pool is created lazily here if the registry didn't
register it at engine construction time — keeps wiring local to
this module.

Bump-pointer allocate a tensor view inside the arena.

Returns a Tensor whose storage is the arena's underlying
``uint8`` buffer at the current 16-byte-aligned offset, viewed
as ``dtype`` at ``shape``. The bump pointer advances by
``prod(shape) * dtype.itemsize``; the tensor's lifetime is the
caller's, but the storage outlives every per-step allocation.

A request that would exceed ``capacity_bytes`` does NOT
raise: it falls back to the default allocator for that one
allocation (logged once, counted in ``overflow_bytes``) so an
under-sized arena costs the fragmentation saving and never
correctness.

MUST be called eagerly — see the module docstring's "Eager-only"
section. A traced call is refused here by name.

Rewind the bump pointer to zero.

Does NOT free the underlying buffer — the buffer stays alive
for the arena's lifetime so subsequent resets reuse it without
re-allocating. The high-water mark is preserved across resets
so :meth:`snapshot` can report the worst-case usage seen so
far.

Called at two granularities: once per LAYER, from the eager
dispatch seam, immediately before that layer's allocations; and
once per STEP, by :meth:`scope`. The per-layer rewind is what
makes the arena's capacity a MAX over layers rather than a sum
— the whole point of a stack allocator — and is sound because
the seam is the only allocation site and the block consumes its
buffers before returning.

``check_escapes`` (default True) raises :class:`ArenaEscape`
when any tensor handed out since the previous reset is still
alive. It is turned off only when an exception is already
propagating, where the in-flight traceback legitimately holds
the frames that reference them and the original error is the
one worth seeing.

Context manager pairing :meth:`reset` with the forward
boundary.

Usage::

    with arena.scope():
        logits = model.forward(batch, arena=arena, ...)

On normal exit *or* on exception, :meth:`reset` rewinds the
offset. The buffer survives — only the bookkeeping resets.
The escape check runs on the normal exit only: on an exception
the traceback holds the forward's frames, so every arena tensor
it touched reads as a survivor and an :class:`ArenaEscape`
would replace the error the caller needs to see.

Return a metric-friendly dict for OTEL gauges and admin
endpoints.

Keys:
  - ``capacity_bytes``: arena buffer size (constant).
  - ``used_bytes``: current bump-pointer offset.
  - ``high_water_bytes``: max ``used_bytes`` ever reached.
  - ``num_live_allocs``: count of allocs handed out since the
    last :meth:`reset` (zero between steps).
  - ``fragmentation_ratio_estimate``: ``1.0 - high_water /
    capacity`` — what fraction of the arena has *never been
    touched*. Not real fragmentation (the arena cannot
    fragment by construction); included because dashboards
    expect a 0..1 utilisation knob and this is the closest
    arena-shaped analog.

Multi-shape activation profiler.

Boot-time helper that runs synthetic forward passes at every shape
the engine actually serves and reports the MAX of their
``torch.cuda.max_memory_allocated()`` deltas — including allocations
that landed in :class:`NamedMemPool` segments. The number is what
:mod:`arbi_serve.engine.memory_budget` and
:func:`arbi_serve.engine.build.verify_gpu_headroom` should subtract
from free VRAM before sizing the KV pool.

Allocated, not reserved
-----------------------
The profiler reads ``torch.cuda.max_memory_allocated(device)``. That
counter is an AGGREGATE over every pool: the caching allocator stamps
each block with the AGGREGATE stat type whichever pool served it, so a
private :class:`NamedMemPool`'s allocations are already inside it. The
registry's ``caching_allocator_allocated_bytes`` delta is added as a
floor under the reading for any pool the counter somehow missed; it is
taken at the END of the forward, so on a normal boot it is ~0 and the
two numbers agree.

Which pool held the peak
------------------------
One counter over every pool is one number where the serving floor needs
two: the probe forwards run inside ``scratch.forward_arena``
(:func:`profile_engine`), and that pool is reserved in full by its own
floor term. Charging its bytes to the floor's base term as well reserves
them twice. The end-of-forward pool delta cannot separate them — it is a
residue, not a high-water — so :func:`_profile_one` replays the
allocator's per-allocation trace, which carries the owning pool id on
every event, and :func:`split_trace_peak` maximises the arena curve and
the off-arena curve each at its OWN moment. Each :class:`ShapeProfile`
then carries both the arena's high-water and its residue, and a shape
whose trace could not be read carries ``None`` — the whole peak stays on
the base term — together with the REASON it could not be read, so a KV
ceiling derived from that inflated floor is never reported to an operator
as a full card (:meth:`MultiShapeProfile.unmeasured_arena_note`).

Every shape is profiled twice: an unrecorded warmup pass, whose only
outputs are the shape's allocator-event count and its cold peak, then the
measured pass, recorded at exactly the count the warmup observed. The ring
therefore has no declared capacity at all, and the measured window is the
WARM one — which matters because the decode shape is a boot's first
decode-class forward and the TKV split-K autotune fires inside it (see
:func:`_profile_one_measured` for the numbers).

Why MAX-of-shapes
-----------------
Profiling a single synthetic shape under-estimates the peak: a
moderate ``max_batched_tokens`` shape can report a modest activation
profile while the real worst-case (a large prefill chunk against a
half-full KV pool) blows through it and OOMs on a small allocation.
The fix is to profile several representative shapes and take MAX:

    1. Pure prefill at ``max_batched_tokens``       (single big chunk)
    2. Pure decode at ``max_batch``                 (B requests × 1 tok)
    3. Mixed: half-prefill + half-decode in one fwd (chunk_prefill)

These are the three shape families the engine schedules; activation
peaks behave very differently across them (prefill: O(seq² heads),
decode: O(B), mixed: O(seq² + B*overhead)). MAX bounds the peak of any
ONE step, because the shapes are time-exclusive — a step is prefill-class
or decode-class, never both at once.

What consumes this profile: the post-capture serving grow floor
(:func:`arbi_serve.engine.inprocess_capture.serving_floor_for_grow`),
the optional activation arena's capacity, and telemetry. The KV budget's
own activation line comes from the separate single-shape probe in
:mod:`arbi_serve.engine.profile`, not from here.

``safety_fraction`` scales the serving floor's base term. It is an
operator knob for a workload whose live step outgrows every profiled
shape; the per-step transients that have a known cause are reserved by
their own explicit terms instead (the RoPE table in ``scratch.rope``, the
vocab-scale sampler tail in the serving floor's additive verify reserve).
First-call JIT / cuBLAS workspace has no term and is unreserved. The
serving boot passes ``cfg.memory_safety_fraction`` here, whose shipped
default is 0.0 — the ``0.10`` defaults on this module's own signatures
apply only to a caller that passes nothing.

Output shape
------------
Returns a :class:`MultiShapeProfile` carrying:

  - the byte peak per shape (with shape label);
  - the MAX across all shapes;
  - the safety margin in bytes and ``MAX + safety``, both reported for
    telemetry — the floor applies the fraction itself, over the shapes
    it actually sizes against.

The metrics module reads ``per_shape`` to emit
``arbi_serve.gpu.peak_activation_bytes_profiled{rank,profile_shape}``.

Device milliseconds spent between two CUDA events, or ``0.0``.

A zero is "not measured" — no CUDA device, or the event pair itself
raised. It is never a claim that the work was free, and every consumer
of :attr:`ShapeProfile.device_ms` treats it as a refusal to derive.

Time the enclosed device work with a CUDA event pair.

Wall time cannot stand in here. The probe's host thread also runs the
allocator's ``empty_cache`` and the trace replay, and the forward it
wraps is asynchronous, so a host clock measures launch and bookkeeping
as readily as it measures the step. The event pair measures the stream.

The timer never raises: a device that cannot record events yields a
zero, and the caller reads that as unmeasured.

Allocator events this device has recorded since the process started.

The counter behind the trace ring. ``free`` ticks the trace twice
(``free_requested`` + ``free_completed``), so the reading is doubled for
frees; over-counting is the safe direction for a capacity.

Returns 0 when the device is not an active CUDA device, so a CPU-only
caller reads a window of length 0 rather than raising.

Ring capacity for a window, derived from a MEASURED event count.

``reference_events`` is what the SAME shape emitted on its warmup pass,
counted with :func:`allocator_event_count`. The capacity is that number
plus one, because a window is only replayable while it is strictly shorter
than its ring — equality is indistinguishable from a wrap.

There is no headroom multiplier and no floor: every one would be a number
nobody measured. The measured pass is the COLD one and the recorded pass is
the warm one, so the recorded window is the smaller of the two by
construction; a warm pass that somehow exceeds its cold pass is an anomaly
the profile reports rather than absorbs.

0 means "nothing measured yet", which reads as "do not record" — the seam
that makes the first pass a measurement instead of a guess.

One allocator-trace window: its entries, or why it has none.

``entries`` is the replayable window; ``unmeasured_reason`` is set instead
whenever the window cannot be replayed. Exactly one of the two is
meaningful, and a caller that reads ``entries`` without checking the reason
gets an empty list — never a truncated one, which would report a peak
measured from the middle of a forward.

Token count of the pure-prefill synthetic shape.

Shape #1 is ONE sequence, and no sequence can be longer than its own
context window, so the probe is ``min(max_batched_tokens, max_context)``
— the same ``prefill_cap`` the capture ladder resolves with. This value
is the ONLY way ``max_context`` reaches any synthetic shape's token /
sequence / KV-length count: every other per-sequence length in
:func:`profile_engine` is ``min(x, max_context)`` for some ``x <=
max_batched_tokens``, which equals ``min(x, prefill_cap)``. The boot
manifest's fingerprint folds THIS number rather than ``max_context`` so
the key does not fork on an ``auto`` context the boot narrows.

``max_context`` of ``None`` / ``<= 0`` means "no context bound" — the
probe is then the full ``max_batched_tokens``.

Longest q-span ONE prefill row can carry at serving.

A prefill row is bounded by BOTH of the scheduler's per-row limits:
``chunk_prefill`` (every path sizes a row at ``min(remaining_prompt,
chunk_prefill, token_budget)`` — see ``Scheduler.schedule``, which chunks
recurrent rows too) and ``max_context`` (no sequence holds more KV than
its own window). ``max_batched_tokens`` bounds the STEP, not the row, so a
budget wider than this cap is reachable only by SEVERAL rows.

Either bound may be absent (``None`` / ``<= 0``), meaning unbounded; with
neither present the result is ``0`` == no bound.

Widest token count ONE step can actually carry.

The step budget is not reachable on its own. A step is built from rows,
every prefill row is capped at ``chunk_prefill``, and there are at most
``max_batch`` rows — so a step tops out at ``max_batch * chunk_prefill``
however large ``max_batched_tokens`` is. Probing the raw budget past that
measures a forward the scheduler cannot assemble, and the serving floor
then holds out a peak no request can reach.

Both bounds are real: ``max_batched_tokens`` because the scheduler stops
admitting rows at it, ``max_batch * chunk_prefill`` because it runs out of
rows to admit. The reachable width is the smaller. With no
``chunk_prefill`` the row cap is unknown and the budget stands alone.

Worst-case ``num_seqs`` the boot prefill probe builds.

The pure-prefill synthetic shape packs ``max_batched_tokens`` q-tokens —
a per-STEP aggregate. One row carries at most
:func:`prefill_row_cap`, so the probe spreads the tokens across
``ceil(max_batched_tokens / cap)`` sequences (q-span == kv-span per row),
which is how the scheduler reaches that budget: several competing
prefills, each capped at its own chunk. That row count is the ``num_seqs``
the live metadata builders' ``TQBufferPool.ensure`` sees at boot —
:func:`profile_engine`'s ``_forward`` derives its ``B`` from THIS function
so the two can never diverge, and ``capture_pool_num_seqs`` pre-sizes the
pool to cover it.

``chunk_prefill`` is keyword-only and required: a caller that omitted it
would silently model the whole budget as ONE row, which no scheduler path
can build, and size the serving floor from that unreachable peak.

One synthetic forward shape to profile.

Args:
    label: stable name (``prefill``, ``decode``, ``mixed``).
    num_tokens: total tokens in the synthetic batch (sum across
        sequences).
    num_seqs: number of sequences in the batch (decode = N rows
        of one token; prefill = 1 row of N tokens; mixed = a mix).
    decode_rows: how many of ``num_seqs`` carry exactly ONE token.
        The rest share the remaining tokens. ``0`` means "spread the
        tokens evenly", which is what the pure prefill and pure decode
        shapes want.

        This field is what makes a MIXED spec a mixed step rather than
        a wide uniform prefill: a fused mixed forward is a handful of
        single-token decode rows beside a chunk-sized prefill row, and
        its per-row transients (the GDN recurrent gather/scatter above
        all — it scales with ROWS, not tokens) are paid at the full row
        count while its per-token transients are paid at the chunk.
        Modelling the same tokens spread evenly measures neither.
    query_len_per_row: tokens EVERY row carries in a DECODE-CLASS batch.
        ``0`` (the default) means the shape is classified by the token /
        row ratio as it always was. A positive value is the verify width:
        ``num_seqs`` rows of ``query_len_per_row`` tokens, dispatched as
        ``is_prefill=False`` with ``max_query_len == query_len_per_row``,
        which is what :func:`~arbi_serve.spec_decode.mtp_verify_plan.
        build_verify_plan` stages and what the eager verify rung then runs.

        It has to be STATED rather than inferred, because the ratio alone
        cannot tell the two apart: a batch of ``B`` rows at ``K + 1`` tokens
        carries more tokens than rows and would classify as a prefill, which
        dispatches a different attention path and a different
        ``max_query_len`` from the forward this shape exists to measure.

Result of profiling one :class:`ShapeSpec`.

Attributes:
    peak_reserved_bytes: retained under this name for
        metric-cardinality stability; the *value* is the
        allocated-bytes peak inclusive of named-pool segments
        (default-allocator delta + sum-of-named-pool deltas). See
        module docstring for the rationale.
    peak_allocated_bytes: the global allocated peak. It already
        counts allocations made inside a private ``MemPool`` — the
        caching allocator stamps every block with the AGGREGATE stat
        type regardless of which pool served it — so this is a
        whole-process number, not a default-pool-only one.
    arena_peak_bytes: the ``scratch.forward_arena`` share of this
        shape's peak. It is ``peak_reserved_bytes`` less what the
        allocator trace replay saw peak OUTSIDE the pool, so the
        arena share and the off-arena residue are two halves of ONE
        total rather than two instruments' readings — the trace
        supplies the split at the peak MOMENT and nothing else.
        ``None`` when the trace was unavailable or wrapped, which
        callers must treat as "the whole peak may be non-arena".
    arena_residue_bytes: the same pool's allocated delta at the END
        of the forward. It reads ~0 once the transients are freed;
        it is carried beside the high-water so a boot where the two
        diverge is visible rather than inferred.

The arena's share of a shape peak could not be measured.

The serving floor charges a probe's bytes to exactly one term: whatever the
arena high-water accounts for goes to ``arena_regrow``, and only the
remainder stays in ``base``. Without the split nearly every byte stays in
``base`` AND is reserved again by ``arena_regrow``, i.e. the floor holds one
forward's activations twice.

Continuing without the measurement is therefore not a safe default, it is
the defect this split exists to remove, re-applied silently. Fail here so a
boot that cannot measure says so, rather than quietly sizing the KV pool
from a number known to be wrong.

A profile probe could not be measured, so no floor can be derived.

Raised instead of recording a zero peak: the serving floor is a max over
the profiled shapes, and a zero for a shape serving can build makes that
max an under-count rather than a bound.

The widest FUSED mixed step the scheduler can hand one forward.

A mixed step is decode rows co-admitted with a chunked prefill row, and
the engine runs it as ONE forward whenever
:func:`~arbi_serve.runtime.split_mixed.is_splittable` declines the split
— which is every step on an MTP deployment (a decode row there is a
multi-token verify row, refused by name), and every step inside the cost
gate's threshold on any deployment. That fused forward is therefore a
shape production builds constantly, and it is the shape this spec makes
the boot profile measure.

Its two independent maxima are reached TOGETHER, which is why one spec
covers both:

* ROWS. The scheduler admits at most ``max_batch`` rows, so a fused
  mixed step tops out at ``max_batch`` rows — one (or a few) prefill
  rows and the rest single-token decode rows. Per-row transients scale
  with this number: the GDN / Mamba recurrent gather and scatter index
  one slab row per sequence, the page table is ``rows x pages``, and
  every varlen metadata builder allocates per row.
* TOKENS. The step still carries :func:`reachable_step_tokens` tokens,
  essentially all of them in the prefill row(s), so the chunk-scale
  transients (attention scratch, the MLP intermediate) are unchanged
  from the pure-prefill probe.

A fused step at ``R`` rows and ``N`` tokens therefore holds ``R`` rows'
worth of per-row state and ``N`` tokens' worth of per-token state at the
same instant. Neither of the other two probes does: the prefill probe is
``ceil(N / row_cap)`` rows (ONE row at the shipped
``max_batched_tokens == chunk_prefill``), and the decode probe is
``max_batch`` rows of one token, which is a DECODE-class forward — a
different kernel set, no chunked-prefill scratch, and
``max_query_len == 1``.

Composition. ``P`` prefill rows is the FEWEST that can carry the step's
prefill tokens once the other ``R - P`` rows are single-token decode
rows, at :func:`prefill_row_cap` tokens per row — the same per-row bound
the scheduler chunks against. Fewest, because that is what the scheduler
produces (it fills a row to its chunk before opening another) and it is
also the composition that leaves the most rows decoding, which is the
axis this spec exists to reach.

Every input is read, none is assumed: ``max_batch`` sets the row
ceiling, ``max_batched_tokens`` and ``chunk_prefill`` set the reachable
token width through :func:`reachable_step_tokens`, and ``max_context``
joins ``chunk_prefill`` in the per-row cap.

The synthetic shapes the profile measures, as a pure function.

The serving floor is sized from these peaks, so the widths ARE the
contract and are testable without running a forward.

  prefill: the reachable step width, spread over the rows it takes;
  decode:  ``max_batch`` sequences x 1 token;
  mixed:   the widest FUSED mixed step — ``max_batch`` rows carrying the
           reachable step width, decode rows and prefill row(s) in one
           forward (:func:`fused_mixed_shape`);
  verify:  ``max_batch`` decode-class rows x ``mtp_k + 1`` tokens, the
           width every served MTP verify step runs. Emitted only when
           ``mtp_k`` is given — at ``K == 0`` it is the decode shape.

The verify shape is the ONLY one that runs a decode row at more than one
token, which is the axis both consumers of this profile were short on. It
is capped at ``max_batch`` rows and not at the step token budget: a verify
step's width comes from the draft depth, so its ``B x (K + 1)`` tokens are
reachable whether or not that product fits ``max_batched_tokens``, and
clamping it to the budget would measure a narrower forward than the one the
driver builds.

Widths come from :func:`reachable_step_tokens`, not from
``max_batched_tokens`` directly. ``chunk_prefill`` caps ONE request's row
and the scheduler bounds the STEP by the budget, so competing prefills
co-batch up to it and decode rows ride along — but only ``max_batch`` rows
exist, so the step also stops at ``max_batch * chunk_prefill``. Probing
narrower than the reachable width measures a step smaller than serving can
build; probing wider measures one it cannot build at all, and the floor
then holds a peak out of the KV pool for a forward no request reaches.

Run the canonical synthetic shapes; return MAX + safety.

Args:
    forward: callable that runs one synthetic forward at the
        given :class:`ShapeSpec`. Owns its own batch construction
        (the profiler doesn't know the model). Must be a real
        ``model.forward`` against a real synthetic
        :class:`ScheduledBatch` so the caching allocator sees the
        same alloc / free pattern as production.
    max_batched_tokens: ``cfg.batch.max_batched_tokens`` —
        shape #1's token count.
    max_batch: ``cfg.batch.max_batch`` — shape #2's seq count.
    max_context: ``cfg.batch.max_context``. Shape #1 is ONE sequence, and
        a sequence cannot exceed its own context window, so the prefill
        probe is ``min(max_batched_tokens, max_context)`` — the same
        ``prefill_cap`` clamp the capture ladder resolves with. Optional
        only for callers that have no context bound; passing it is what
        keeps the probe inside the served context window.
    chunk_prefill: ``cfg.batch.chunk_prefill``; defaults to
        ``max_batched_tokens // 2`` when None — shape #3's
        prefill-chunk size.
    mtp_k: this boot's resolved MTP draft depth
        (:func:`verify_forward_geometry`), or None / 0 when the engine
        does not speculate. When given, a fourth shape runs the SAME
        ``forward`` callable at ``max_batch`` decode-class rows of
        ``K + 1`` tokens — the width every served verify step runs and the
        one no other shape reaches. It goes through ``forward``, not
        ``verify_forward``: it is a real ``model.forward``, where the
        stochastic tail beside it is the rejection sampler.
    device: the CUDA device the forward runs on.
    safety_fraction: extra margin as a fraction of the MAX peak,
        reported on the profile and applied by the serving floor to
        its base term. The serving boot passes
        ``cfg.memory_safety_fraction`` (shipped default 0.0), so this
        signature's ``0.10`` reaches only a caller that omits it.
        Range 0.0–0.5; values outside are clamped with a warning.
    named_pools: optional :class:`NamedPoolRegistry` (or any
        object exposing ``snapshot_all() -> list[dict]`` where
        each dict carries ``caching_allocator_allocated_bytes``).
        When provided, each shape's peak is augmented with the
        sum of named-pool allocated-byte deltas observed while
        the synthetic forward ran. Without it, per-step activations
        that land in a named pool (e.g. the ``activation_arena`` or
        ``tkv_scratch_pool``) would be missing from the peak
        entirely.

Returns:
    :class:`MultiShapeProfile` with per-shape peaks + the MAX +
    the chosen safety margin in bytes.

Profile one shape as a SERVED step, with a trace ring it measured itself.

Two passes, always:

  1. **Warmup, unrecorded.** Its only outputs are the shape's cold-call
     event count — which is what sizes pass 2's ring — and the cold peak,
     kept for the receipt.
  2. **Measured, recorded** at exactly the count pass 1 observed.

Why the first pass is not the measurement
-----------------------------------------
The profile's decode shape is the FIRST decode-class forward of a boot, so
the TKV split-K decode autotune fires lazily inside it. MEASURED on a 27B
GDN hybrid at ``chunk_prefill=2048``: the decode probe emitted **4,452,533
allocations in 347.9 s**, of which 95.5% came from ONE call site
(``tkv/kernels/turbo_attn_simt.py`` under
``tkv/runtime/_autotune_timing._time_decode_call``), against 5,276
allocations in 1.23 s for the 2048-token prefill probe and 5,276 for the
mixed probe on the same boot. The count is a property of the autotune
sweep, not of the slate — which is why halving ``chunk_prefill`` did not
move it, and why a ring of ``1 << 20`` and one of ``1 << 23`` both filled.

Two things follow, and both are why the warmup pass exists rather than a
bigger ring:

  * the ring would have to hold millions of entries that describe a kernel
    SEARCH rather than a step, and
  * ``max_memory_allocated`` over that window counts the search's candidate
    buffers, so the decode shape's "activation peak" is the autotune's
    scratch. The serving floor is a MAX over the decode / mixed / verify
    shapes — the prefill peak is explicitly excluded as reclaimable — so
    that scratch is held out of the KV pool for the process's life.

Serving never pays it: the boot warms the same table explicitly at Phase 4d
(``_warm_decode_autotune``) and BOOT-FAILS if a warmable config does not
warm, so a served decode step runs against a resolved table. Measuring the
cold call measures something serving cannot do.

The peak carried forward is pass 2's — the served step — with the cold peak
recorded beside it so a boot where the two diverge says so out loud instead
of being inferred from a KV pool that came out small.

Peak allocated bytes over one allocator-trace window, split by pool.

Returns ``(total_peak, off_pool_peak, pool_peak)``, each a byte delta over
the window's start. The caching allocator records an ``alloc`` /
``free_requested`` pair per tensor carrying the id of the pool that served
it, so replaying the window in order reconstructs both curves and lets each
be maximised at its OWN moment — which is the whole point: the two peaks do
not coincide, and a subtraction of one maximum from the other is not the
maximum of the difference.

Sizes come from the trace as REQUESTED bytes, while the allocator's stats
count the rounded block, so every one of these three numbers runs slightly
under the ``max_memory_allocated`` delta for the same window. That is why
the caller uses ``off_pool_peak`` alone and subtracts it from the peak the
COUNTERS reported (:func:`_arena_share_from_trace`): a split taken from
this function's own total would be a second reading of a quantity the
budget already has, and the difference between the two instruments would
land on whichever consumer subtracts one from the other. There is no
conservative direction for that difference — one consumer holds the arena
share back and another spends it — so it is not allowed to exist.

Pure function over the trace dicts — no CUDA, no allocator state.

Record the allocator's per-allocation trace for the duration.

Yields a :class:`TraceWindow` filled in on exit. It carries the entries
when the window is replayable, and an ``unmeasured_reason`` when it is not
— the recorder was unavailable, the snapshot raised, or the ring filled (a
wrapped ring has lost its head, so replaying it would report a peak
measured from the middle of a forward). Callers treat a reason as
"unknown", never as "zero", and the reason is what the budget path later
names so an operator is never told the card is full when it is not.

``capacity`` is derived from a MEASURED event count
(:func:`trace_capacity_for`); a capacity of 0 means nothing has been
measured yet, which is itself a reason rather than a silent skip. The
window also records its OWN event count from the O(1) counters, so an
overflow is reported with the exact number it needed rather than with the
number it was allowed to keep.

The recorder is process-global and this clears it on entry, so it runs at
boot around the shape probes only.

The arena's share of one probe's peak, and why it is not known.

Returns ``(share, unmeasured_reason)`` with exactly one side populated.
``window`` is what :func:`_alloc_trace_window` handed back: a replayable
window, or a reason it is not one. A reason means UNKNOWN — the caller must
then charge the whole peak to the non-arena side — which is why the share
is ``None`` rather than 0 on every failure path, and why the reason travels
with it instead of being logged and dropped.

ONE TOTAL, SPLIT ONCE. ``peak_total`` is the shape's peak as the caller
reports it and as every consumer prices it — the allocator's own counters,
which count the ROUNDED BLOCK the allocator handed out. The replay measures
the same forward from the trace, whose sizes are the REQUESTED bytes
(:func:`split_trace_peak` says so), so its total runs under the counters'
for the same window. The trace is therefore asked for the SPLIT alone: what
peaked OUTSIDE the pool. The share is the caller's own total less that, so
``peak_reserved_bytes == arena_peak_bytes + off-arena residue`` holds by
construction.

Returning the replay's own total as the share instead made the two halves
of one step's accounting readings from two instruments: the serving floor
and the admission gate price a step from the counters
(:meth:`StepActivationModel.worst_admissible_bytes` fits them exactly) and
credit the arena from the replay, so the block-rounding difference between
the instruments became a shortfall no term named — spent as prefill chunk
width every step. Which instrument reads larger is not the point; that two
of them meet inside one subtraction is.

Sum ``caching_allocator_allocated_bytes`` across every NamedMemPool.

Returns 0 when ``named_pools`` is None / lacks ``snapshot_all``.
Best-effort: a snapshot exception is logged and treated as 0 so
the profile path never crashes on diagnostics.

Run one synthetic forward, return its peak allocated bytes.

``max_memory_allocated`` is a WHOLE-PROCESS counter: the caching
allocator stamps every block with the AGGREGATE stat type whatever
pool served it, so a private ``MemPool``'s allocations are already
inside it. The registry delta added on top is taken at the END of
the forward and reads ~0 once the transients are freed; it is kept
as a floor under the reading and reported separately so a boot
where it is NOT ~0 is visible.

Empties the caching allocator BEFORE each profile so the
counter starts from a clean slate — without this, segment bytes
from a prior shape stay cached and would inflate the named-pool
delta for the next shape.

THE ARENA SPLIT. The peak this returns is one number over two
allocators: the probe forward runs inside ``scratch.forward_arena``
(:func:`profile_engine`), and that pool is reserved in full by its own
serving-floor term. Charging its bytes to the floor's base term as well
reserves them twice. The end-of-forward pool delta cannot separate them —
it is a residue, not a high-water — so the window is replayed from the
allocator's per-allocation trace, which carries the owning pool id on every
event, and the arena's share is maximised at its own moment
(:func:`split_trace_peak`). A shape whose trace is unavailable carries
``arena_peak_bytes=None`` and its whole peak stays in the base term.

The legacy ``peak_reserved_bytes`` field name is preserved on
:class:`ShapeProfile` for metric-cardinality stability; the value
it carries is the allocated-bytes-including-named-pools number.

The engine's PERSISTENT verify operands at the ``(K, B)`` slate, or None.

A served stochastic verify step reads its ``(K+1, B, V)`` logits out of
:class:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers`'s persistent
gather scratch and writes ``p_target`` / ``draft_probs`` / Gumbel noise into
the persistent sampler scratch, so none of those vocab-scale tensors is a
per-step allocation. Returning them here makes the profiled shape measure
what the served step actually allocates FRESH (the fp32 upcast and the
sampling-params chain intermediates) instead of the fresh-allocation
fallback the pool replaced.

Returns None when the pool is absent or was not built at this exact
worst-case slate (``vocab_size == 0`` stubs, a narrower pool), so the caller
keeps the conservative fresh-allocation synthetic.

A pool with no dense-``q`` scratch is a boot whose drafters all propose a
point mass, so the served ``q`` operand is
:data:`~arbi_serve.spec_decode.drafter.POINT_MASS_Q` — handing the profile a
synthetic dense ``q`` there would measure a tensor the step never builds and
re-book the very bytes the pool skipped.

``(K, V)`` of the stochastic verify tail, or ``None`` when it is not profiled.

THE single expression for "does this boot measure
:data:`SHAPE_VERIFY_STOCHASTIC`". Two readers share it and must not drift:
:func:`_build_verify_forward`, which decides whether the shape is measured,
and :func:`arbi_serve.runtime.boot_manifest.build_fingerprint`, which must
key the manifest on that decision.

They have to be one expression because the manifest caches the profile's
SHAPE SET, not just its numbers. When only the builder knew the condition,
``rejection_sampling_enabled`` could turn the shape off without moving the
key: a boot with it OFF persisted a profile with no verify tail, and a later
boot with it ON hit that entry, read ``stochastic_measured == 0``, and sized
the serving floor from the config-dims ``analytic`` fallback instead of the
measurement (:mod:`arbi_serve.engine.inprocess_capture`), which over-reserves
against the measured tail and takes the difference straight out of KV.

Both callers run at the same seam on the same engine
(``_build_activation_profile`` builds the fingerprint and profiles
back to back), so the live drafter depth this reads is the same one the
profiled shape is built at.

The loaded ``lm_head``'s geometry, or ``None`` when it cannot be read.

THE single reader of the head for memory purposes, so the boot floor's
verify-slate row and the admission gate's per-drafted-position charge
cannot describe different heads. Both go through
:func:`~arbi_serve.engine.memory_budget.pool_predictors.
verify_logits_epilogue_peak_bytes` with what this returns; they differ only
in the row count they ask about.

Reads the MODULE first and the config dims only where the module does not
declare a width — a tied head projects through the embedding and has no
``out_features`` of its own, and its vocab is the model's by construction.

``None`` means the head is unreadable, which is a floor that cannot be
computed rather than a head that is free: callers must fall back to a
measurement or leave their existing number standing, never to 0.

This boot's MTP draft depth ``K``, or ``None`` when there is no verify width.

THE single expression for "does this boot measure
:data:`SHAPE_VERIFY_FORWARD`", for the same reason
:func:`verify_tail_geometry` is the single expression for its sibling: the
boot manifest caches the profile's SHAPE SET, so the predicate that decides
whether the shape is emitted and the one that keys the cache have to be one
function or a boot can hit an entry that measured a different set.

Deliberately WEAKER than :func:`verify_tail_geometry`, on both of the
conditions that one adds. It does not require ``rejection_sampling_enabled``
— that flag selects the SAMPLER, and this shape is ``model.forward``, which
a greedy verify step runs at exactly the same width. And it does not require
a vocab size, which the main forward's activation footprint does not depend
on. What it does require is a resolved ``K >= 1``: at ``K == 0`` the engine
does not speculate, every decode row is one token, and the shape would
duplicate ``SHAPE_DECODE``.

Reads the LIVE drafter's depth ahead of the config, the same order its
sibling reads it in, so the profiled width is the width the driver will
build (the K-calibration sweep can resolve a depth below ``n_draft``).

Rows ONE verify step's ``lm_head`` call is handed, or 0.

``max_batch x (K + 1)``: the scheduler buckets a verify slate to a single
``step_k`` (``_bucket_mtp_k_uniform``) and every row in it contributes its
committed tail plus ``K`` drafts, ALL of which
:mod:`arbi_serve.spec_decode.mtp_verify_offload` pushes through the head in
one call. Concurrency is what widens it — which is why the reserve cannot
be sized from the single-stream step that boots the server.

``K`` comes from :func:`verify_tail_geometry`, the same predicate that
decides whether the verify tail is measured at all, so a boot cannot reserve
for a head call it will never make. The head's own widths come from
:func:`logits_epilogue_geometry`, which reads the module that is loaded.

The closed-form ``lm_head`` epilogue peak at ``rows``, off the LOADED head.

0 when the head cannot be read
(:func:`~arbi_serve.runtime.activation_profile.logits_epilogue_geometry`),
which is a bound that could not be computed — never a head that is free.
Callers pair it with a measurement and take the larger.

The shape labels :func:`profile_engine` will emit for ``eng``, in order.

The three base shapes are unconditional; the verify tail is not. Folding
this list into the boot-manifest fingerprint keys the cache on the profile's
SHAPE SET, so a config that measures a different set of shapes gets its own
entry instead of hitting one that measured another set.

Returning the SET rather than the individual predicates is what closes the
class: any future conditional shape joins the key the moment it is emitted
here, with no second place to remember to update.

Build the synthetic stochastic-verify-tail forward, or None.

Returns a callable that runs ``verify_and_accept(greedy=False, ...)`` at the
serving max shape (B=max_batch, the configured MTP draft depth K, the model
vocab V) so the profiler measures the rejection sampler's per-step peak.
Returns None — and the profiler then skips the shape entirely — on any of
the conditions in :func:`verify_tail_geometry`: spec-decode off,
``rejection_sampling_enabled`` off, or K/V unknown. That the SHAPE SET turns
on a config flag is why the manifest fingerprint keys on
:func:`expected_shape_labels` rather than on the numbers alone.

Measures the SERVED allocation, not a stand-in for it: when the engine's
persistent verify / sampler scratch exists it is threaded in exactly as
:mod:`arbi_serve.spec_decode.mtp_verify_accept` threads it, so the profiled
peak is the tail's FRESH per-step footprint on top of a pool that is already
resident before the KV grow measures free VRAM. Without that pool the
sampler really does allocate every vocab-scale tensor per step, and the
fresh-allocation synthetic measures that.

Both rejection-tail branches run under one profile so the recorded peak is
the MAX over them: the fused single-pass path (``ARBI_MTP_FUSED_REJECTION``,
which scales the whole ``(K+1, B, V)`` block) and the materialized
``p_target`` chain (whose sampling-params intermediates are slot-chunked
into the persistent destination). Either is reachable at serve time, so the
floor must cover the larger.

The engine's RESOLVED context ceiling (``cfg.cache.max_context``).

``"auto"`` is turned into the model ceiling by
``build_config_resolve._resolve_auto_max_context``, which runs before the
activation profile. Anything else here means the profile is running against
an unresolved config — fail loud rather than return a value that reads as
"no context bound", because the caller's only use for this is a clamp whose
absence is an illegal memory access.

The batch-side ``MTPMeta`` a verify-width probe must carry, or ``None``.

NOT bookkeeping the probe could skip. Three recurrent backends resolve
``meta.verify_pass = (not batch.is_prefill) and (batch.mtp_meta is not
None)`` (GDN, Mamba, DSv4), and the captured-graph lookup gates refuse an
``S > 1`` batch that carries no ``mtp_meta`` as "something else — a chunked
prefill". A verify-width probe without one is therefore not a cheaper
verify step: it is a DIFFERENT forward, dispatched down the recurrent
half's non-verify path at ``B x (K + 1)`` tokens, and the peak it measures
would be a peak serving never reaches.

The three fields are exactly what :meth:`~arbi_serve.spec_decode.mtp_meta.
MtpStepPlan.as_batch_meta` hands the batch for a UNIFORM-K slate, which is
the only kind admission builds (``_bucket_mtp_k_uniform``): every row takes
the same ``K``, its leading flat token is the committed tail and the ``K``
behind it are drafts.

``None`` for every other shape, which leaves their batches exactly as they
were.

The synthetic batch's per-row q- and kv-spans for one shape.

Returns ``(num_rows, q_lens, kv_lens)``. Pure, so the widths the boot
probe forwards ARE the contract and are testable without a GPU or a
built engine — which matters because a probe that quietly builds a
narrower batch than its spec asked for is invisible in every byte the
profile reports.

Row count. ``prefill_probe_num_seqs`` is the FLOOR, not the answer: a
row cannot carry more than :func:`prefill_row_cap` tokens, so a shape's
tokens need at least that many rows. A spec asking for MORE rows is
asking for a wider slate — which is what the mixed shape does — and
taking the floor as the answer collapses it onto the prefill probe's
row count. At the shipped ``max_batched_tokens == chunk_prefill`` that
count is ONE.

Composition. ``spec.decode_rows`` rows carry exactly one token each and
the prefill rows share the rest; that is how the scheduler builds a
fused mixed step and it is what makes the per-row and per-token
transients co-resident in one forward. A spec with no decode rows
spreads its tokens evenly, with the last row absorbing the remainder —
what the pure prefill and pure decode shapes want.

KV spans. A fresh prefill row's kv-span equals its q-span (the Turbo
prefill paged loader derives ``S_kv`` from the per-seq q-span, and a q != kv
batch indexes it out of bounds). A decode row is modelled at ONE KV
token, the same way the decode probe models its rows, so the mixed shape
adds exactly one variable to the profile: the ROW count.

The even-spread branch keeps its ``seq_lens`` UNIFORM at ``N // B`` even
where the last row's q-span absorbs the token remainder: that is the
kv-span every kernel on this path reads, and widening it would change
what the pure-prefill probe measures.

A VERIFY-WIDTH spec (``query_len_per_row > 0``) is uniform by construction:
every row carries the same ``K + 1`` tokens, which is what
``build_verify_plan`` stages — the slate is uniform-K by admission
(``_bucket_mtp_k_uniform``), so a ragged verify batch is not a shape the
scheduler can present. Its kv-span equals its q-span for the reason a fresh
prefill row's does: the paged loader derives ``S_kv`` from the per-seq
q-span. That is deliberately NOT the served row's accumulated context — this
shape isolates the QUERY-WIDTH axis, which nothing else prices; the KV axis
is already carried by the model's gather term and the floor's tkv staging
term, and widening it here would book those bytes a second time.

Profile the three canonical shapes against a built engine.

Synthesizes a :class:`ScheduledBatch` for each shape using
placeholder token / position / slot data (page id 0 throughout —
the same null-page convention CUDAGraph capture uses), runs
``model.forward`` under the eager path, and reports peaks.

Skipped when CUDA is unavailable (returns an empty profile so
callers can still log).

The batch-level ``max_query_len`` a forward at this shape must set.

A decode row is one token unless the shape states a verify width; a
prefill row is its own span, which only the row plan knows, so this
returns ``0`` there and the caller supplies the measured span.

This shape's peak MINUS the part of it the forward arena held.

``arena_covered_bytes`` is what the caller's OWN arena term already
holds free; the subtraction is capped by it so the two terms together
can never come out below the measured peak. An unmeasured arena share
subtracts nothing.

Shapes whose arena split could NOT be measured, in profile order.

Non-empty means every one of these shapes kept its WHOLE peak in the
serving floor's base term while ``arena_regrow`` reserved the same
pool again — the floor holds one forward's activations twice, and any
KV ceiling derived from it is an artefact of the missing measurement
rather than a statement about how full the card is. The budget path
reads this to keep those two facts apart.

One paragraph naming the probes that came back unmeasured, or ``""``.

THE single rendering of the unmeasured fact. Two boot refusals quote it
(the KV budget's and the deferred resize's), and they must not drift:
an operator who reads one and then the other has to see the same cause
named the same way.

Async output processing — the deferred sampled-token host copy.

The decode step samples on the GPU; instead of blocking the step on
``samples.cpu().tolist()``, the sampled-token tensor is copied to host on a
dedicated side stream with a recorded CUDA event, and the host ``.tolist()``
plus the per-request commit / ``post_token`` are deferred by one engine tick.
The forward of step N+1 does not wait on step N's output copy.

The just-sampled ids are kept on the GPU (``Request.last_sampled_gpu``) so the
next step's ``input_ids`` slot is written by a D2D copy with no host
round-trip.

When CUDA is unavailable the copy stream + event are no-ops and the held
tensor is materialized lazily on the deferred drain; the deferral ordering,
stop-exactness, and abort handling are identical to the CUDA path.

The MTP / spec verify path is not deferred — it runs synchronously and any
pending K=1 output is drained before it commits so per-request token order is
preserved. See ``_async_output_enabled`` callers.

One depth of a :class:`HostCopyRing`.

Owns pinned host staging tensors (one per distinct ``(name, shape,
dtype)`` the egress asks for, allocated on first use and reused after)
and one CUDA event, re-recorded per use. ``held`` is set from
:meth:`HostCopyRing.acquire` until the consumer's ``materialize`` has
read the staged bytes and called :meth:`release`. ``holder`` is a weak
reference to the object that will consume the slot (the parked step,
or the handles before park): a step that raised between issue and
drain drops its holder without consuming, and the ring reclaims such a
slot instead of refusing every later step.

Depth-N ring of reusable host staging tensors and events for one egress.

Replaces the per-step ``torch.empty(..., pin_memory=True)`` and
``torch.cuda.Event()`` of the deferred D2H with slots that live for the
runner's lifetime. A slot is acquired when a copy is issued and released
by the consumer's ``materialize`` after the host read; acquiring a slot
that is still held raises instead of overwriting, because a staging
tensor overwritten before its read would commit another step's tokens.

A step's sampled tokens whose host materialization is deferred.

Holds the ``(B,)`` int64 device tensor (``samples_gpu``) alive across
the deferral so the allocator can't reclaim it before the copy lands,
the host-side staging tensor the copy targets, and the CUDA event
recorded on the copy stream. On CPU ``event`` is ``None`` and
``samples_cpu`` is materialized on demand.

``ready_indices`` maps each entry of the materialized id list back to
its slate row (the ``_sample`` "ready" rows); rows not in it sampled
no token this step (mid-prefill) and stay ``None``.

Issue the non-blocking D2H of the sampled ids on the side stream.

The copy stream waits on the compute (default) stream so the D2H is
ordered after the sample kernel, then the copy + event are issued
non-blocking — the host returns immediately. No ``.tolist()`` here.

With ``ring`` the staging tensor and event come from a
:class:`HostCopyRing` slot the returned sample releases on
``materialize``; without one they are allocated per call.

On CPU there is no stream; we just hold the tensor and defer the
``.tolist()`` to ``PendingSample.materialize`` (same ordering, no
GPU).

A whole deferred K=1 step: its samples + the commit context.

The run loop holds at most one of these. At the top of step N+1 —
after N+1's forward + sample is launched — the loop drains step N:
block on the copy event, materialize the ids, run ``scheduler.commit``
+ ``post_token``. A request finished (stop / max_tokens / abort)
during the drain has its further deferred output dropped.

A whole deferred MTP verify step: its acceptance + the commit ctx.

Held by the engine (``eng._pending_verify``) for exactly one tick.
Structurally parallel to :class:`PendingStep` for the K=1 path, but the
deferred host materialization is the verify acceptance — the greedy
op's device ``n_accepted`` ``(B,)`` and ``main_argmax`` ``(K+1, B)`` —
rather than a single sampled-token vector.

Fields:
  ``slate``: the verify sub-slate (request, n_tokens) in row order.
  ``n_accepted_gpu`` / ``main_argmax_gpu``: the greedy verify op's
    device outputs, kept alive across the deferral so the allocator
    can't reclaim them before the copy event lands.
  ``n_accepted_cpu`` / ``main_argmax_cpu``: the host staging tensors
    the side-stream D2H targets (``None`` on CPU — materialized
    lazily from the device tensors on drain).
  ``event``: the CUDA event recorded on the copy stream (``None`` on
    CPU).
  ``eff_k``: the uniform per-row spec depth K (the optimistic
    advance appended ``K + 1`` placeholders per row).
  ``optimistic_appended``: per-row count of placeholder tokens
    appended at park (``K + 1`` for a committed row, ``0`` for a row
    that finished before park / wasn't advanced).

One verify step's acceptance D2H, already in flight on the copy stream.

Exists so the copy can be ISSUED at the point the acceptance becomes
final and PARKED later. The two are different moments: the copy stream's
``wait_stream(compute)`` captures whatever is queued on the compute
stream *at issue time*, so issuing at park chains the copy behind the
whole remainder of the step (drafter chain, recurrent rollback) even
though its inputs were ready long before.

Issue the acceptance D2H the moment the acceptance is final.

Call this at the last write to ``n_accepted`` / ``committed``, NOT at
park. Everything the step enqueues afterwards (drafter chain, recurrent
reconcile) then overlaps the copy instead of gating it, and the drain's
event wait collapses to the copy's own latency.

Staged into PINNED host memory so ``non_blocking=True`` is a true async
D2H rather than a driver-staged pageable copy. With ``ring`` the staging
tensors and event are a :class:`HostCopyRing` slot the parked step
releases on ``materialize``; without one they are allocated per call.
Returns ``None`` on CPU or with no copy stream — the caller falls back
to holding the device tensors and materializing lazily.

Issue the non-blocking D2H of the verify acceptance on the side
stream and park a :class:`PendingVerifyStep`.

Structurally parallel to :func:`issue_output_copy`, with one difference
that matters: the copy should already be IN FLIGHT by the time we park.
Pass the ``handles`` returned by :func:`start_verify_copy` at the point
the acceptance became final; only a caller with no earlier issue point
leaves them ``None``, and then the copy is issued here as before. No
``.tolist()`` either way — the host pull is deferred to
:meth:`PendingVerifyStep.materialize` in the drain.

On CPU there is no stream; the device tensors are held and
materialized lazily on drain (same ordering, no GPU).

Take the next slot; raises if it still holds an unconsumed copy.

A slot whose holder was dropped without consuming (a step that
raised after issuing its copy) is reclaimed: nothing can read those
bytes any more, so reusing them commits nothing stale.

Release every held slot whose holder is not one of ``live``.

The failure path's counterpart to :meth:`acquire`'s leak check: a
step that raised after issuing its copy never parks the object that
would have consumed the slot, so the slot stays held by an orphan
for as long as anything references it. ``live`` is the set of
objects the engine still intends to drain; a slot held by anything
else — or by nothing — is handed back. Returns the count released.

Block on the copy event, then return per-row sampled ints.

The event sync only stalls if the copy hasn't finished — by the
time the deferred drain runs, the next step's forward has been
enqueued, so the copy overlapped it and this is (near-)free.
Returns a length-``n_rows`` list with ``None`` for non-ready rows.

This row's accepted draft-slot ordinals, or ``None`` (chain).

A method rather than a raw index so the chain's ``None`` — by far
the hot case — costs one attribute load at the call site instead
of a conditional the caller has to get right twice.

Block on the copy event, then return ``(n_accepted, main_argmax)``.

``n_accepted`` is a length-B host int list; ``main_argmax`` is a
``(K+1)``-by-``B`` host int matrix (committed-token candidates).
A TREE step also fills :attr:`accepted_paths` here — from the same
host pull, so the paths and the counts they are sliced against can
never come from different steps.
Unlike :meth:`PendingSample.materialize`, this sync is NOT
overlapped by a following forward: :meth:`MtpStrategy.run_step`
drains at the head of the step, before the next verify plan is
built (the plan needs the corrected lengths). Releases the device
tensors once the host copy has landed.

Boot guards for the native backends the running source expects.

TWO GUARDS, because a dependency is pinned in two different currencies and a
mismatch in each has a different meaning.

VERSION FLOOR (:func:`check_turbo_attn_compatible`)
    Published images resolve every package exactly from ``uv.lock`` and the
    image build verifies that resolution. A source bind mount may carry a newer
    lock than the image while remaining compatible with its installed backend.
    Runtime source therefore checks the minimum version declared by project
    metadata rather than requiring the image to reproduce the bind-mounted lock.

    THE FLOOR MUST BE READABLE WHEREVER THE ENGINE RUNS. A shipped image has no
    ``pyproject.toml`` beside its ``site-packages`` — walking up from this file
    finds nothing there — so a lookup that can only read that file has no answer
    in the deployment this guard exists for, and no answer reads exactly like a
    satisfied floor. The installed ``arbi-serve`` distribution freezes the same
    requirement at build time, so it is the fallback; a floor that neither
    source yields is refused, not skipped.

    ``tkv`` REPORTS A VERSION IT MAY NOT OWN. ``tkv.__version__`` prefers the
    ``tkv/_version.py`` its build wrote and otherwise falls back to the
    ``turbo-attn`` dist-info. A bind mount replaces the package DIRECTORY and
    leaves that dist-info behind (see GIT REVISION below for the same reading
    made about the fork), so for a mounted checkout with no build stamp the
    number is the IMAGE's and says nothing about what will run. Which of the two
    answered is therefore part of the reading, and it decides whether the floor
    comparison is evidence about this tree at all.

    SO THE TREE IS ALSO ASKED WHAT IT DECLARES. The floor is a number and a
    checkout's number can be stale in either direction; the failure itself is a
    name ``arbi_serve`` imports at module scope that the mounted tkv does not
    have, which surfaces as an ``ImportError`` from deep inside an unrelated
    module long after the mount that caused it. Those names are read off the
    tree's own sources — never imported, since importing a tkv kernel module
    triggers a JIT build — so the reading costs no CUDA and holds for a tree
    whose version nothing records.

GIT REVISION (:func:`check_pinned_revisions`)
    A revision is not ordered, so there is no floor to be above: a dependency
    resolved from a git pin either is the commit this source tree was written
    against or it is not. Nor can a version string stand in — the exllamav3
    fork and the PyPI wheel report the same ``__version__``, which is why the
    image build stamps the resolved commit (see the Dockerfile's provenance
    stamp) instead of trusting the version.

    This is the guard for the rig everyone uses here: a CURRENT source tree
    bind-mounted into a PUBLISHED image. That rig deliberately leaves the
    image's git sha differing from the tree's — the sha therefore says nothing
    about staleness — while the image's native closure is frozen at whatever
    the tree pinned on the day it was built. When the tree later moves a pin,
    the mismatch surfaces as an opaque ``TypeError`` from a fork call deep in a
    forward pass, long after the boot has spent its expensive minutes. The
    revision comparison is available in the first second of the boot, so it is
    made there.

    The comparison only needs facts that TRAVEL WITH THE MOUNT: the required
    revision is a constant inside ``arbi_serve`` (single-sourced against
    ``pyproject.toml`` by ``tests/test_exl3_dep_managed.py``). Nothing is read
    from a ``uv.lock`` beside the source, because the bind-mount overlay mounts
    the package directory alone (``compose.bindmount.yaml``) and no lock is
    there to read.

    WHICH ARTIFACT IS MEASURED. The thing that runs is the package DIRECTORY on
    ``sys.path``: the fork's Python modules are imported from it, and its
    ``.cu``/``.cpp`` sources are what ``exllamav3_ext`` JIT-compiles (torch's
    ``build.ninja`` names every translation unit by absolute path inside it, so
    that is checkable rather than asserted). ``dist-info`` records an
    INSTALLATION EVENT and sits in a SIBLING directory, so a bind mount over
    the package directory replaces the artifact without touching the metadata
    beside it — after which the metadata describes something that is no longer
    on ``sys.path``. That is the rig everyone here iterates in, and reading
    dist-info for it measures the wrong artifact.

    So the metadata is authoritative exactly while no mount boundary separates
    it from the tree it claims to describe (``/proc/self/mountinfo`` answers
    that directly, before any import). When it is authoritative, a mismatch is
    the STALE IMAGE and is refused. When a mount has disowned it, the revision
    is unrecorded — a checkout carries a directory, not the commit that names
    it — and the question is put to the tree itself instead: does it DECLARE
    the calls this source makes (:class:`SourceSurface`)? That is not a proxy
    for the failure, it is the failure — the incident was ``exl3_gemm`` without
    its ``size_n_out`` bound raising ``TypeError`` from pybind11.

    UNDETERMINABLE IS NOT A PASS. A dependency that is installed but whose
    revision nothing records, and whose declared surface cannot be read either,
    is exactly the silence this guard exists to break, so it is reported as
    loudly as a mismatch. Only an ABSENT dependency is waved through, and it is
    named in the log when it is.

The same floor, as the INSTALLED ``arbi-serve`` distribution records it.

The authority for a deployment that has no source tree beside it. A wheel
freezes the requirement its ``pyproject.toml`` declared at build time, so
this is the same number by a different route rather than a weaker one.

``(floor, where it was read)`` — the requirement the backend must satisfy.

:data:`TURBO_ATTN_FLOOR` first, because it is the only source that travels
WITH the imported code. A bind mount replaces the package and nothing
else, so a floor read from the image around it is a fact about the tree the
image was built from, and the code calling the backend is the mounted one —
an older image then supplies its own older floor and passes itself.
``pyproject.toml`` and the installed distribution answer for a tree that
predates the constant; :func:`tests.test_turbo_attn_floor_single_source`
holds them equal.

Refuse a ``tkv`` that imported but contains no package.

An empty DIRECTORY named ``tkv`` on the path imports cleanly as a PEP-420
namespace package: ``import tkv`` succeeds, ``__file__`` is ``None`` and
there is no ``__version__``. The version guard treated that as "cannot read
the version" and skipped, so the boot continued and died minutes later on
whichever ``from tkv... import`` ran first — a message about a missing name,
which reads as a version-skew problem and is not one.

It is reachable from one mistake, and it is a common one:
``compose.turbo-attn-bindmount.yaml`` mounts
``${TURBO_ATTN_SRC_ROOT}/tkv`` over the installed package, and Docker
CREATES a missing bind source as an empty root-owned directory rather than
refusing. A ``TURBO_ATTN_SRC_ROOT`` pointing anywhere that does not exist
therefore replaces tkv with nothing, silently.

Checked here because this guard already runs at boot and already owns the
question "is the backend under this engine the one we think it is".

Match a Python source file that BINDS ``name``.

Deliberately not a bare substring: a docstring or a comment naming the
symbol is how a source scan comes to report present for a tree that no
longer defines it. Definition, assignment, ``as`` re-export and plain
``from … import`` are the four forms turbo-attn actually uses to publish a
name, so all four count and nothing else does.

Every ``tkv`` name ``arbi_serve`` imports AT MODULE SCOPE.

Module scope is the selection, not a sample: those imports run while the
engine's own modules are being loaded, so a missing one ends the boot with
an ``ImportError`` naming a module that has nothing to do with the mount
that caused it. A name imported inside a function fails the one call that
reaches it, which is a different and much louder failure.

``tests/test_backend_version_guard.py`` holds this list against both the
INSTALLED tkv and the ``arbi_serve`` sources, so an entry that has gone
stale fails there rather than by refusing a good tree.

Where the loader WOULD import tkv from, without executing it.

Asked after an import has already failed, so it separates the two causes an
exception cannot: a tkv that is not installed at all, and one that is there
and refuses to load.

The loaded ``tkv``, or a refusal naming the cause the exception cannot.

An import that fails has two very different causes and one exception text.
tkv NOT INSTALLED is a closure problem — the engine extra was not installed,
or this is the client-only distribution. tkv INSTALLED BUT NOT IMPORTABLE is
something placed in front of it, which in this repo means the bind mount.
Refusing with the origin in hand beats letting the same exception surface
later from whichever engine module happened to import tkv first.

``git <sha> (clean|dirty)`` for a tree inside a work tree, else why not.

Best-effort and bounded: the answer is provenance for a human reading the
boot log, never a gate. The canonical bind mount grafts the package
directory alone, so the checkout's ``.git`` is NOT inside the container and
this correctly reports that there is no work tree here — which is itself
the fact an operator needs when the version stamp is the only identity the
mount carries.

``(version, where it came from, whether it describes THIS tree)``.

``tkv._resolve_version`` prefers the ``_version.py`` its own build wrote and
falls back to the ``turbo-attn`` dist-info, so the presence of that file in
the tree on ``sys.path`` is what separates the tree's own claim from a
number belonging to whatever was installed underneath it.

State which turbo-attn this process loaded, and refuse one it cannot call.

Three readings, all available before the first CUDA call and none of them
able to report success without having run:

* PROVENANCE — the file tkv imported from, the version and WHICH artifact
  supplied it, git identity when a work tree is reachable. Logged at INFO
  whether or not anything is wrong, so "verified" is distinguishable from
  "never ran".
* FLOOR — the version against the ``turbo-attn>=`` requirement, read from
  project metadata rather than hardcoded. A floor that cannot be read at
  all is refused; a version that cannot be read is refused; a version that
  belongs to a dist-info a mount has disowned is not compared, because it
  is not a fact about the tree that will run — and is said so in the log.
* SURFACE — whether the tree declares the tkv names ``arbi_serve`` imports
  at module scope. This is the instrument for the disowned case, and it is
  not a proxy: a name that is missing here IS the ImportError the boot
  would otherwise take minutes to reach.

Exact agreement with ``uv.lock`` is an image-build invariant enforced by
``scripts/ci/assert_image_backend_matches_lock.py --exact``.

Fail the boot, unless the deliberate-A/B override is set.

ONE refusal path for every guard in this module: a second copy is a second
place for the override to be honoured differently, and the override's whole
value is that an operator can predict what it does.

One backend API this source CALLS, as the backend's own sources declare it.

A revision comparison asks about an installation event. This asks the
question the incident actually turned on — does the tree that will be
compiled and imported declare the call ``arbi_serve`` makes? — and it is
answerable from the tree alone, which is what makes it the instrument for a
tree whose revision nothing records. Read, never imported: these are the
``.cuh`` declarations the JIT build compiles, so reading them costs no CUDA
and no build, and they are the same text ``nvcc`` will see.

A dependency this source tree resolves from a git commit, not a version.

``required`` is read from a constant inside ``arbi_serve`` rather than from
a lock file, because a bind mount carries the package directory and nothing
else — see this module's docstring.

Every git-pinned native dependency, as a SET rather than one instance.

A guard naming a single dependency only ever catches that dependency's
instance of the failure; the reasoning covers every package whose surface
the engine calls and whose identity a version string cannot express. Today
the fork is the only such package — the tuple is where the next one goes,
not a second copy of this function.

The commit the environment ACTUALLY carries for ``pin``.

Three outcomes, deliberately distinct — collapsing the last two into a
single falsy value is how a staleness check comes to report fresh:

``(None, reason)``
    No such distribution. There is no backend here to be stale.
``("", reason)``
    Installed, but nothing records WHICH revision. Undeterminable.
``(commit, source)``
    Determined.

dist-info is preferred over the image stamp because it is the INSTALLED
truth: the stamp describes the image as built, and a package replaced
inside a running container would leave the stamp describing a revision that
is no longer there. The image build asserts the two agree when it writes
the stamp, so the fallback only ever supplies a revision dist-info lost.

BOTH describe an installation, not a directory. Whether that installation
is still the directory on ``sys.path`` is a SEPARATE reading —
:func:`_metadata_describes_loaded_tree` — because a bind mount replaces the
one without touching the other.

Directory the loader would import ``pin`` from, WITHOUT importing it.

``find_spec`` resolves the location from the finder and does not execute
the package ``__init__``, so this stays free of the CUDA-extension build
that importing ``exllamav3`` triggers.

Whether the dist-info still describes the directory on ``sys.path``.

Takes a distribution NAME rather than a pin: the reading is the same one for
every native dependency this engine loads (turbo-attn's package directory is
bind-mounted by exactly the same mechanism as the fork's), and a second copy
is a second place for the mount question to be answered differently.

The metadata and the package directory are siblings one installer wrote
together, so they are the same artifact exactly while nothing has been
mounted between them. The mount table answers that directly, instead of
inferring it from the very disagreement the guard is trying to judge.

FAIL LOUD, NOT OPEN: when neither the mount table nor a device id can be
read the answer is ``True`` — the metadata is taken as authoritative, so a
drifting revision still refuses. An unreadable mount table must not become
a way to skip the check.

Directory the CACHED JIT build of ``extension`` compiled its sources from.

``torch.utils.cpp_extension`` writes a ``build.ninja`` naming every
translation unit by absolute path, so the built ``.so``'s provenance is
readable without loading it — which is what makes "the compiled half and
the imported half are one tree" a reading here rather than an assumption.

Say whether the compiled half was built from the tree that will be imported.

Reported, never refused: torch regenerates ``build.ninja`` and re-runs
ninja on every load, so a build dir left over from another tree is rebuilt
rather than served. The disagreement is still a real condition an operator
wants named — it is the one shape in which the Python half and the CUDA
half of a single package sit at different revisions.

Judge a tree whose metadata no longer describes it, by what it declares.

The refusal this can still produce is the SAME incident the revision
comparison catches — a backend without the surface this source calls — read
off the artifact that will actually be compiled and imported instead of off
a record of what was once installed.

Fail loudly when a git-pinned backend is not the revision this source needs.

Runs for every boot, not only an EXL3 one. A moved pin is evidence about
the IMAGE — its native closure predates this source tree — and that fact
does not become false because the model about to load happens not to reach
the fork. Gating on the model would put a proxy in front of the quantity
and hand back the silent window this guard exists to close; the override
covers the operator who has decided the divergence is irrelevant to the run
in front of them.

WHAT it compares is settled per pin by
:func:`_metadata_describes_loaded_tree`: the recorded revision while the
installation is still the directory on ``sys.path``, and that directory's
own declared surface once a mount has replaced it. Neither reading needs an
import or a CUDA call, so both are available in the first second.

The hot per-step batch materializer for :class:`EagerModelRunner`.

``build_batch`` is the per-step flat-tensor materialization path: it
allocates page slots, gathers per-token ids / positions / slot mappings,
and writes them either into the engine's persistent piecewise-buffer
ring (the cudagraph hot path) or into freshly-allocated activation-arena
tensors (CPU stub / oversized / capture-off fallback).

This hot method is a free function taking the runner by reference
(never a delegating wrapper in the loop). The runner keeps a one-line
``_build_batch`` method that forwards to here so the engine path
(``eng.model_runner._build_batch``) and the incremental-equivalence
tests are preserved. All persistent caches (``runner._bt_occupant`` /
``_bt_occupant_pb_id`` / ``_bt_epoch`` / ``_time_build_batch`` /
``_build_batch_ns_total`` / ``_build_batch_calls``) remain instance
state on the runner and are read/written through it.

This module is the public façade for the per-step batch builder. The
concrete implementations live in sibling modules and are re-exported
here so every import path (``batch_build.<name>``) keeps working:

  * :mod:`arbi_serve.runtime._batch_build_helpers` — shared helpers.
  * :mod:`arbi_serve.runtime._batch_build_persistent` — the pure-decode
    persistent fast path (:func:`try_persistent_decode`).
  * :mod:`arbi_serve.runtime._batch_build_materialize` — the full-rebuild
    slate gather + tensor materialization.
  * :mod:`arbi_serve.runtime._batch_build_plan` — StepPlan-side bridges
    and the captured-prefill replay staging.

Host-side ``cu_seqlens_q`` for FLA's ``cu_seqlens_cpu``, or ``None``.

Thin adapter over :func:`~arbi_serve.runtime._batch_build_helpers.
host_cu_seqlens_q_twin` — the ONE decision point every batch
construction (serving and boot-time synthetic alike) shares. See that
function for why the decision may not be duplicated.

Cacheable cold-boot artifacts that ride alongside the BootManifest.

The :mod:`~arbi_serve.runtime.boot_manifest` caches the single biggest
cold-boot phase — the multi-shape activation profile. This module adds the
RoPE table computation as a cacheable artifact so a warm restart can skip
it too:

  - :class:`RopeCacheArtifact` — the RoPE cos/sin tables built on first
    forward. A warm boot loads the cached tensors so the model's lazy
    ``_ensure_rope`` / ``RoPECacheRegistry.get`` paths copy them in
    instead of recomputing the trig tables.

Each artifact follows the same fail-safe contract as the BootManifest:

  * **signature** — a content hash of the artifact's inputs. Two boots
    at the same configuration produce the same signature; any relevant
    change bumps it.
  * **persist** — atomic write (``mkstemp`` + ``fsync`` + ``os.replace``;
    ``torch.save`` to a tmp file then ``os.replace`` for the pickled
    graph / tensors). A crash mid-write leaves a ``.tmp`` behind, never a
    half-written final file.
  * **load** — self-healing read. A miss (file absent, unreadable,
    version / signature mismatch, malformed payload) returns ``None`` so
    the caller recomputes live. An artifact can only ever SKIP work that
    reproduces the same value; it can never serve a wrong
    layout / graph / RoPE table.

Cache layout
------------
This artifact shares the BootManifest fingerprint key — it lives in
``<boot-manifest-dir>/<fingerprint-key>/`` so it invalidates together
with the build fingerprint (model / dtype / shapes / GPU arch / kernel
versions) AND carries its own internal signature stamp for defense in
depth. The enable flag (``serve_boot_manifest``) and the cache directory
(``serve_boot_manifest_dir``) are reused from the BootManifest — no new
env knob.

Cached RoPE cos/sin tables built on first forward.

The cos/sin tables are deterministic functions of ``(head_dim /
rotary_dim, rope_theta, partial_rotary_factor, max_seq_len, dtype)``.
Caching them as raw tensors saves the first-forward trig-table
construction.

The warm path deposits the cached pairs on
``eng._cached_rope_tables`` keyed by ``(head_dim, theta, partial)``;
the model's lazy ``_ensure_rope`` (Qwen3.5) and
``RoPECacheRegistry.get`` (mixin arches) copy them into the
freshly-built buffers, skipping the trig compute.

``signature`` = ``(head_dims_used, max_seq_len, dtype)``, where
``max_seq_len`` is the row count the tables are actually BUILT to —
``dims.effective_rope_cache_seq_len()``, the same number the per-arch
RoPE ctors and the ``scratch.rope`` budget predictor read — not the raw
``max_position_embeddings``. The two differ whenever the serving gate
narrows the table (a pinned ``--max-context`` below the model ceiling),
and a signature that ignored the difference would match an artifact of
the wrong SHAPE: the copy into the fresh buffers then raises and is
swallowed, so the miss is silent and the warm path degrades with no
signal. A model code change that adds a new RoPE shape, or any change to
the built row count, bumps the signature.

Cold-boot step skipped on a hit: the first-forward trig-table
construction of the RoPE cos/sin caches.

Collect every distinct ``(head_dim, rope_theta, partial)`` triple.

Walks ``eng.model.layer_specs``; ``partial`` is ``rotary_dim /
head_dim`` when the model exposes a partial ``_rotary_dim`` (Qwen3.5 /
3.6), else ``1.0``. Non-attention layers (head_dim ≤ 0) are skipped.

Warm path: hydrate the RoPE-cache artifact from cache onto the engine.

Returns a ``{name: hit}`` map. The artifact is fail-safe — a miss /
corrupt entry returns ``False`` and the cold computation runs live.
Never raises.

Cold path: persist the RoPE-cache artifact from the live engine.

Returns a ``{name: persisted}`` map. Best-effort — a persist failure
(RoPE not yet built) is logged and swallowed; the next boot recomputes
it live. Never raises.

Snapshot every built RoPE table from the model and pickle them.

Walks the model's ``_rope_registry._caches`` (mixin arches) or
the Qwen3.5-style single ``_rope_cache`` attribute. A registry
that hasn't been built (no first forward yet) produces nothing —
the next boot misses + re-caches.

Warm path: load the tables and stash them on the engine; ``None`` on miss.

Deposits the pairs on ``eng._cached_rope_tables`` keyed by
``(head_dim, theta, partial)``. The model's lazy ``_ensure_rope``
/ ``RoPECacheRegistry.get`` copy them into the freshly-built
buffers.

The pairs stay on the HOST. Their only consumer is a ``copy_`` into a
device buffer already allocated inside ``scratch.rope``, and ``copy_``
takes a CPU source — so a device-resident staging copy buys nothing
while costing a full SECOND set of tables in the torch DEFAULT pool,
outside every named pool, for the process lifetime. Nothing ever
released them, so they were bytes the VRAM ledger could not place and
the KV budget could not reclaim.
:func:`arbi_serve.engine.build_phases_kv._build_activation_profile`
drops the host copy once every RoPE table has been built.

Boot-time cache hit/miss registry for the admin startup panel.

Records, per boot phase, whether a phase reused its on-disk cache (a **hit**,
cheap) or (re)built it (a **miss**, the expensive phases). The compile caches
come from :func:`arbi_serve._compile_cache_env.compile_cache_report`; this
registry captures the ones that live outside that report — the flat-weights
(flat tensor) dump and the activation profile — so the startup panel can colour
each phase line by its own cache's status.

Boot-only bookkeeping (a dict write per phase); never touched on the hot path.

Stamp one boot cache's hit/miss. ``phase`` links it to a startup-phase
slug so the UI can colour that phase line.

The slug is checked against the metrics layer's canonical set: a phase the
UI cannot resolve colours as "not cache-backed", so a typo hides a rebuild
rather than reporting one.

Which served process answered — one value, minted once, on every response.

A benchmark cell asks a URL for tokens and reports what came back. Nothing in
that exchange says the answers came from ONE server. They routinely do not:

* a container name is global to the docker daemon rather than scoped to a
  compose project, so a neighbour running the same recipe from another
  directory owns the same name and ``up -d --force-recreate`` replaces the
  running server;
* an out-of-band SIGTERM (an OOM guard configured ``--prefer python`` takes a
  27B engine first) makes the engine drain cleanly and exit 0, after which a
  ``restart:`` policy brings a NEW process back on the same port.

Both are silent on both sides. The replaced run keeps the requests it had
already completed and simply stops getting new ones, so the cell it reports is
SHORT and FAST -- the requests that survived are the ones that finished, and
the ones that dropped out are the slow ones. A destroyed run produces the
best-looking number in the table.

The identity is per PROCESS, not per container or per host: it is the thing
that is actually different after any of those events, it needs nothing from
the docker daemon (so a client on another machine can check it), and it moves
on a restart that leaves the container's name, id and image unchanged. A
caller that captures it on its first response and compares it on every
subsequent one cannot mistake a replaced server for the one it measured.

BootManifest — persist the activation profile across cold boots.

The multi-shape activation profile (:func:`arbi_serve.runtime.
activation_profile.profile_engine`) runs three synthetic forwards —
``max_batched_tokens`` prefill, ``max_batch`` decode, and a mixed
chunked-prefill batch — and reports the MAX of their allocated-byte
peaks. It is the single biggest cold-boot phase: three real
``model.forward`` passes at the worst-case shapes, eager (no compile),
against a freshly built engine.

The profile output is a pure function of the things that determine the
forward's allocation footprint: the model arch, dtype, the three shape
sizes (``max_batched_tokens`` / ``max_batch`` / ``chunk_prefill``), the
KV layout/backend, the GPU arch, and the kernel-library versions (torch /
turbo-attn) that decide the per-kernel scratch. Two boots at the SAME
fingerprint produce the SAME peaks — so we cache the
:class:`~arbi_serve.runtime.activation_profile.MultiShapeProfile` keyed on
a SHA of those inputs and skip the live profiling on a hit.

This mirrors the existing graph-pool budget cache
(:func:`arbi_serve.engine.memory_budget.predict_graph_pool_bytes` /
:func:`~arbi_serve.engine.memory_budget.measure_and_persist_graph_pool`):
same atomic-write discipline (``mkstemp`` + ``fsync`` + ``os.replace``),
same self-heal on a corrupt / version-mismatched / garbage JSON (fall
through to live profiling), and the same cache-dir resolution convention.

Fail-safe, not fail-open
------------------------
A manifest MISS — file absent, unreadable, version mismatch, fingerprint
mismatch, or a malformed payload — ALWAYS falls back to live profiling.
The manifest can only ever SKIP work that would reproduce the same
number; it can never cause the engine to serve a wrong KV size. Sizing
the KV pool against a stale/wrong activation peak would OOM at capture or
serve corrupt KV, so every uncertain path returns ``None`` (miss) and the
caller profiles live.

Installed version of ``dist_name`` or ``""`` when unresolvable.

Kernel-library versions belong in the fingerprint: a torch /
turbo-attn bump can change per-kernel scratch allocation, so a
manifest captured under an older library must NOT be reused. Best-
effort — a missing dist degrades to ``""`` (still a stable, hashable
component) rather than raising.

Pick the boot-manifest directory.

Resolution order:
  1. Explicit ``build_dir`` argument (tests).
  2. ``ARBI_SERVE_BOOT_MANIFEST_DIR`` (``serve_boot_manifest_dir``).
  3. ``/cache/arbi-serve/boot-manifest`` when writable — the canonical
     convention on EVERY host: a compose volume in containers, a
     ``/cache -> /mnt/k8scache/cache`` symlink on bare-metal boxes
     (shared NFS, so manifests warm across the fleet).
  4. ``~/.cache/arbi-serve/boot-manifest`` — defense-in-depth when the
     ``/cache`` mount is broken; logged as a WARNING because a missing
     mount means the fleet-shared cache is silently not being used.

Assemble the fingerprint that determines the activation profile.

Every value is JSON-serializable so the SHA reproduces across boots
at the same configuration. The inputs are exactly the things the
three synthetic forwards' allocation footprint depends on:

  * ``model_id`` — model arch + weights identity (path / name).
  * ``dtype`` — activation/compute dtype (``str(eng.dtype)``).
  * ``tp_size`` — tensor-parallel degree (per-rank shapes shrink).
  * ``kv_backend`` — paged-KV backend class name (the TKV codec and
    ``tkv-bypass`` allocate different scratch; the KV bits live inside
    the codec).
  * ``max_batched_tokens`` / ``max_batch`` / ``chunk_prefill`` — the
    three synthetic shape sizes.
  * ``prefill_probe_tokens`` — ``min(max_batched_tokens, max_context)``,
    the clamp that fixes every synthetic shape's token / sequence /
    KV-length count. This is the only route by which the served context
    reaches the measured peak, and it is folded INSTEAD of the raw
    ``max_context``: under ``--max-context auto`` that number is narrowed
    to whatever KV the boot realizes, so folding it forks the key on
    every boot and the persisted profile is never read back.
  * ``max_context`` — the operator's PIN, or ``"auto"`` when the engine
    may narrow it (``max_context_is_auto`` /
    ``max_context_is_recipe_default``). A pin is config and reproduces
    across boots; a narrowable context reaches the profile only through
    ``prefill_probe_tokens``.
  * ``block_size`` — page size: KV layout + the probe's page count.
  * ``mtp_enabled`` / ``mtp_n_draft`` — verify/drafter shapes.
  * ``drafter_source`` — WHICH drafter, from config
    (:meth:`~arbi_serve.config_groups.MtpConfig.source_identity`). The
    depth above says how many tokens are proposed; this says what proposes
    them, and the two are independent. The drafter is attached BEFORE
    ``_build_activation_profile``, so attaching one changes what the
    profiled forwards allocate — a DFlash drafter arms the target model's
    tap slab and its per-layer writes; an EXL3 drafter re-runs
    ``reserve_exl3_reconstruct_scratch`` / ``pin_exl3_kernel_shapes``
    against its own linears, which can widen the shared reconstruct
    scratch and move the pinned GEMM leg the profiled prefill dispatches.
    A CONFIG read, deliberately, not a read of the loaded drafter: the
    fingerprint is built at two points of a boot (boot artifacts, then the
    profile) and only a config read gives the same key at both.
  * ``recurrent_capacity`` / ``sliding_window`` — recurrent slab +
    SWA path selection in the synthetic forward.
  * ``safety_fraction`` — folds into the stored budget margin.
  * ``gpu_arch`` — compute capability (kernel selection).
  * ``torch_version`` / ``turbo_attn_version`` — the kernel-library
    identity (per-kernel scratch can move on a bump; torch_version
    also carries the CUDA build tag).
  * ``shape_labels`` — the SHAPE SET the profile will contain
    (:func:`~arbi_serve.runtime.activation_profile.expected_shape_labels`).
    The three base shapes are unconditional but the stochastic verify tail
    is not, and its conditions include ``mtp.rejection_sampling_enabled``,
    which no other component here carries. Without this the cache serves a
    profile that measured a DIFFERENT SET of shapes: a tail-less entry read
    back as ``stochastic_measured == 0``, and the serving floor fell to the
    of KV silently on every boot at that key.
  * ``profile_cache_version`` — explicit bump knob for changes to the
    measurement logic itself.

What this fingerprint deliberately does NOT carry is the arbi-serve git
SHA: the profile depends on shapes and kernels, not on every edit, and a
key that moved per commit would re-profile the ~19 s activation phase on
every boot. The compensating control for a change in our own MEASUREMENT
code is the manual ``_BOOT_MANIFEST_VERSION`` bump.

Reconstruct a :class:`MultiShapeProfile`; ``None`` on any malformation.

Strict: a payload that does not type-check in EVERY field returns
``None`` (a miss → live profiling), never a partial/garbage profile
that could under-size the KV pool.

Load a cached activation profile for ``fingerprint``, or ``None``.

Returns ``None`` (a MISS → caller profiles live) for ANY of: file
absent, unreadable, JSON garbage, version mismatch, fingerprint
mismatch, or a malformed profile payload. Self-heals — a corrupt
file never raises, it just falls through to live profiling.

Atomically write the activation profile for ``fingerprint``.

Atomic write: ``mkstemp`` in the same dir, ``fsync``, then
``os.replace`` — a crash mid-write leaves a ``.tmp`` behind but
never a half-written ``<key>.json``. A write failure is logged and
swallowed (next boot re-profiles); persistence is best-effort and
never blocks boot.

Returns ``True`` on a successful persist, ``False`` otherwise.

Which GPU architectures this build's CUDA kernels were compiled for.

A CUDA extension carries SASS for the architectures its build named and for
no others. Ask a card outside that set to launch one and the driver answers
``no kernel image is available for execution on the device`` — from inside
whichever kernel happened to be first, after the model has loaded and minutes
into a boot. The arch is a property of the IMAGE, known before a single weight
is read, so the refusal belongs where the fact is: at the point a kernel
surface is resolved, naming the card, its ``sm_XY``, and the list the image
was built for.

**Where the list comes from — the BINARIES, not the build arg.** The image
records at ``/opt/venv/share/arbi-serve.gpu-arch`` the archs its baked
extensions actually carry: ``scripts/ci/assert_baked_arch_coverage.py
--record`` runs ``cuobjdump --list-elf`` over the arch-fanned extensions in
the builder stage and writes the archs every one of them holds a cubin for.
The record is trustworthy at runtime because it is DERIVED from the artifact
— it restates no intent, so there is no step whose absence could make it
false. That is the whole of the justification, and it does not depend on any
other build step having run (#2196).

``GPU_ARCH`` keeps its own, separate job: it is the build's REQUEST, and the
same script's assertion mode is what fails the image when the request was not
honoured. This module reads neither. Two extensions are deliberately
single-arch (``arbi_serve_nvfp4_marlin``, ``arbi_serve_exl3_i8_gemm_v*``) and
are excluded from the derivation for the reason that script's docstring gives,
so the record describes the fanned surface — the one a card must be able to
launch — and not those two.

**No record is not a refusal.** A source checkout, a dev box, an editable
install — none of them write that file, and none of them are wrong. Absence
means "unknown", and an unknown build cannot be shown to exclude this card, so
the gate stays out of the way and the old failure mode is what remains. Only a
POSITIVE record that excludes the device refuses.

**The compatibility rule is CUDA's, not an equality.** A cubin built for
``sm_X.y`` executes on a device of compute capability ``X.z`` for ``z >= y``
— forward within a major family, never backward and never across families. So
an ``8.6`` cubin runs on a 4090 (8.9) and an ``8.9`` cubin does NOT run on an
A6000 (8.6). PTX would be the escape hatch, since the driver can JIT it for a
newer device; the shipped extensions carry none (verified with
``cuobjdump --list-ptx`` on every baked ``.so``), which is why this file's rule
has no PTX term.

Parse a ``GPU_ARCH`` record into sorted ``(major, minor)`` pairs.

Accepts the raw file contents (``GPU_ARCH=8.9;12.0``) or a bare list
(``8.9;12.0``). Entries that are not ``<int>.<int>`` are dropped rather
than raised on: a ``+PTX`` suffix or a stray blank is a shape this record
may legitimately grow, and a parse error here would refuse a boot over a
formatting detail — the opposite of what the gate is for.

The archs this build compiled for, or ``None`` when unrecorded.

``None`` and ``()`` are different answers and callers must not conflate
them: ``None`` is "this build left no record" (a source checkout), ``()``
is "the record parsed to nothing", which is a broken record and equally
not grounds to refuse.

Raise if this build has no kernels for ``device``. No-op otherwise.

Silent on every host where the question cannot be answered — no CUDA, no
arch record — because an unanswerable question is not a refusal (see the
module docstring). ``component`` is the surface being resolved and leads
the message, so an operator reads which loader stopped and why before any
weight is touched.

Per-kind capture/replay impls.

Each kind has its own typed pool — different graph classes with
different replay signatures, so polymorphic unification never paid.
``engine.cudagraph_pools`` constructs each pool directly:

  * :mod:`.decode` — whole-forward decode capture (``capture_decode``).
    The ``CapturedGraph`` / ``CapturedGraphPool`` shape-keyed pool lives
    in :mod:`.decode_graph`; the pre-flight gates (``_can_capture_decode``
    / ``_can_capture_prefill`` / ``_can_capture_mtp_verify``) in
    :mod:`.preflight`; whole-forward prefill capture (``capture_prefill``)
    in :mod:`.prefill`; the failed-capture stream drain
    (``_force_end_stream_capture``) in :mod:`._common`.
  * :mod:`.drafter` — K-step drafter chain capture (greedy argmax;
    true-stochastic slates run the live chain), ``DrafterChainGraph``
    and its pool.
  * :mod:`.dispatch` — per-layer piecewise capture used for variable-
    length prefill and hybrid models — ``LayerCapturedGraph`` /
    ``LayerCapturedGraphPool`` / ``PiecewiseBuffers`` /
    ``model_dispatch`` / ``dispatch_layer``.

Shared low-level helpers for the capture/replay subsystem.

Holds :func:`_force_end_stream_capture` — the post-failure CUDA
stream-capture drain used by every capture path
(:mod:`.decode`, :mod:`.dispatch`, :mod:`.drafter`). It lives here
rather than in any one capture module so the lazy importers in the
sibling modules don't introduce an import cycle.

Torch-reserved bytes, for capture-cost deltas. Boot-only, never hot.

``torch.cuda.memory_reserved`` IS the whole truth for a capture delta —
including on the cuMem path. Every pool a capture allocates into
(``capture.cudagraphs``, ``capture.io_buffers``, ``scratch.*``) is a
:class:`~arbi_serve.runtime.named_pool.NamedMemPool`, i.e. a **torch**
``MemPool`` whose segments happen to be *backed* by the cuMem pluggable
allocator. A pluggable allocator changes WHO calls the driver, not whether
the caching allocator books the segment: ``memory_reserved`` counts it
either way.

THIS FUNCTION USED TO ADD ``CuMemPoolAllocator.mapped_bytes(None)`` on top,
on the docstring's claim that ``memory_reserved`` "CANNOT see a cudagraph
pool backed by the cuMem allocator". That claim is FALSE, and the wrong
allocated into a cuMem-backed ``NamedMemPool``, deltas bracketed around the
allocation)::


The two terms are the same physical bytes seen through two counters, so the
sum double-counted every captured graph. ``CapturedGraph.
capture_memory_bytes`` is a delta of this function, ``*Pool.
total_memory_bytes()`` sums those, and ``Engine.total_cudagraph_bytes()``
the real one. GPU-measured on a Qwen3.5-0.8B TP1 bf16 mb64 boot: reported
Those phantom bytes were held out of KV.

``mapped_bytes`` is NOT redundant everywhere — it is the right counter for
the DIRECT-cuMem consumers that bypass torch entirely (``GrowableRegion``,
which backs the KV slabs; see ``grow_kv_after_capture``'s mem-triage, which
subtracts torch-reserved and the growable KV separately for exactly this
reason). It is redundant HERE because captures allocate through torch
MemPools, never through raw ``map_range``.

Return the persistent-buffer allocation context for a capture.

Persistent kernel-input/output buffers route through ``buffers_pool`` —
a MemPool distinct from the cudagraph capture pool ``graph_pool`` (see
the BUG A address-reuse rationale in
:func:`arbi_serve.runtime.capture.decode.capture_decode`). Falls back to
``graph_pool`` when ``buffers_pool`` is None (legacy callers / unit
tests), or a null context when neither pool is supplied (CPU / no-GPU
stub). Shared by the decode / prefill / drafter capture paths so the
pool-defaulting cannot drift between them.

Return the ``compute_logits`` value for a (possibly verify) forward.

On the MTP verify shape (``S = K+1 > 1``) the pass recomputes
``lm_head`` over ALL K+1 flat tokens itself, so the model's baked
last-token ``lm_head`` GEMV (+ its TP all-gather) is pure per-step
waste — both verify consumers (``mtp_verify_offload`` /
``mtp_verify_spmd``) assign the first return slot to ``_logits_last``
and discard it. The verify shape therefore ALWAYS returns
``compute_logits=False`` so the CAPTURED verify graph never bakes that
GEMV + all_gather (which would be replayed and thrown away every spec
step). This is unconditional — the skip is proven byte-identical (the
only consumed output, ``hidden``, is computed before the head).

Plain decode (``is_verify_shape=False``) ALWAYS keeps
``compute_logits=True`` — its last-token logits ARE the sampler
input, never skipped. Shared by the capture path
(:func:`arbi_serve.runtime.capture.decode.capture_decode`) and the
SPMD live verify forward (``mtp_verify_spmd``) so the skip decision
cannot drift between capture and replay.

Default the ``torch.cuda.graph(pool=...)`` id to the named pool's id.

When ``mem_pool`` is unset and a ``graph_pool`` is supplied, route the
captured working set into the named pool's ``MemPool.id``. Routing goes
through the ``torch.cuda.graph(pool=...)`` arg ONLY (NOT
``NamedMemPool.use()`` simultaneously — torch raises on
double-registration of the same pool). Returns ``mem_pool`` unchanged
otherwise. Shared by every capture path.

Tracks pre-capture stable-VA buffer migrations + writeback binding.

Under the stable-VA sleep backend every persistent buffer the captured
graph references must already live on its FINAL (stable-VA) ``data_ptr``
BEFORE the capture region bakes that pointer into kernel args (see
:func:`arbi_serve.runtime.capture.decode.capture_decode` section 1b).
This helper centralises the gate, the per-buffer migration, and the
post-capture host binding so the three sites cannot drift.

Usage::

    migrator = _PreCaptureMigrator(sleep_pool)
    if migrator.active:
        buf = migrator.migrate(buf, "buf")
    ...  # build the captured-graph container
    migrator.bind_hosts(captured)

Record on ``exc`` that it unwound out of an open cudagraph record.

:func:`_force_end_stream_capture` is the only caller — it is the one place
that asks CUDA whether the stream was still capturing, so it is the one
place that knows. Best-effort: an exception whose class forbids attribute
assignment simply carries no mark.

True iff ``exc`` (or something it was raised from) aborted an open record.

Follows ``__cause__`` / ``__context__`` for the same reason
:func:`~arbi_serve.engine.boot_degradation.caused_by_oom` does: a capture
helper that re-raises a typed refusal ``from`` the original must not lose
the fact that a stream capture was torn down underneath it.

Force-drain CUDA stream-capture state on a stream after a failed
capture, so PyTorch's allocator ``captures_underway`` set is cleared
and the process can shut down cleanly.

Returns True when the stream was STILL INSIDE A CAPTURE at the moment
this ran — i.e. the failure unwound out of an open record region rather
than from before it opened. ``failed_with`` (defaulting to the exception
currently in flight, since every caller is an ``except`` block) is
stamped with that fact via :func:`mark_capture_aborted`, so a sweep
handler far from here can tell a bucket that could not be captured from
a bucket whose capture was torn down half-recorded. The default is what
makes it un-forgettable: a new capture site inherits the classification
by calling this the way every other site already does.

PyTorch's :class:`torch.cuda.graph` adds the capture pool to the
allocator's ``captures_underway`` set at ``capture_begin`` and
removes it at ``capture_end``. When ``capture_end`` itself throws
(because a kernel inside the region invalidated the capture), the
set entry is never removed and the next allocator-deinit hits an
``INTERNAL ASSERT FAILED: captures_underway.empty()``.

Recovery (in order):

  1. ``cudaStreamEndCapture`` on the capture stream regardless of
     current state. Returns success on a successfully-ended capture
     and ``cudaErrorStreamCaptureInvalidated`` on a stream that's
     already in the failed-capture state — both leave the stream
     capture-flag-clear afterwards. We swallow the return value
     since there's no recovery to attempt at that point.
  2. ``torch._C._cuda_endAllocateToPool`` (always) + optionally
     ``_cuda_releasePool`` (gated by ``release_pool``) on the
     supplied ``graph_pool``'s :class:`torch.cuda.MemPool` id.
     ``torch.cuda.graph(graph, pool=mem_pool)`` calls
     ``_beginAllocateToPool`` on entry and the matching
     ``_endAllocateToPool`` / ``_releasePool`` on exit; when
     ``__exit__`` raises before those run, the allocator's pool-
     routing entry stays live and the pool's
     :class:`at::cuda::MemPool` destructor at process shutdown
     hits the ``release_block`` /
     ``CUDACachingAllocator.cpp:3524`` assert and aborts. These
     private APIs are the public-equivalent of what
     :func:`torch.cuda.use_mem_pool`'s ``finally`` block runs and
     are ABI-stable on PT 2.4+.

     ``release_pool=False`` skips the ``_cuda_releasePool`` call
     while still running ``_endAllocateToPool`` — this is the path
     the boot-time piecewise capture sweep uses when one bucket
     fails: the cleanup must clear the allocator's recording state
     (so the next bucket's ``_beginAllocateToPool`` doesn't trip
     "already recording"), but MUST NOT decrement the named pool's
     ``use_count`` since subsequent buckets will reuse the same
     pool. Decrementing here would take ``use_count`` from 1 to 0
     (the ``MemPool.__init__`` increment), destroy the named pool's
     allocator entry, and the next bucket would be operating on a
     stale pool ID — surfacing again as ``beginAllocateToPool:
     already recording`` (with no recovery this time).
  3. ``CUDAGraph.reset()`` on the partially-captured graph object,
     when supplied. The C++ ``CUDAGraph::reset`` zeroes any
     partial capture state the graph holds; without it, the
     graph's ``__del__`` may try to ``releasePool`` a second time
     on a pool the allocator already considers released and
     abort with ``c10::Error: uc >= 0 INTERNAL ASSERT FAILED at
     CUDACachingAllocator.cpp:2698``. Defense-in-depth against
     a secondary crash path: even if the recovery in
     step (2) succeeded, a stale CUDAGraph that escapes the
     try/except can re-enter releasePool at __del__. Best-
     effort: any failure here logs + continues.

Implemented via a direct ``ctypes`` call to ``cudaStreamEndCapture``
because PyTorch does not expose a public Python hook to clear the
allocator's ``captures_underway`` set without going through
``CUDAGraph.capture_end``, which has already failed by the time we
enter recovery. The ``_cuda_*Pool`` calls in step (2) are private
PyTorch internals (``torch._C.``-prefixed) but documented through
:func:`torch.cuda.use_mem_pool`'s implementation in
``torch/cuda/memory.py`` and stable across PT 2.4 → current.

Args:
    stream: the CUDA stream that was the capture target.
    graph_pool: optional named MemPool wrapper supplying
        ``mempool().id``; passed when the caller routed the
        capture-time allocations through a named pool. ``None``
        falls back to step-1-only recovery (legacy behaviour)
        unless ``mem_pool_id`` is supplied.
    mem_pool_id: optional ``MemPool.id`` tuple. Used by per-layer
        recovery in :mod:`arbi_serve.runtime.capture.dispatch`
        where the caller has only the raw pool id stamped onto
        :attr:`_PiecewiseBucketState.mem_pool` and not the
        wrapper. Mutually exclusive with ``graph_pool`` (when
        both are supplied, ``graph_pool`` wins).
    device: optional device index / torch.device the pool lives
        on; required only when ``graph_pool`` / ``mem_pool_id``
        is passed. Defaults to the current device when omitted.
    cuda_graph: optional partially-captured :class:`torch.cuda.CUDAGraph`
        object. When supplied, ``reset()`` is called on it after
        the stream + pool drain so its destructor sees a clean
        state and does NOT re-enter releasePool. Best-effort.
    release_pool: when True (default — legacy / shutdown-recovery
        behaviour), call ``_cuda_releasePool`` after
        ``_cuda_endAllocateToPool`` so the allocator's pool entry
        is destroyed and process shutdown does not trip the
        ``release_block`` assert. When False, skip ``_releasePool``
        so the named pool stays alive for subsequent reuse — this
        is the boot-time piecewise capture sweep recovery path
        where one bucket failing must NOT tear down the pool the
        remaining buckets are about to reuse.
    failed_with: the exception being unwound. Defaults to the one
        currently in flight (``sys.exc_info()``); pass it explicitly
        only when calling from outside an ``except`` block.

Returns:
    True when the stream was inside a capture (active or already
    invalidated) when this ran. False when it was not — including the
    degenerate build where ``libcudart`` cannot be loaded, which logs at
    WARNING because the classification is then UNAVAILABLE rather than
    negative.

Migrate one persistent buffer to its stable VA, tracking the entry.

Only valid when :attr:`active` (the call sites gate on it). Migrates
the tensor, records ``(name, entry)``, and returns the migrated
tensor.

CUDAGraph capture + replay for the canonical decode shape.

Replaces the per-step kernel-launch storm on the canonical decode
shape (batch_size=1, seq_len=1) with one captured graph replay. tkv's
:class:`tkv.runtime.cg_state.TQCGState` codifies the data-pointer-
stability contract that makes it safe.

What gets captured
------------------
For decode at a fixed ``(batch_size, seq_len)``:

  - the per-layer Q/K/V projections + RMSNorms + RoPE + attention +
    o_proj all-reduce
  - the page-metadata kernel that populates
    :class:`tkv.runtime.page_metadata.TQBufferPool` slots
    (``indptr`` / ``indices`` / ``lpl``) from the persistent
    ``seq_lens`` / ``block_table`` buffers
  - the final RMSNorm + lm_head and the last-token gather → fp32 logits

What stays out of the graph
---------------------------
  - request → token-list scheduling and slot allocation
  - the host-side ``torch.tensor([...], device=cuda)`` materialisations
    (copy_() into persistent buffers instead)
  - sampling (different shape per step depending on which seqs are
    "ready" — and the sampler is fast enough to not need capture)
  - per-batch ``compute_bypass_safe`` (canonical decode → always False;
    we hard-code that branch off)

Frozen host branches (read this before adding one)
--------------------------------------------------
A CUDAGraph records the branch taken on the CAPTURE batch and replays it
forever. So any predicate a backend evaluates on the HOST — a Python
``bool`` — is FROZEN at capture: if the live batch would have flipped it,
the graph still runs the recorded kernel, silently — e.g. a captured
PREFILL graph that records TKV's first-chunk bypass branch and replays
it on continuation chunks attends over the wrong KV window. The class of
bug is general.

TKV has two such host predicates, both pure math over the batch metadata
mirrors: ``compute_bypass_safe`` (batch-global; True iff EVERY row has
``q_len == seq_len``) and ``compute_row_split`` (``TKV_PER_ROW_BYPASS``,
per-row class). Decode / verify capture freezes ``bypass_safe = False``
and one homogeneous row class. That is SAFE, and provably so rather than
by luck, for two independent reasons:

  1. Polarity. The frozen value is the GENERAL branch (paged codec
     decode), not the optimistic one. The bypass branch is the special
     case (it reuses ``cu_seqlens_q`` as ``cu_seqlens_k``, valid only when
     each row's whole K range is its Q range). Freezing the general branch
     can at worst leave a fast path unused; freezing the optimistic one is
     what corrupts.
  2. Reachability. ``bypass_safe`` can only flip True if EVERY row has
     ``q_len == seq_len`` — zero cached KV. The lookups only ever hand a
     decode/verify graph a NON-prefill batch (``lookup_captured_graph``
     requires ``not is_prefill`` and ``max_query_len == 1``; the verify
     lookup requires the uniform ``K+1`` shape), and every such row owns a
     prompt already in the cache, so ``seq_len > q_len`` always. Decode-pad
     scratch rows (``seq_lens = 1``) classify ``ROW_DECODE`` — ``q == 1``
     is checked first — so they do not perturb the row class either.

``tests/test_captured_host_gate_audit_live.py`` is the empirical half: it
recomputes both predicates from the buffers ACTUALLY fed to every captured
replay of a real greedy generation (real weights, TKV codec) and fails if
either ever disagrees with what capture baked. If you add a host-computed
branch to a captured path, extend that audit — or make the decision
device-side and data-driven, which is the durable fix.

Capture covers PAGED_KV decode (single-token + MTP-verify shapes,
including sliding-window-attention layers — autotune picks the
deterministic per-shape kernel at capture time) under any TP world
size: the capture region runs through
:mod:`arbi_serve.distributed.graph_capture` so row-parallel layers'
``all_reduce`` collectives are routed onto the captured side stream.
Prefill capture, LoRA capture, recurrent-backend capture
(MLA / Mamba / GDN / short-conv) are open.

Module layout. The captured-graph container + shape-keyed pool live in
:mod:`.decode_graph` (:class:`CapturedGraph` / :class:`CapturedGraphPool`);
the persistent kernel-input/output buffer allocation (section 1) in
:mod:`.decode_buffers` (:func:`allocate_capture_buffers`); the pre-flight
gates in :mod:`.preflight`; the failed-capture stream drain in
:mod:`._common`; whole-forward prefill capture in :mod:`.prefill`. This
module keeps :func:`capture_decode`.

Capture one decode graph at the given ``(batch_size, seq_len)``.

Supports ``seq_len`` ∈ ``[1, tkv_max_verify_block_m()]`` (the SERVED
verify ceiling). ``seq_len=1`` is plain decode (block_m=1 in tkv
terminology); ``seq_len>1`` corresponds to MTP verify shapes
(block_m=K+1 where K is the per-row draft count). ``TKVCore.forward``
routes block_m within the split-K register ceiling
(``tkv_max_block_m()``) to the unified split-K decode kernel and
block_m ABOVE it (high-K) to the Turbo prefill verify route — both are
stream-capture-safe, so a single ``capture_decode`` covers either
(the synthetic batch drives whichever route ``forward`` selects). The
synthetic batch built here drives warmup + capture against that route.

``lora_bucket`` is the quantized active-LoRA count this graph is
captured for (one of :data:`arbi_serve.adapters.lora.ACTIVE_LORA_BUCKETS`).
``0`` is the canonical no-LoRA capture (every parallel linear
short-circuits ``_maybe_apply_lora``); larger buckets attach a
synthetic :class:`LoraBatchState` view that references
``lora_capture_pool``'s persistent buffers — replay copy_()s real
adapter weights into the same slots before each replay.

``verify_buffers`` is the engine's :class:`VerifyBuffers`.
When supplied AND ``seq_len > 1``, the captured graph references
prefix slices of those persistent buffers instead of fresh
allocations — replay then skips the inner ``copy_()`` calls
because the live verify path's :meth:`VerifyBuffers.load_step`
already populates the same storage before invoking the forward.
For ``seq_len == 1`` the parameter is ignored: plain-decode batches
don't go through ``VerifyBuffers``. ``None`` (default) preserves
the legacy fresh-allocation behaviour for callers that don't have
a verify buffer pool (CPU smokes, profile probe).

``sleep_pool`` is the engine's :class:`SleepableTensorPool`. When
supplied AND the stable-VA backend is active, every persistent
kernel-input / output buffer is migrated to a stable cuMem VA
BEFORE the capture region (section 1b) so the captured graph bakes
the stable ``data_ptr`` directly. ``None`` (default) or the
CPU/no-GPU stub leaves the buffers on the caching allocator — the
behaviour for CPU smokes and the profile probe.

Lifecycle:

  1. Allocate every persistent buffer the captured graph needs.
  2. Build a synthetic :class:`ScheduledBatch` referencing those
     buffers (zero-filled but plausibly valid: page id 0, position
     0, etc.).
  3. Run the metadata builder ONCE at capture-prep so the per-kind
     meta dataclass references point at persistent pool buffers.
  4. Run a warmup forward (un-captured) to JIT every kernel + warm
     every per-layer first-call path; CUDAGraph capture rejects
     any cudaMalloc / cudaFree, so first-call init must be done
     before ``with torch.cuda.graph(g)``.
  5. Capture: call the page-metadata kernel + ``model.forward``
     under a stream-scoped graph context.

The synthetic batch uses page id 0 throughout. Token ids are 0;
positions are ``[seq_lens-S, seq_lens-S+1, ..., seq_lens-1]`` per
row (the canonical "this row already has ``seq_lens-S`` prefix
tokens, now feeding S new ones" layout that matches every real
verify-pass step). Real KV content is irrelevant — capture only
cares about kernel-launch shape; replay overwrites every persistent
input the kernels read.

Persistent kernel-input/output buffer allocation for decode capture.

Section 1 of :func:`arbi_serve.runtime.capture.decode.capture_decode`:
allocate every persistent buffer the captured decode/verify graph
references (input ids, positions, slot mapping, seq lens, cu_seqlens,
block table, per-kind recurrent ``state_indices`` mirrors, and the
persistent logits/hidden outputs).

The captured graph bakes each buffer's ``data_ptr`` into its kernel args,
so these MUST be persistent (route through the ``buffers_pool`` MemPool,
distinct from the cudagraph capture pool) — see the inline comments below
and ``capture_decode``'s module docstring for the BUG A / BUG #5 / BUG #7
address-reuse failure modes this discipline prevents.

ONE ``(max_batch, vocab)`` logits buffer every captured decode shape
slices, instead of one full-vocab buffer per shape.

Why. The decode B-ladder captures ~15 shapes; each one used to allocate
its OWN ``(B, vocab)`` persistent output. At Qwen3.5's 248,320-entry
Qwen3.5-0.8B TP1 bf16 mb64 boot, against a whole capture reserve of

Only ``max(B)`` is ever needed. Graph replay is strictly serial (see
:class:`~arbi_serve.runtime.capture.graph_exec.SleepableCUDAGraph`'s
``_replay_in_flight`` tripwire, and
:func:`~arbi_serve.engine.capture_admin.decode.refuse_verify_overlap_with_shared_capture_pool`
which boot-refuses the one config that could break that), and
:meth:`CapturedGraph.replay` already documents its output buffer as
"the next replay overwrites it" — ``borrow_output`` is built on exactly
that contract. So no two shapes' logits are ever live at once, and one

Stable VA. The buffer lives in ``capture.io_buffers``, a cuMem-backed
named pool whose VAs survive sleep/wake — the same reason
:class:`~arbi_serve.runtime.capture.dispatch.PiecewiseBuffers` and
:class:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers` are safe
to reference from a captured graph without a per-capture
``_PreCaptureMigrator`` hop. Slices must NEVER be migrated: migration
copies the bytes to a FRESH VA and returns a standalone tensor, which
would silently un-share every shape (and re-spend the memory).

Per model. ``capture.io_buffers`` is per-resident-model
(``ENGINE_MODEL_ATTRS``), so a parked member's graphs can never alias
a live member's buffer — a parked member's pools are unmapped and its
graphs are never replayed while another member is active.

The persistent buffers a single decode/verify capture references.

Every field is a persistent storage whose ``data_ptr`` the captured
graph bakes into its kernel arguments; the recurrent dicts / optional
outputs are ``None`` when the arch/shape does not use them.

Allocate every persistent buffer the captured decode graph needs.

Section 1 of :func:`~arbi_serve.runtime.capture.decode.capture_decode`.
``use_shared_verify`` (``verify_buffers is not None and S > 1``) routes
the persistent inputs through prefix slices of the engine's
:class:`VerifyBuffers` instead of fresh allocations.

The ``(B, vocab)`` prefix view, or ``None`` when this buffer does
not fit the request (caller then allocates its own — correct, just
unshared). Shape/dtype are re-checked rather than assumed: a
cross-arch pool member has a different vocab.

Captured-graph container + shape-keyed LRU pool.

Holds :class:`CapturedGraph` (persistent kernel-input/output buffers +
a captured ``torch.cuda.CUDAGraph`` + the ``replay`` refresh/replay
path) and :class:`CapturedGraphPool` (the
``(B, S, lora_bucket, is_prefill, kv_pages_bucket, prefill_context)``
→ graph LRU map
with VRAM-budget-bound + legacy count-bound eviction).

The capture *functions*
(:func:`~arbi_serve.runtime.capture.decode.capture_decode`,
:func:`~arbi_serve.runtime.capture.prefill.capture_prefill`) construct
these objects. The eviction-warning logger keeps the
``arbi_serve.runtime.capture.decode`` name (asserted by
``tests/test_cudagraph_lru.py``).

Owns persistent kernel-input/output buffers + a captured graph.

The captured graph references ``data_ptr``s of every persistent
tensor below; replay copy_()s fresh content into the same buffers
and reads the sampled-state output back from them.

Field shapes (sized for the captured ``(B, S)``):

  - ``input_ids``    : int32  ``(B*S,)``
  - ``positions``    : int32  ``(B*S,)``
  - ``slot_mapping`` : int64  ``(B*S,)``
  - ``seq_lens``     : int32  ``(B,)``
  - ``cu_seqlens_q`` : int32  ``(B+1,)``  pre-filled to [0, S, 2S, ...]
  - ``cu_seqlens_k`` : int32  ``(B+1,)``  refreshed per replay
  - ``block_table``  : int32  ``(B, max_pages_in_table)``
  - ``logits_out``   : fp32   ``(B, vocab_size)``  graph-output ref

``lora_bucket`` is the quantized active-LoRA count (``0`` for the
no-LoRA capture, ``1``/``2``/``4``/``8`` per
:data:`arbi_serve.adapters.lora.ACTIVE_LORA_BUCKETS` for LoRA-aware
captures). The captured kernels reference the engine's
:class:`arbi_serve.adapters.lora.LoraCapturePool` persistent buffers
sliced at this bucket; replay copy_()s the per-step adapter
weights into those slices.

The :class:`tkv.runtime.page_metadata.TQBufferPool` lives on the
metadata builder and is already pre-sized; the captured graph
references its scratch tensors via stable pointers.

LRU map of
``(batch_size, seq_len, lora_bucket, is_prefill, kv_pages_bucket,
prefill_context)`` → :class:`CapturedGraph`.

The third key dim is the quantized active-LoRA bucket
(:func:`arbi_serve.adapters.lora.quantize_active_loras`). LoRA-disabled
deployments only ever insert ``lora_bucket = 0`` keys; LoRA-
enabled deployments capture the boot sweep at ``lora_bucket = 0``
and a configurable subset of larger buckets so steps with active
LoRAs can also replay.

The fourth key dim is ``is_prefill`` (added with the whole-forward
prefill capture path). Decode + MTP-verify captures use
``is_prefill=False``; whole-forward prefill captures use
``is_prefill=True``. Without this dim a ``(B=1, S=128)`` prefill
graph would shadow a ``(B=128, S=1)`` decode graph keyed at the
same ``B*S`` flat-token count, and the live lookup could pick the
wrong shape.

The fifth key dim is ``kv_pages_bucket`` (Turbo prefill captured-graph over-
read fix), whose ``bucket * page_size`` upper-bounds the captured
kernel's K-gather read volume. The runner picks the smallest bucket
whose coverage exceeds the live ``max_seq_len``; smaller buckets read
fewer pages and recover the bounded block-table read win inside
captured-graph replay (the production hot path). ``0`` keys the graphs
that are not bucket-bound.

Capture currently mostly produces ``(1, 1, 0, False, bucket)``
plus a small batched-decode sweep + LoRA bucket sweep × KV-page
bucket sweep, but mixed-shape workloads can populate many keys
over time.

Eviction policy. The pool is **VRAM-budget-bound** when
``max_bytes`` is set (production): every captured shape that fits
within the graph-pool VRAM budget is retained, and :meth:`put`
only evicts the LRU entry when adding a graph would actually push
``total_memory_bytes()`` past ``max_bytes``. In this mode the
decode/verify shape space is bounded by
``max_batch × (K+1) × kv-buckets × lora`` and almost always fits,
so eviction is rare — and when it does fire it is logged LOUDLY
(which shape, current pool bytes vs budget). ``max_shapes`` is NOT
a binding constraint in budget mode (it never evicts); VRAM is.

When ``max_bytes`` is ``None`` the pool falls back to the legacy
**count-bound** LRU: bounded by ``max_shapes`` (default ``None`` =
unbounded for sleep-mode bench scripts); when set non-zero,
:meth:`put` evicts the LRU entry on count overflow. This path is
retained for unit tests and callers that have no VRAM budget.

Either way :meth:`get` marks the touched entry as recently-used.
The underlying store is an :class:`OrderedDict` so ``move_to_end``
is O(1).

Eviction releases the :class:`CapturedGraph` strong reference; the
underlying ``torch.cuda.CUDAGraph`` + persistent buffers are GC'd
once their last reference goes out of scope, and
``torch.cuda.empty_cache()`` (called by the engine after a sweep)
returns the segments to the driver.

The ``(model, params_hash)`` dimension is captured one level up by
:class:`MultiGroupCapturedGraphCache`, which holds one
:class:`CapturedGraphPool` per group. The eviction policy here
applies WITHIN a single group.

Every persistent buffer the graph references by ``data_ptr``.

Returned as a list of ``(name, tensor)`` pairs so the caller can
register them with a sleepable pool (the names land in the
:class:`SleepableTensorPool` entry name field for diagnostics).
Names are stable across captures so a pool can map updates to
replays one-to-one.

Refresh persistent buffers + replay + return cloned outputs.

For decode + MTP-verify captures (``is_prefill=False``),
``cu_seqlens_q`` is fixed at capture (``[0, S, 2S, ..., B*S]``
is invariant for uniform decode at the captured shape) and is
never re-copied. Callers do not pass ``cu_seqlens_q``.

For whole-forward prefill captures (``is_prefill=True``),
callers MAY supply ``cu_seqlens_q`` to enable padding-replay:
the captured graph at ``num_tokens=bucket_N`` accepts a request
with ``real_N <= bucket_N`` tokens by padding ``input_ids``
/ ``positions`` / ``slot_mapping`` to ``bucket_N`` and refreshing
``cu_seqlens_q = [0, real_N]``. The model's last-token gather
(``hidden[cu_seqlens_q[1:] - 1]``) then reads position
``real_N - 1`` — the correct position for sampling — instead of
the captured-baked ``bucket_N - 1`` (which would be a pad-token
position with off-distribution logits). The Turbo prefill kernel honours
the per-row Q range from ``cu_seqlens_q`` and the per-row K
length from ``seq_lens`` via its ``mSeqUsedK`` knob, so padded
positions ``[real_N..bucket_N)`` are processed for shape but
ignored for content. Padded ``slot_mapping[real_N..bucket_N)``
must point at the page-0 null sentinel so the K/V scatter does
not corrupt valid cache rows.

``return_hidden_state`` requires a graph that retained
``hidden_out``: S>1 verify captures store the per-token (N, H)
hidden; prefill mtp-fill captures store the LAST-SLOT (B, H)
hidden. Others (S=1 decode, non-fill prefill) raise.

Shared-buffer captures (:attr:`shares_verify_buffers` set) skip
the per-tensor ``copy_()`` calls below: the persistent buffers
ARE slices of the engine's :class:`VerifyBuffers`, and the
live path's :meth:`VerifyBuffers.load_step` already wrote the
per-step content into those same storages before this call.
Saves a redundant D2D copy on the verify hot path.

Output safety. ``logits_out`` and ``hidden_out`` are
persistent buffers the next replay will overwrite. We
``.clone()`` before returning so the caller can't accidentally
hold a stale view across a subsequent replay (which would
silently corrupt — the views would alias the *next* replay's
output). The clone cost (~B×V fp32 + (S>1) ~B×S×D bf16) is
worth paying for the foot-gun elimination.

Read "the next replay" as ANY next replay, not this shape's.
``logits_out`` is normally a ``[:B]`` slice of ONE ladder-wide
:class:`~arbi_serve.runtime.capture.decode_buffers.SharedLogitsBuffer`,
so every captured decode/verify shape writes the same storage —
which costs ``max(B)`` full-vocab rows instead of ``sum(B)``. That
is sound because replay is strictly serial and every return path
here clones. The one caller that does NOT clone is
``borrow_output`` below, whose contract already requires the
borrowed tensor be consumed before the next replay of anything.
``split_mixed_forward`` is the site that proves the distinction
matters: it holds a decode replay's logits across a *prefill*
replay, and is safe precisely because it takes the clone.

``borrow_output`` (decode critical-path opt-in): return the
persistent :attr:`logits_out` buffer DIRECTLY, skipping the
per-step ``.clone()`` (a fresh ``B×V×4 B`` fp32 alloc + D2D
caller can assert the returned tensor is fully consumed BEFORE
the next :meth:`replay` overwrites the buffer — the c=1/c=N
plain-decode ``execute`` path qualifies because the sampler's
argmax/sample kernel reads ``logits_out`` on the compute stream,
strictly ordered before the next replay's writes on the same
stream, and the sampled ids it produces are an INDEPENDENT
allocation (no view aliasing ``logits_out``). Rejected when
``return_hidden_state=True`` (the verify/MTP path retains hidden
and re-gathers logits across its own follow-up work — it must
keep the clone). Default ``False`` preserves the clone for every
existing caller (MTP verify, drafter, worker-bridge dispatch,
split-mixed, decode-pad, cudagraph-admin), none of which may
borrow.

A capture that recorded NO ``lm_head`` write
(:attr:`logits_recorded` ``False`` — the MTP verify shape, whose
forward runs with ``compute_logits=False``) is the one
``return_hidden_state`` case that also skips the logits clone:
the buffer holds bytes no caller reads, so cloning them is a
full ``(B, vocab)`` allocation + D2D for a discarded slot. The
hidden clone is unaffected — that IS the consumed output.

Returns:
    ``logits`` clone by default (or the borrowed persistent
    buffer when ``borrow_output=True``). ``(logits, hidden)``
    when ``return_hidden_state=True``: ``hidden`` always a
    clone, ``logits`` a clone unless the graph recorded no head
    write.

Return the captured graph for
``(batch_size, seq_len, lora_bucket, is_prefill,
kv_pages_bucket, prefill_context)`` and mark it as
recently-used. ``None`` on miss.

``is_prefill`` defaults to ``False`` so existing decode +
MTP-verify call sites retain their semantics; whole-forward
prefill lookups must pass ``is_prefill=True`` explicitly.

``kv_pages_bucket`` defaults to ``0`` (the legacy "no bucket"
sentinel — pre-bucketing captures register at this key dim
and pre-bucketing lookups continue to hit them). Bucket-aware
lookups pass the resolved bucket explicitly; the runner picks
the smallest configured bucket that covers the live
``max_seq_len``. Pre-bucketing capture paths (legacy callers,
unit tests, the prefill capture path) still produce
``kv_pages_bucket = 0`` keys and bucket-aware lookups fall
back to ``0`` when the operator-pinned bucket sweep is empty.

Register ``graph``, evicting the LRU entry only when the
binding constraint would otherwise be breached.

Inserts at the "most recently used" end so a fresh capture
survives at least one eviction round.

Budget-bound mode (:attr:`max_bytes` set): retain EVERY shape
that fits the graph-pool VRAM budget. Evict (LRU) only while
adding this graph would push ``total_memory_bytes()`` past
``max_bytes`` — and never evict the graph we just inserted (a
single graph larger than the whole budget is kept and the
over-budget condition is logged loudly, not silently dropped).
``max_shapes`` does NOT evict in this mode — VRAM is the only
binding constraint, so a fixed serving config keeps all of its
``max_batch × (K+1) × kv-bucket × lora`` decode/verify graphs
resident instead of LRU-thrashing at the count cap.

Count-bound mode (:attr:`max_bytes` is ``None``, legacy): drop
entries from the LRU end until the pool is at
:attr:`max_shapes`.

Every eviction is logged loudly (which shape, pool bytes vs
budget) — there is no silent count-cap drop of a needed graph.

After a real eviction (``did_evict``), call
``torch.cuda.empty_cache()``. Dropping the strong reference is
necessary but not sufficient — the caching allocator holds the
segments until ``empty_cache()`` runs. It returns them only for
a graph captured OUTSIDE a private ``MemPool``: ``empty_cache``
does not visit a private pool's block pools (pytorch#145168), so
a graph captured into ``capture.cudagraphs`` gives its blocks
back to that pool's own free list, where later captures reuse
them, and the physical stays mapped until the pool itself is
destroyed. Gated to the ``did_evict=True`` path so steady-state
``put`` (no overflow) stays free.

Sorted ``seq_len`` rungs ACTUALLY captured for the single-row
whole-forward prefill family ``prefill_context``.

The bucket-selection source of truth for
:func:`~arbi_serve.runtime.captured_lookup_gates.
lookup_captured_graph_for_prefill`. Selecting out of the raw
``cfg.prefill_cudagraph_buckets`` tuple instead lets the configured
ladder and the captured ladder diverge silently in BOTH directions —
a rung the sweep auto-added (the ``chunk_prefill`` cover rung) is
captured but never selectable, and a rung the ``prefill_cap`` clamp
dropped is selectable but never captured. Reading the pool cannot
diverge from the pool.

The scan is over the whole captured pool (tens of entries) and runs
once per prefill chunk — no allocation, no device work.

Snapshot of the cached graphs, oldest first — WITHOUT the LRU
bump :meth:`get` applies. Boot-time audits
(:func:`arbi_serve.engine.cudagraph_admin.assert_captured_variant_policy`)
walk every entry; touching them through :meth:`get` would
reorder eviction priority as a side effect.

Drop every captured graph + release its persistent buffers.

Teardown and pool-rebuild only. :meth:`Engine.release_memory_occupation`
does not call this: every graph's ``data_ptr`` is a stable VA that
survives release / resume, so the graphs stay valid.

Per-layer piecewise CUDAGraph dispatcher.

This is the eager / no-``torch.compile`` capture path; the optional compile
pipeline (gated by ``ARBI_COMPILE_ON``) is additive and orthogonal —
piecewise cudagraph capture runs whether or not compile is enabled.
Closes two capture gaps in one shot:

  1. **Variable-length prefill** — full-forward capture is impossible
     because attention's varlen scheduler has token-count-varying
     metadata. Per-layer captures cover the shape-stable pieces (Q/K/V
     proj, RMSNorms, MLP, residuals) plus the per-layer attention
     call; replay sees fresh per-step ``attn_meta`` because both the
     capture and the live path route through engine-owned persistent
     batch buffers (:class:`PiecewiseBuffers` — see below).
  2. **Hybrid models** (Mamba / GDN / ShortConv mixed with attention —
     LFM2, LFM2-MoE, NemotronH, Qwen3.5-MTP, Qwen3.6) —
     attention AND recurrent layers capture under piecewise prefill.
     The recurrent block prefill forwards
     (:meth:`Mamba2Block._forward_kernel_prefill`,
     :meth:`ShortConvBlock.forward` (prefill branch),
     :meth:`GDNBlock._forward_prefill_fla`) are vectorized for the
     B=1 capture bucket: per-row ``int(state_indices[i].item())``
     host syncs and per-row ``state_view.X[slab_row].copy_`` scatter
     calls are replaced by ``index_copy_`` over the persistent
     :attr:`PiecewiseBuffers.recurrent_state_indices` buffer.

Persistent-buffer redirect (the data_ptr correctness contract)
--------------------------------------------------------------
The captured graph holds the ``data_ptr``s of every tensor the
captured kernel launches reference (``cu_seqlens_q``, ``seq_lens``,
``slot_mapping``, ``block_table``, ``positions``, ``input_ids``).
At capture time the engine builds a synthetic prefill batch via
:func:`_capture_one_bucket`; at replay time the live hot path
constructs a real batch through the scheduler. Without redirect,
each side allocates fresh tensors with different ``data_ptr``s —
the captured kernels then read from the synthetic-zeros tensors
forever and KV-cache scatter / paged-attention gather walk past
their valid bounds.

The fix mirrors :class:`VerifyBuffers` (R3) for the verify-pass
capture pool and :class:`CapturedGraph`'s persistent buffers for
the whole-forward decode pool: one engine-level
:class:`PiecewiseBuffers` instance, sized worst-case at boot, holds
every per-step batch tensor at stable ``data_ptr``s. The synthetic
batch in :func:`_capture_one_bucket` references prefix slices of
those buffers, and :func:`arbi_serve.runtime.model_runner._build_batch`
``copy_()``s the live per-step content into the same prefix slices
before invoking ``model.forward``. Captured kernels see fresh
content automatically — no per-graph internal copy required.

Ping-pong hidden-state buffers (the per-bucket sharing contract)
----------------------------------------------------------------
Each captured layer-graph reads a hidden tensor and writes a hidden
tensor; both ``data_ptr``s are baked in at capture time. A naive
implementation allocates a fresh ``(hidden_in, hidden_out)`` pair
per ``(layer_idx, num_tokens, num_seqs, is_prefill)`` key — for an
8B model with 40 layers and 9 buckets, that's 360 buffer pairs ×

Since the model runs layers serially (layer K's output feeds layer
K+1's input), only TWO hidden buffers are needed per bucket. They
ping-pong across layers: layer 0 reads ``buf_a`` → writes ``buf_b``;
layer 1 reads ``buf_b`` → writes ``buf_a``; etc. Captured graphs
hardcode the ``data_ptr``s of ``buf_a`` and ``buf_b`` based on the
parity at capture time (``parity = layer_idx % 2``). At replay the
dispatcher copies the live ``hidden`` into the captured-parity read
buffer, replays, and clones the write buffer.

A single ``_PiecewiseBucketState`` (``buf_a``, ``buf_b``) per
``(num_tokens, num_seqs, is_prefill)`` bucket replaces the per-graph
buffers. Memory drops from O(num_layers × num_buckets × buf_pair)
to O(num_buckets × buf_pair) — for granite-8B that's

Captured graphs key on ``(num_tokens, num_seqs, is_prefill)``: the
captured kernel launches bake in fixed tensor shapes from the
synthetic capture batch (B=1 for prefill capture), so a live B=4
step at the same ``num_tokens`` must NOT replay the B=1 capture or
it walks past the ``cu_seqlens_q`` of length 2. ``num_seqs`` is
therefore a key dim; mismatched ``num_seqs`` falls through to the
eager body.

The dispatcher's API contract is the ``layer_call`` thunk:

    hidden = dispatch_layer(
        layer_idx=i,
        layer=layer,
        hidden=hidden,
        num_seqs=batch.num_seqs,
        layer_call=lambda h: layer(h, positions, state_view, attn_meta, ...),
    )

The thunk takes the hidden buffer as its only arg and returns the new
hidden. Per-layer extras (positions, state_view, attn_meta, rope_cache,
attn_op, lora_state, ...) are bound by closure — the dispatcher stays
backend-agnostic. On capture the dispatcher passes the bucket-state
read buffer; on replay it ``copy_()``s the live hidden into the
captured-parity read buffer and returns a clone of the write buffer.

Eligibility
-----------
``getattr(type(layer), "_cudagraph_eligible", True)``.

Recurrent layer classes opt out at class scope:

    class _NemotronHMambaLayer(nn.Module):
        _cudagraph_eligible: bool = False

When the dispatcher sees an ineligible layer it always runs
``layer_call(hidden)`` eagerly — no capture, no pool entry. The
ping-pong contract still holds across the eager hop: ineligible
layers receive a ``hidden`` (cloned from the previous capture's
write buffer) and return a fresh tensor; the next eligible layer
copies that tensor into ITS captured-parity read buffer.

Persistent + ping-pong buffers for per-layer piecewise capture.

Holds the batch-tensor and hidden-state buffer dataclasses the
piecewise dispatcher binds captured graphs against:

  * :class:`PiecewiseBuffers` — engine-level worst-case-sized batch
    tensors at stable ``data_ptr``s (the persistent-buffer redirect).
  * :class:`_SplitAttnInterBuffers` / :class:`_SplitGdnInterBuffers` —
    inter-graph buffers for the dual-graph split-attn / split-gdn
    capture paths.
  * :class:`_PiecewiseBucketState` — the per-bucket ping-pong
    ``(buf_a, buf_b)`` pair + shared CUDAGraph mem-pool handle.

See :mod:`arbi_serve.runtime.capture.dispatch` (the package docstring)
for the data_ptr-stability + ping-pong correctness contracts.

Worst-case-sized persistent batch tensors for piecewise capture.

Mirrors :class:`arbi_serve.spec_decode.verify_buffers.VerifyBuffers`
for the prefill-shape capture path: one engine-level instance
holds every per-step batch tensor (``input_ids``, ``positions``,
``slot_mapping``, ``seq_lens``, ``cu_seqlens_q``, ``cu_seqlens_k``,
``block_table``) at stable ``data_ptr``s, sized to the engine's
worst-case bounds.

Both the synthetic batch built by :func:`_capture_one_bucket`
AND the live :func:`_build_batch` reference prefix slices of
these buffers. Capture binds the graph against the persistent
``data_ptr``s; replay reads whatever the live ``_build_batch``
just ``copy_()``d into the same prefix slice.

Sizing:
    * ``max_num_tokens`` — token-flat dims (``input_ids``,
      ``positions``, ``slot_mapping``). Set to
      ``cfg.batch.max_batched_tokens``.
    * ``max_num_seqs`` — per-row dims (``seq_lens``,
      ``cu_seqlens_*``, ``block_table`` rows). Set to
      ``cfg.batch.max_batch``.
    * ``max_pages`` — block-table columns. Set to
      ``ceil(cfg.cache.max_context / cfg.cache.block_size)``.

Lifecycle. Built unconditionally in
:func:`arbi_serve.engine.build.build` when CUDA is available so
both capture (boot) and the live hot path can read through the
same storage. Phase1 sleep release drops the engine reference
along with ``verify_buffers``; resume rebuilds in the same pre-
capture sequence. Phase2 stable-VA registers each tensor with
:class:`SleepableTensorPool` so captured graphs survive
release/resume untouched.

Inter-graph buffers used by the split-attn dual-graph capture path.

Sized to ``(num_tokens, num_heads, head_dim)`` (and KV variants)
so they can hold the QKV projection output that ``_pre_attn``
produces and the attention output that ``_attn_eager`` writes
between captured-graph replays. ``data_ptr``s are baked into both
captured graphs at capture time — the pre_graph writes Q/K/V/gate
into these buffers via in-place ``copy_()``, the eager attn call
reads from them + writes attn_out via ``copy_()``, the post_graph
reads attn_out + gate from them and writes the layer-output into
the bucket's ping-pong write buffer.

Buffers live in the engine's ``graph_buffers_pool`` (NOT the
captured-graph pool) — exactly the same separation that
:attr:`_PiecewiseBucketState.buf_a` / ``buf_b`` use. Mixing
persistent allocations with capture-time intermediates in the
captured-graph pool breaks the cudagraph allocator's address-
reuse invariant (BUG A in the persistent-buffer redirect docs at
:func:`_PiecewiseBucketState.alloc`).

Inter-graph buffers used by the split-gdn dual-graph capture path.

GDN-flavored counterpart to :class:`_SplitAttnInterBuffers`. The
GDN seam is hidden-sized in/out (the eager ``_gdn_eager`` call
is the entire :class:`GDNBlock` forward — projections, conv,
FLA chunk_gated_delta_rule / fused_recurrent_gated_delta_rule,
materialize), so the inter-buffers are just two ``(N, hidden_size)``
tensors:

  * ``h_buf`` — pre_graph writes here (the input_layernorm
    output); the eager GDN call reads from it.
  * ``gdn_out_buf`` — the eager GDN call writes here (via an
    explicit ``copy_()`` since :class:`GDNBlock` returns a fresh
    ``out_proj`` allocation); post_graph reads from it.

Same persistent-pool routing as :class:`_SplitAttnInterBuffers`:
these live in the engine's ``graph_buffers_pool``, NOT the
captured-graph pool. Shared across every GDN layer in a bucket
(every GDN layer in a Qwen3.5/3.6 stack has the same hidden_size
by construction).

Shared ping-pong hidden-state buffers for one ``num_tokens`` bucket.

A bucket is keyed by ``(num_tokens, num_seqs, is_prefill)``: every
layer captured at the same key shares ONE pair of ``buf_a`` /
``buf_b`` tensors. Captured graphs at parity 0 read ``buf_a``,
write ``buf_b``; parity 1 swaps. At replay the dispatcher copies
the live hidden into the captured-parity read buffer and reads
the write buffer back via clone (see :meth:`LayerCapturedGraph.replay`).

Without sharing, each captured layer-graph would carry its own
``(hidden_in, hidden_out)`` pair — for an 8B model with 40 layers
memory before any work happens. With sharing, the pool fits in

Per-layer kernel intermediates allocated INSIDE
``torch.cuda.graph(graph, pool=...)`` are also shared across
layer captures within the same bucket via :attr:`mem_pool`: the
first capture creates a fresh CUDAGraph mem-pool handle, every
subsequent capture in this bucket passes the same handle so
torch's caching allocator overlays dead allocations from prior
captures (CUDAGraph trees memory-management, see torch docs at
https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs).
Without this, each layer's per-capture intermediates would be
distinct residual.

Allocated once per bucket by
:meth:`LayerCapturedGraphPool.ensure_bucket_state` (called from
:func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`
BEFORE the per-layer capture sweep at this bucket); released
when the last :class:`LayerCapturedGraph` referencing this bucket
is evicted from the LRU pool (see :meth:`LayerCapturedGraphPool.put`).

H2D that reads the pinned
``h_*`` staging has EXECUTED on-device.

Call BEFORE writing any ``h_*`` buffer for a new step. Waits on
the event recorded by the previous step's
:meth:`mark_staging_consumed`; a no-op when nothing is armed
(first step, CPU stub) or the event already fired (host slower
than the GPU — the common case, ~µs). The wait covers only the
copies themselves (they sit at the FRONT of the step's stream
work), so the host still overlaps the previous step's forward.

The event is queried before it is waited on, so
:data:`_STAGING_REUSE_BLOCK` separates the two outcomes and
:attr:`staging_block_ns` charges the host stall to the steps that
actually paid it. The query replaces the wait in the non-blocking
case, so the instrumentation is not an addition to the hot path.

Record the staging-consumed point on the CURRENT stream.

Call right after the last ``copy_(non_blocking=True)`` whose
SOURCE is an ``h_*`` pinned buffer, on the same stream that
issued it. The next :meth:`wait_staging_free` then blocks until
those copies executed. No-op on CPU.

Allocate the q/k/v/gate/attn_out buffers under ``buffers_pool``.

The pool routing matches :meth:`_PiecewiseBucketState.alloc` —
persistent capture-input buffers MUST live in a mempool
DISTINCT from the cudagraph capture pool. Fresh ``torch.zeros``
so first-replay fences see deterministic content (the live
path overwrites every cell on every step before any capture
replay reads it).

Allocate ``buf_a`` / ``buf_b`` under the supplied named pool.

``buffers_pool`` is the engine's ``graph_buffers_pool`` —
persistent capture-input buffers MUST live in a mempool
DISTINCT from the cudagraph capture pool ``graph_pool``.
Mixing persistent allocations with capture-time intermediates
in the same pool breaks the cudagraph allocator's address-
reuse invariant and corrupts captured-graph replays (BUG A).
Falls back to ``graph_pool`` when ``buffers_pool`` is None
for backwards compatibility with legacy callers / tests.

Allocate the persistent per-replay input buffers for split capture.

Currently the RoPE ``positions_buf`` (the split-attn corruption
fix — see the attribute docstring). Allocated under
``buffers_pool`` (the engine's ``graph_buffers_pool``) so it
sits OUTSIDE the cudagraph capture pool, matching ``buf_a`` /
``buf_b`` and the split inter-buffers. Pre-filled with a valid
synthetic ``arange(num_tokens)`` so the FIRST capture-time
``cos[positions_buf]`` gather indexes a self-consistent
position ramp (the live path overwrites ``[:real_num_tokens]``
before any replay reads it).

Idempotent — re-invoking is a no-op once ``positions_buf`` is
allocated. Called by
:func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`
BEFORE the per-layer split capture sweep so the captured
graphs bind against a stable ``data_ptr``.

Allocate (or fetch) split-attn inter-graph buffers for one
attention shape inside this bucket.

Keyed by ``(num_heads, num_kv_heads, head_dim)`` — homogeneous
attention models reuse one entry across every attention layer
in the bucket. Mismatched dtype against an existing entry
raises (engine bug; the bucket's primary dtype already
constrains it).

Idempotent. Splits the buffer cost across as many distinct
attention shapes as the model carries; each entry is bounded
in size at ``O(num_tokens × max_heads × head_dim × dtype)``
which is small relative to the captured-pool
footprint (e.g. Qwen3.5-0.8B at N=2048 / 8 heads × 256

Allocate (or fetch) the split-gdn inter-graph buffer set.

Idempotent. Sized to ``(num_tokens, hidden_size)`` from this
bucket's primary fields — every GDN layer in a Qwen3.5/3.6
stack shares the same hidden_size, so one entry per bucket
is sufficient. Mismatched ``dtype`` against an existing entry
raises (engine bug; the bucket's primary dtype constrains it).

Boot-time per-layer CUDAGraph capture functions.

Holds the three ``capture_layer*`` entry points driven by the
boot-time precapture sweep
(:func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`):

  * :func:`capture_layer` — whole-forward single-graph capture.
  * :func:`capture_layer_split` — split-attn dual-graph (pre/post).
  * :func:`capture_layer_split_gdn` — split-gdn dual-graph (pre/post).

See :mod:`arbi_serve.runtime.capture.dispatch` (the package docstring)
for the data_ptr-stability + ping-pong correctness contracts.

Capture one layer's forward at one ``(num_tokens, num_seqs)`` bucket.

Picks a parity from ``layer_idx % 2`` (even reads ``buf_a`` / writes
``buf_b``; odd reads ``buf_b`` / writes ``buf_a``) so consecutive
captured layers ping-pong through the bucket's shared ``(buf_a,
buf_b)`` pair without per-layer hidden buffers. Runs two warmup
forwards (the JIT first-call paths must not fire under
``torch.cuda.graph``), then captures the third pass against the
bucket state's parity-selected read buffer. The returned
:class:`LayerCapturedGraph` carries the parity + a reference to
the shared bucket state.

Captures within the same bucket SHARE the CUDAGraph mem-pool
handle (via ``bucket_state.mem_pool``). The first capture creates
a fresh handle and stores it on the bucket; subsequent captures
pass it back into ``torch.cuda.graph(graph, pool=...)`` so torch's
caching allocator overlays per-capture intermediates from prior
captures. Without this, every layer's per-capture residual
(RMSNorm scratch, attention Q/K/V projections, MLP swiGLU
intermediates, the layer-output tensor) would be distinct
captured memory and the net pool footprint would be O(num_layers
× per_capture_residual), defeating the per-bucket sharing.

Capture writes into the parity-selected WRITE buffer; the layer-
call thunk's return value flows into the bucket's persistent
``write_buf`` via a final captured ``copy_()`` so the written
``data_ptr`` is always the bucket's stable address. Replay
``copy_()``s the live hidden into ``read_buf``, replays the
captured graph (which re-issues the ``copy_(write_buf, out)``
node), and clones ``write_buf``.

The ``layer_call`` thunk MUST consume its argument as the hidden
input and return the new hidden. The thunk's closure should
already reference :class:`PiecewiseBuffers` slices for the
per-step ``attn_meta`` tensors (see :func:`_capture_one_bucket`);
replay sees fresh values automatically because :func:`_build_batch`
writes into the same persistent storage every step.

TP-aware capture (``capture_stream``). Under TP>1 the layer's
``RowParallelLinear.o_proj`` (and ``down_proj`` / MoE down-shard)
issues ``dist.all_reduce`` to sum the per-rank partials. NCCL's
first call on a ``(communicator, stream)`` pair lazily allocates
internal communication buffers via ``cudaMalloc``, which
``torch.cuda.graph`` rejects mid-capture
(``cudaErrorStreamCaptureUnsupported`` / device-side assert).
The orchestrator (:func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`)
enters :func:`arbi_serve.distributed.graph_capture.graph_capture`
ONCE PER BUCKET to (a) pre-warm NCCL on a side stream and (b)
pin :class:`GroupCoordinator._capture_stream` so every per-layer
``all_reduce`` enqueues on that side stream rather than the
default. ``capture_stream`` threads the orchestrator's side
stream into this function so every layer's
``torch.cuda.graph(layer_graph, stream=capture_stream)`` records
against the same stream the AR enqueues onto — the cross-rank
op gets baked into the captured layer graph instead of crashing
capture. REQUIRED at every world size: torch's caching allocator
keys free-block reuse by stream, so captures on per-layer fresh
streams strand every capture's dead transient working set in the
shared pool (reserved grows as the SUM across captures instead of
the MAX). ``None`` raises.

Capture one split-attn layer's pre/post pieces as TWO cuda graphs.

Splits :func:`capture_layer`'s body around the eager ``_attn_eager``
seam. Inter-graph buffers (q / k / v / gate / attn_out) live on
:attr:`_PiecewiseBucketState.split_attn_buffers` keyed by the
layer's attention shape. The pre_graph writes into them via
captured ``copy_()`` nodes; the post_graph reads from them via
the same persistent-data_ptr contract used for ``read_buf`` /
``write_buf``.

The eager ``_attn_eager`` call is NOT captured — it runs between
pre_graph.replay() and post_graph.replay() at every replay. Its
workspace allocations (FA varlen scratch, paged_attention
intermediates, the freshly-allocated out_buf in the layer's
body) land in the regular allocator pool, NOT the captured-graph
mempool.

``block_args`` matches the same tuple :func:`capture_layer`
consumes — ``(positions, batch_meta, state_view, rope_cache,
attn_ops, residual_buf)`` for Qwen3.5 attention layers. We pull
the pieces the layer's pre/eager/post methods actually read out
of that tuple by name (positions, residual_buf, rope_cache for
pre; state_view, batch_meta, attn_op for eager).

Capture one split-gdn layer's pre/post pieces as TWO cuda graphs.

GDN-flavored counterpart to :func:`capture_layer_split`. Splits
the GDN decoder layer's body around the eager ``_gdn_eager`` seam
(=the entire :class:`GDNBlock` forward). Inter-graph buffers
(``h``, ``gdn_out``) live on
:attr:`_PiecewiseBucketState.split_gdn_buffers` keyed only on
the bucket's primary ``hidden_size`` (one entry per bucket).

The eager ``_gdn_eager`` call is NOT captured — it runs between
pre_graph.replay() and post_graph.replay() at every replay. Its
workspace allocations (FA chunk-fwd intermediates / fused-recurrent
state, conv buffer, projection results, the ``out_proj`` allocation)
land in the regular allocator pool, NOT the captured-graph
mempool.

``block_args`` matches the same tuple :func:`capture_layer`
consumes — ``(positions, batch_meta, state_view, rope_cache,
attn_ops, residual_buf)`` for Qwen3.5 GDN layers. We pull
state_view / batch_meta / residual_buf out of that tuple by name
(positions / rope_cache / attn_ops are unused by GDN blocks).

Piecewise prefill cudagraph-coverage hit/miss accounting.

A PERMANENT (always-on) fail-loud guardrail for the piecewise prefill
dispatch path. arbi's no-silent-footguns rule (NORTH STAR) requires
that a bucket miss that drops a multi-thousand-token prefill forward to
EAGER can NEVER pass unnoticed: a future shape / scheduler change that
re-opens a coverage hole must surface in production logs + metrics, not
silently regress TTFT.

The dispatcher
(:func:`arbi_serve.runtime.capture.dispatch.dispatch.dispatch_layer_args`)
records every PREFILL layer call here: the live ``(num_tokens,
num_seqs)`` key + which of THREE outcomes it took. The record is a
single ``dict`` increment on the steady-state path (no kernel, no
allocation, no device sync) — negligible against the many-millisecond
prefill forward it gates.

THREE outcomes (the cap-and-eager design: capture what fits, run the
rest eager, never chunk):

  * **HIT** — replayed a captured graph (incl. via PAD-UP to the next
    captured rung). The cudagraph fast path was taken.
  * **too-large-eager** — ``num_tokens`` EXCEEDS the top captured rung
    (the VRAM-feasible cap; large rungs OOM the post-capture gate on a
    forwards run EAGER by design). This is a NAMED, COUNTED,
    NON-WARNING bucket: deliberate, not a footgun.
  * **unexpected-miss** — ``num_tokens`` is WITHIN the captured range
    (``<=`` the top rung) but STILL missed the pool (and pad-up found no
    rung ``>=`` it). This SHOULD have hit a captured rung and did not —
    a real coverage hole. It logs a WARNING + increments the Prometheus
    miss counter (one line per shape, rate-limited).

The distinction matters: ``too-large-eager`` is expected at high
concurrency (c64/ISL=8192 packs >2048 tokens per forward); an
``unexpected-miss`` is a regression. Acceptance asserts
``unexpected_misses == 0`` (not ``too_large_eager == 0``).

Why prefill-only: decode coverage has its own pad-up ladder + the
``ARBI_DEBUG_CAPTURE_LOOKUP`` histogram. Prefill is where the
long-context high-batch TTFT gap lives (a num_tokens up to
``max_batched_tokens`` × num_seqs > 1 forward off the ``(1..512,
num_seqs=1)`` boot ladder ran fully eager).

Surfaces:

  * :func:`record` — called from the dispatch miss/hit branch.
  * :func:`snapshot` — the ``{(num_tokens, num_seqs): [hits, misses]}``
    distribution + totals; the ground-truth dump + acceptance test
    (``total_misses == 0``).
  * :func:`reset` — zero the counters (between benchmark phases).
  * the first miss of each distinct shape logs a WARNING (rate-limited
    to one line per shape so a sustained hole doesn't flood the log).

A Prometheus counter
(``arbi_piecewise_prefill_cudagraph_misses_total``) is incremented in
lock-step when ``prometheus_client`` is importable so the guardrail is
visible on the Grafana board, not just the log.

Record one prefill layer dispatch (one of four outcomes).

* ``hit=True`` → replayed a captured graph (incl. via pad-up).
* ``hit=False, too_large=True`` → ``num_tokens`` exceeded the top
  captured rung (the VRAM-feasible cap) and ran EAGER BY DESIGN —
  the cap-and-eager path. Counted in the ``too_large_eager`` bucket
  (named, non-WARNING — deliberate, not a footgun).
* ``hit=False, pad_waste=True`` → the smallest covering rung sits at
  or beyond the pad-waste bound (``graphs._PAD_WASTE_BOUND`` x the
  real width), so the forward ran EAGER BY DESIGN rather than pay
  the padded rung's fake GEMM work. Named + counted, non-WARNING.
* ``hit=False, too_large=False, pad_waste=False`` → an UNEXPECTED
  miss: a shape within the captured range that should have hit a
  rung but did not. The first of each distinct ``(num_tokens,
  num_seqs)`` shape logs a single WARNING; the Prometheus
  miss-counter increments.

Return the current coverage distribution + totals.

Shape::

    {
      "shapes": {(num_tokens, num_seqs): {"hits": h,
                   "too_large_eager": t, "unexpected_misses": m}, ...},
      "total_hits": int,
      "total_too_large_eager": int,
      "total_unexpected_misses": int,
      "unexpected_miss_shapes": [(num_tokens, num_seqs), ...],
      "too_large_eager_shapes": [(num_tokens, num_seqs), ...],
    }

The acceptance gate is ``total_unexpected_misses == 0`` —
``total_too_large_eager`` is the deliberate above-the-cap eager
bucket and is expected to be non-zero at high concurrency / long
context.

Install the SIGUSR1 → coverage-dump handler (idempotent, main-thread).

Called from the engine build path so the running server can be
poked (``kill -USR1 <pid>``) to emit the live prefill coverage
distribution + write it to the diagnostic dump path (see
:func:`_dump_path` / :data:`_DUMP_PATH_ENV`).

Per-layer piecewise dispatch hot path.

The steady-state entry points the model forward calls per layer:

  * :func:`model_dispatch` / :func:`model_dispatch_args` — read the
    piecewise context off the model and delegate.
  * :func:`dispatch_layer` / :func:`dispatch_layer_args` — captured-
    replay or eager pass-through on a pool hit/miss.

plus the default bucket set (:data:`DEFAULT_PIECEWISE_BUCKETS`). See
:mod:`arbi_serve.runtime.capture.dispatch` (the package docstring) for
the full design — data_ptr stability, ping-pong buffers, eligibility.

The ``num_seqs`` value to KEY the captured-graph pool on.

Collapses the ``num_seqs`` axis to :data:`_PREFILL_CANON_NUM_SEQS` for
PREFILL split-attn / split-gdn layers (their captured pieces are
``num_seqs``-independent — see the constant's docstring); returns the
live ``num_seqs`` unchanged otherwise (decode, or whole-forward prefill
that bakes the ``num_seqs``-dependent attention op into the graph).

Positional fast-path counterpart to :func:`model_dispatch`.

Identical semantics; the layer call is ``layer(hidden, *block_args)``
rather than ``layer_call(hidden)`` — no closure-as-arg, no nested
``def``. This is the entry point Dynamo traces through under
``ARBI_COMPILE_ON=1``: the model's ``_dispatch_layer`` bound method
delegates here, the eager fast path (when ``_piecewise_pool`` is
None) returns ``layer(hidden, *block_args)`` directly with no
intermediate Python callable.

Positional-args variant of :func:`dispatch_layer`.

Eager fast path (the steady-state production hot path with
``ARBI_COMPILE_ON=1`` AND ``_piecewise_pool=None``) is one call:
``layer(hidden, *block_args)``. No layer_call thunk constructed.
Capture / replay branches share the legacy mechanism; only the
eager fall-through differs.

Convenience wrapper that reads piecewise context off the model.

Engine-set attributes (all optional, default to off):

  * ``_piecewise_pool``: :class:`LayerCapturedGraphPool` that
    captures land in / replays come from. ``None`` disables
    piecewise capture entirely (every call is eager).
  * ``_piecewise_capture_on_miss``: ``True`` during precapture
    sweep, ``False`` (default) during the hot path.
  * ``_piecewise_capture_stream``: optional :class:`torch.cuda.Stream`
    the bucket-level orchestrator pinned via
    :func:`arbi_serve.distributed.graph_capture.graph_capture`.
    Set during the boot-time precapture sweep so every layer's
    ``capture_layer`` records its CUDAGraph against the same
    stream NCCL's all-reduces enqueue on; cleared on the steady-
    state hot path.

Required keys propagated to the dispatcher:

  * ``num_seqs`` — live batch's ``num_seqs``
    (``batch.seq_lens.shape[0]``).
  * ``is_prefill`` — live batch's ``is_prefill`` flag. Capture
    runs at ``is_prefill=True``; live decode / verify steps
    miss the pool and fall through to the eager body.

All four engine attributes default to safe no-ops; a model whose
engine has not enabled piecewise capture pays exactly one
``getattr`` per attribute and then calls the eager ``layer_call``
thunk.

Per-layer captured-replay or eager pass-through.

Lookup order:

  1. Layer ineligible (``_cudagraph_eligible = False`` on the
     class) → eager. No capture, no pool entry.
  2. ``pool is None`` (engine running without piecewise capture
     enabled) → eager.
  3. Pool hit on ``(layer_idx, hidden.shape[0], num_seqs,
     is_prefill)`` → replay against the bucket-state
     parity-selected buffer + return a clone.
  4. Pool miss + ``capture_on_miss=True`` (precapture sweep) →
     capture the layer at this ``(num_tokens, num_seqs,
     is_prefill)`` bucket using the existing bucket state
     (allocated by :func:`_capture_one_bucket` before the sweep),
     register, replay.
  5. Pool miss + ``capture_on_miss=False`` (steady-state hot
     path) → eager. The engine's pre-capture sweep
     (:func:`arbi_serve.engine.cudagraph_admin.precapture_layer_graphs`)
     pre-populates the bucketed token counts at ``B=1``; steps
     at B>1 (batched decode) replay the whole-forward decode
     pool instead, and steps at non-bucketed N run eager.
     Lazy capture mid-step would synchronize the device and
     torque per-step latency.

The ``hidden`` arg is the live activation tensor; returning a
captured-graph clone preserves the tensor-identity contract the
caller's residual stitching depends on (subsequent uses of
``hidden`` keep their own storage).

``capture_stream`` is the bucket-level orchestrator's shared
capture stream (set by
:func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`
while it holds a :func:`graph_capture(device)` context). On
pool-miss capture, it is forwarded into :func:`capture_layer` so
every per-layer ``torch.cuda.graph`` records against the ONE
stream NCCL's ``all_reduce`` enqueues onto and the allocator's
stream-keyed free-block reuse overlays per-capture transients.
``None`` with ``capture_on_miss=True`` raises inside
:func:`capture_layer`; on the steady-state hot path (no capture)
it is never read.

Captured-layer graph owners + the LRU graph pool.

Owns the three non-polymorphic captured-graph classes
(:class:`LayerCapturedGraph`, :class:`LayerSplitCapturedGraph`,
:class:`LayerSplitGdnCapturedGraph`) and the
:class:`LayerCapturedGraphPool` LRU map keyed on
``(layer_idx, num_tokens, num_seqs, is_prefill)`` plus its parallel
:class:`_PiecewiseBucketState` ping-pong-buffer pool.

See :mod:`arbi_serve.runtime.capture.dispatch` (the package docstring)
for the data_ptr-stability + ping-pong correctness contracts.

Owns one captured-layer graph + the parity it captured at.

Capture binds the graph against the bucket-state read buffer's
``data_ptr`` AND the engine-level :class:`PiecewiseBuffers`'s
per-step batch tensors — replay ``copy_()``s the live hidden
into the bucket-state read buffer (selected by ``parity``),
replays, and clones the bucket-state write buffer. The per-step
batch tensors are refreshed OUT-OF-BAND by :func:`_build_batch`
(which writes into the same persistent storage before forward),
so :meth:`replay` itself only handles the hidden tensor.

The captured shape includes ``num_seqs`` (B) AND ``is_prefill``
in addition to ``num_tokens`` (N): kernel launches bake fixed
tensor shapes for ``cu_seqlens_q[:B+1]`` etc. at capture time
AND the prefill-vs-decode dispatch picks DIFFERENT kernels (Turbo
prefill varlen vs paged-decode kernel). A live decode at the
same (B, N) cannot replay a captured prefill graph, and vice
versa. Pool lookup is therefore keyed on
``(layer_idx, num_tokens, num_seqs, is_prefill)``. The MTP
verify pass — which runs at ``is_prefill=False`` with
``mtp_block_m > 1`` — also misses by this key (its captures, if
any, would live in the whole-forward decode pool, not here).

Hidden buffers live on a SHARED :class:`_PiecewiseBucketState`
keyed at the same ``(num_tokens, num_seqs, is_prefill)`` (every
layer captured at one bucket shares the bucket's ping-pong pair).
See :class:`_PiecewiseBucketState` for the memory-savings
rationale.

Owns a pre/post graph PAIR for one split-attn layer at one bucket.

Captures :meth:`_pre_attn` and :meth:`_post_attn` as separate
cuda graphs with the eager :meth:`_attn_eager` call running
between replays. Inter-graph data flow:

  * pre_graph — reads from the bucket's parity-selected read
    buffer, writes (q, k, v, gate) into the bucket's per-shape
    ``_SplitAttnInterBuffers`` entry via in-place ``copy_()``
    nodes baked into the captured graph.

  * eager attn — reads (q, k, v) from inter-buffers, allocates
    its own ``out_buf`` (lands in the regular allocator, NOT the
    captured-graph mempool), writes the attention output back
    into ``inter.attn_out_buf`` via an explicit ``copy_()``.

  * post_graph — reads (attn_out_buf, gate_buf) from inter-
    buffers, writes the layer-output into the parity-selected
    write buffer via the same ``write_buf.copy_(out)`` pattern
    :class:`LayerCapturedGraph` uses.

Captured-pool footprint vs whole-forward capture: pre/post
together leave OUT the attn op's softmax / FA varlen scratch /
paged-attention intermediates / output allocation — these run
eager between replays and live in the regular allocator.

Same parity / bucket-state contract as
:class:`LayerCapturedGraph` — pre_graph reads buf_a (parity 0)
or buf_b (parity 1); post_graph writes the OPPOSITE buffer (so
the layer's net behaviour is "read parity-0 read_buf → write
parity-0 write_buf", matching the legacy single-graph contract).

Owns a pre/post graph PAIR for one split-gdn layer at one bucket.

GDN-flavored counterpart to :class:`LayerSplitCapturedGraph`.
Captures :meth:`_pre_gdn` and :meth:`_post_gdn` as separate
cuda graphs with the eager :meth:`_gdn_eager` call (the entire
:class:`GDNBlock` forward) running between replays. Inter-graph
data flow:

  * pre_graph — reads from the bucket's parity-selected read
    buffer, writes ``h`` (the input_layernorm output) into the
    bucket's :class:`_SplitGdnInterBuffers` ``h_buf`` via an
    in-place ``copy_()`` baked into the captured graph.

  * eager GDN — reads ``h`` from ``inter.h_buf``, runs the
    :class:`GDNBlock` forward (projections, conv, FLA chunk-fwd
    / fused-recurrent kernel, materialize), copies the
    ``out_proj`` output into ``inter.gdn_out_buf`` so post_graph
    reads from the data_ptr it baked at capture time.

  * post_graph — reads ``gdn_out`` from ``inter.gdn_out_buf``
    and (under cross-layer fusion) ``residual_buf``, writes the
    ``mlp_out`` (the pending-add for the next layer) into the
    parity-selected write buffer.

Captured-pool footprint vs whole-forward GDN capture: pre/post
together leave OUT the FLA chunk-fwd workspace, projection
intermediates, conv buffer alloc, and the ``core_attn_out``
return — these run eager between replays and live in the
regular allocator pool. On Qwen3.5-0.8B (18 GDN + 6 attention
layers) the FLA workspace is pinned to a separate
NamedMemPool, but the projection / materialize intermediates and
``core_attn_out`` allocation would otherwise land in the
captured-graph pool — split-gdn moves them out.

Same parity / bucket-state contract as
:class:`LayerSplitCapturedGraph` and :class:`LayerCapturedGraph`
— pre_graph reads buf_a (parity 0) or buf_b (parity 1);
post_graph writes the OPPOSITE buffer.

LRU map of ``(layer_idx, num_tokens, num_seqs, is_prefill)`` →
:class:`LayerCapturedGraph`, plus a parallel pool of shared
:class:`_PiecewiseBucketState`s keyed at
``(num_tokens, num_seqs, is_prefill)``.

Layer graphs are LRU-evicted; their bucket states are released
when their last referencing layer-graph is evicted (so the
bucket's ``(buf_a, buf_b)`` pair returns to the allocator).

Keyed on ``num_seqs`` (B) AND ``is_prefill`` in addition to
``num_tokens`` (N):

* Kernel launches captured at one B bake fixed tensor shapes for
  ``cu_seqlens_q[:B+1]`` etc. and cannot replay at a different B.
* The prefill-vs-decode dispatch picks different kernels (Turbo
  prefill varlen vs paged-decode kernel) so a captured prefill
  graph is not interchangeable with a captured decode graph at
  the same (B, N).

The boot-time precapture sweep covers ``B=1`` prefill (the
canonical piecewise-target shape — variable-length prefill is
the gap whole-forward decode capture cannot close); decode
steps fall through to the eager body and the whole-forward
decode pool covers them via :class:`CapturedGraph`.

Bounded by ``max_shapes`` (total layer-graph entries across
keys). Default ``None`` leaves it unbounded; the engine
constructor sets a non-zero cap so production deployments don't
bloat the pool indefinitely.

Eviction releases the strong reference; the underlying CUDAGraph
GCs when its last reference goes out of scope. The named
``graph_pool`` MemPool returns the segment to the driver when
``empty_cache()`` runs.

Copy ``hidden`` into the parity-selected read buffer, replay,
return a clone of the parity-selected write buffer.

Sizes must match the captured ``num_tokens``; non-conforming
token counts must take the eager path (the caller decides via
bucket lookup before reaching :meth:`replay`).

We ``.clone()`` the write buffer so the caller can safely
chain into the next layer (whose dispatch may also replay
against the same bucket state at the OPPOSITE parity, which
would overwrite our write buffer when its replay runs). The
clone is one ``hidden_size × num_tokens × dtype`` allocation
per layer-replay — handful of microseconds for typical
granite-8B / qwen3-4B sizes.

The per-step ``attn_meta`` tensors the captured kernels read
are NOT refreshed here: they live on the engine-owned
:class:`PiecewiseBuffers` and are populated by
:func:`_build_batch` once per step before ``model.forward``
runs. This replay is one ``copy_()`` + one ``graph.replay()``
+ one ``clone()``; the live ``attn_meta`` is already at the
right ``data_ptr``s the captured kernel launches reference.

Run pre_graph → eager attn → post_graph and return the layer output.

``per_layer_view`` / ``attn_meta`` / ``attn_op`` are the LIVE
arguments the eager ``_attn_eager`` consumes — they are
fetched fresh per step (same contract as the persistent-
buffer redirect for the captured kernels' ``attn_meta``
tensor data_ptrs).

``positions`` is the live forward's per-seq-reset RoPE
positions (``(real_num_tokens,)`` int). The captured
``_pre_attn`` applies RoPE against the bucket's PERSISTENT
``positions_buf`` (a stable ``data_ptr`` baked at capture);
we refill ``positions_buf[:real_num_tokens]`` with the live
positions HERE — BEFORE ``pre_graph.replay()`` — so the
captured ``cos[positions_buf]`` gather sees this step's actual
positions. THIS is what lets ONE capture (at the canonical
``num_seqs``) replay correctly for any concurrent
``num_seqs > 1`` batch without cross-sequence RoPE
contamination — mirroring the decode cudagraph path's
persistent-input-buffer discipline. ``None`` skips the refill
(legacy / pure-test callers + models with no RoPE); the
captured gather then reads whatever the buffer last held.

PAD-UP. ``real_num_tokens`` (when set and ``< self.num_tokens``)
is the live forward's TRUE token count: ``hidden`` was padded up
to this rung's ``num_tokens`` by the dispatcher. The captured
pre/post graphs run at the full padded width (token-flat matmuls /
norms — the padding rows are isolated), but the EAGER attention op
runs on the ``[:real_num_tokens]`` SLICE of q/k/v so it sees
exactly the rows the live ``attn_meta`` / ``cu_seqlens`` describe.
The padding rows' q/k/v are never read by attention and their
post-attn layer-outputs are never sampled. ``None`` (or ``==
num_tokens``) is the exact-rung path (no slicing). The live
``positions`` refill below writes only ``[:real_num_tokens]``;
padding rows keep their synthetic capture-time arange (never
read by the eager attention, never sampled).

We ``.clone()`` the write buffer for the same reason
:class:`LayerCapturedGraph.replay` does (next layer's replay
may overwrite the bucket's write buffer at the opposite
parity).

Run pre_graph → eager GDN → post_graph and return the layer output.

``state_view`` and ``batch_meta`` are the LIVE arguments the
eager ``_gdn_eager`` consumes. The GDN block treats
``state_view`` as the per-layer slab handle directly — pre-
resolved by :meth:`LayerStack._resolve_per_layer_views` from
:meth:`MultiStatePool.per_layer_views` (avoids Dynamo
de-specialization; having the block call
``state_view.layer_view(layer_idx)`` itself would force a
recompile per ``layer_idx``). The captured kernels in
pre_graph + post_graph use the persistent-buffer ``data_ptr``s
baked at capture time; only the eager middle sees fresh per-
step values.

PAD-UP. ``real_num_tokens`` (when set and ``< self.num_tokens``)
is the live forward's TRUE token count: ``hidden`` was padded up
to this rung. The captured pre/post graphs run at the full padded
width; the EAGER GDN mixer runs on the ``[:real_num_tokens]`` slice
of ``h_buf`` so its conv + chunk-recurrent kernel sees exactly the
rows the live ``batch_meta`` (``cu_seqlens`` / ``state_indices``)
describes — the recurrent state is NOT contaminated by padding
tokens (they fall outside every sequence's ``cu_seqlens`` span).
The padding rows' GDN output stays stale in ``gdn_out_buf`` and
post_gdn's garbage layer-outputs there are never sampled.

Allocate (or fetch) the shared ``(buf_a, buf_b)`` pair for one
``(num_tokens, num_seqs, is_prefill)`` bucket.

Called from :func:`arbi_serve.engine.cudagraph_admin._capture_one_bucket`
BEFORE the per-layer capture sweep at this bucket so that
every layer's :func:`capture_layer` references stable
``data_ptr``s on ``buf_a`` / ``buf_b``. Idempotent — re-
invoking with the same key returns the existing state.

Mismatched ``hidden_size`` / ``dtype`` against an existing
state for the same key is an engine bug (the bucket key
already encodes shape constraints) and raises so the boot
path fails loud rather than corrupting the captured graphs.

Sorted captured PREFILL ``num_tokens`` rungs at ``num_seqs``.

Used by the dispatcher's PAD-UP path: a live prefill forward at an
off-rung ``num_tokens`` pads its hidden up to the smallest rung
``>=`` it and replays that rung's captured graphs (the captured
pre/post pieces are token-flat; the padding rows fall outside the
live ``cu_seqlens`` so the eager attention / GDN middle — which runs
on a ``[:real_num_tokens]`` slice — never reads them, and their
garbage layer-outputs are never sampled). Derived from the bucket-
state pool so it reflects exactly which rungs the boot sweep
actually captured (a rung that OOM'd mid-sweep is absent). Cached
per ``num_seqs`` and invalidated whenever a bucket state is added /
dropped (the pool is otherwise stable after boot).

SINGLE source of truth for the prefill pad-up width.

Returns the smallest captured prefill rung ``>= num_tokens`` whose
per-layer capture GRAPHS are actually present AND whose width stays
under ``num_tokens * _PAD_WASTE_BOUND`` (the pad-waste floor) — so
the dispatcher replays it — else ``num_tokens`` unchanged (the
dispatcher runs the layer EAGER at its real width). BOTH the per-layer hidden pad-up
(:func:`dispatch_layer_args`) AND the model's cross-layer
``residual_buf`` sizing
(:meth:`arbi_serve.models.qwen3_5.Qwen3_5Model._residual_pad_rung`)
call THIS, so the padded-hidden width and the residual-slab width
can never drift. A drift makes the fused ``input_layernorm`` add a
residual of the WRONG ``num_tokens`` against the hidden
(``fused_add_rms_norm`` size mismatch → out-of-bounds → CUDA illegal
access — the ``c>=16`` mixed decode+prefill tkv crash).

Distinct from :meth:`prefill_rungs`, which reports the rung set from
the bucket-STATE pool: a rung can carry an allocated bucket state
(so it shows up in ``prefill_rungs``) while its per-layer capture
graphs are absent or only PARTIALLY present — the boot sweep
captured the rung for SOME decoder layers but not all (an
interrupted / OOM'd sweep, or a rung that only a subset of layers
reached). The dispatcher's per-layer ``get()`` then hits for the
covered layers and MISSES (→ eager at real width) for the rest.
Because the model sizes ONE shared ``residual_buf`` for the whole
layer stack, a rung that is only partially covered is
UNSATISFIABLE: the covered layers replay at rung width (need the
residual at rung width) while the uncovered layers run eager at
real width (need it at real width). So the ONLY residual width that
is correct for EVERY layer at a partially-covered rung is the real
``num_tokens`` — which requires the whole forward to run eager
(no pad-up). This method therefore returns a rung ONLY when it is
FULLY covered — captured for every eligible prefill layer — so the
model and the dispatcher make the SAME all-or-nothing pad decision
for the whole forward (uniform pad-to-rung replay, or uniform
eager-at-real), and the residual width can never drift from the
per-layer hidden width. A partially-covered or absent rung falls to
eager (returns ``num_tokens``); dispatch's per-layer ``get()`` then
also misses uniformly, so every layer runs eager at real width.

Register ``graph`` and evict LRU entries on overflow.

After eviction, call ``torch.cuda.empty_cache()`` so the
evicted graph's reserved segments return to the driver.
Bucket states whose last referencing graph was just evicted
are also dropped (their ``(buf_a, buf_b)`` pair GCs).

Gated to the ``did_evict=True`` path so steady-state ``put``
stays free.

Drop every captured graph + every bucket state.

Releases both the per-graph capture residual AND the shared
ping-pong buffer pairs. Sleep-mode release uses this hook
to drop the entire piecewise pool before resume rebuilds.

CUDAGraph capture + replay for the MTP drafter chain.

The MTP drafter (per-arch ``MtpHead`` — today
:class:`arbi_serve.spec_decode.mtp_head.Qwen3_5MtpHead`) runs a
``K``-step autoregressive chain to produce ``K`` speculative tokens
each verify cycle. The Python loop in
:meth:`arbi_serve.spec_decode.mtp.MtpDriver.draft` invokes the head
once per step; each invocation issues a fresh kernel-launch graph
plus host-side allocs (``embed`` lookup, ``_lm_head`` projection,
per-step token reduction). At ``K=3`` and decode ``B=1`` the per-step
launch storm dominates the chain wall-clock.

This module captures the entire chain (``K`` head invocations + ``K``
per-step token reductions) into a :class:`torch.cuda.CUDAGraph`
per ``(B, K)`` shape. Replay copies fresh inputs into persistent
buffers, replays, and returns the ``(K, B)`` int64 draft tokens.

Two capture flavors: GREEDY-ARGMAX and TRUE-STOCHASTIC
------------------------------------------------------
:class:`DrafterChainGraph` / :class:`DrafterChainGraphPool` capture
the greedy argmax chain. Persistent inputs: ``hidden_in``,
``last_token_in``, plus K per-step paged-attn metadata slices
(``slot_mappings[k]``, ``seq_lens[k]``, ``cu_seqlens_k[k]``,
``positions[k]``) and a shared ``cu_seqlens_q`` / ``block_table``.
Each step's feedback token is the argmax of its logits, baked into
the graph.

Stochastic requests draft greedily through this SAME captured chain
by default (the verify-pass rejection sampler corrects the one-hot
proposal losslessly). Under ``ARBI_TRUE_STOCHASTIC_DRAFT`` a
stochastic slate instead samples each step from the drafter's own
filtered distribution, and :func:`capture_drafter_chain` with
``stochastic=True`` captures THAT chain into a parallel pool
(``engine.cudagraph_pools["drafter_stoch"]``): each step draws its
carry token via the counter-keyed Philox Gumbel-max kernel, whose
only inputs are the persistent device step seed (``st_seed``,
``tl.load``-ed in-kernel — capture-safe by design) and the per-row
sampling-parameter tensors (``st_temps`` / ``st_top_k`` /
``st_top_p``, refreshed by ``copy_()`` at replay). The per-step
filtered distribution lands in the persistent ``q_chain (K, B, V)``
buffer for the verify-pass residual sampler. The counter-keyed kernel
needs no noise input, so the stochastic graph costs one extra fp32
``(K, B, V)`` q buffer per shape and is captured ONLY when
``ARBI_TRUE_STOCHASTIC_DRAFT`` is set at boot.

Real-attention paged-KV plumbing
--------------------------------
The bundled head (Qwen 3.5 / 3.6) runs REAL paged
self-attention over its own per-layer KV slab. Each chain step ``k``
consumes a fresh per-step :class:`AttnPagedKVMeta` whose
``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` / ``positions``
differ from step to step (request length grows by one each step;
each step's K/V scatters into the next slot in the MTP slab).

Capture threads K persistent per-step slices through the captured
graph: at warmup + capture the head reads these slices via K
:class:`AttnPagedKVMeta` instances pre-built once, each pointing at
a different slice. Replay :meth:`copy_` ``s`` the live driver's
per-step values into the persistent slices before invoking
``graph.replay()``; the captured kernel reads whatever the host
wrote (graph-capture-safe by construction).

The shared ``cu_seqlens_q`` is invariant across steps
(``[0, 1, 2, ..., B]`` — one query per row); ``block_table`` is also
shared (the K draft slots reside within already-allocated pages of
each request). Only the four per-step buffers vary across steps —
``B + (B+1)`` int32s + ``B`` int64s + ``B`` int32s = ``5B + 1`` ints
per step, sub-microsecond ``copy_()`` cost at decode shapes.

Real-attention coverage
-----------------------
The real-attention head is captured end-to-end: the per-step
slices documented above thread per-step ``attn_meta`` /
``slot_mapping`` through the captured pool's persistent buffers.

What stays out of the graph
---------------------------
  - Per-request slot allocation. The driver allocates K draft slots
    per request before the chain runs and writes the slot indices
    into the captured ``slot_mappings[k]`` slices via ``copy_()``.
    Allocation itself can't run inside the graph (it touches a host
    page-table data structure); replay re-assigns slots each step.

  - Per-row sampling PARAMETERS. The stochastic graph reads every
    knob from device tensors (:class:`DrafterSamplingTensors`
    encoding) refreshed at replay; the Python ``SamplingParams`` →
    tensor encoding runs on the host, outside the graph. ``min_p > 0``
    slates are NOT capturable (the routing layer sends them to the
    live chain, see :meth:`MtpDriver._draft_captured`).

Lifecycle
---------
Capture lives next to the decode-shape capture (one path per shape):

  1. Allocate persistent input + output buffers.
  2. Run ONE warmup chain (un-captured) to JIT every kernel
     (``embed`` / ``RMSNorm`` / ``Linear`` / sampler ops) against the
     persistent buffers.
  3. Capture the K-step chain under
     :func:`arbi_serve.distributed.graph_capture.graph_capture`
     (TP-aware — the head's row-parallel ``o_proj`` issues
     ``all_reduce`` collectives that the side stream pins).
  4. Persist the graph + buffers in the appropriate pool.

Sleep mode
----------
The stable-VA backend keeps both pools live across
``release_memory_occupation`` / ``resume_memory_occupation``;
persistent buffers' ``data_ptr``s survive the release/resume cycle so
the captured graphs stay callable.

Module layout
-------------
The shared capture scaffold (:class:`PersistentBufferHost`,
:class:`_ChainGraphPool`, :func:`_capture_chain_scaffold`, and the per-step
paged-KV slice helpers) lives in
:mod:`arbi_serve.runtime.capture.drafter_common`; the external draft-model
chain (:class:`ExternalDrafterChainGraph` / ``…Pool`` /
:func:`capture_external_drafter_chain`) lives in
:mod:`arbi_serve.runtime.capture.drafter_external`. Both are re-exported from
this module for import stability.

Owns persistent buffers + the captured K-step drafter graph.

Field shapes (sized for the captured ``(B, K)``):

  - ``hidden_in``         : ``dtype``  ``(B, hidden)``
  - ``last_token_in``     : int64      ``(B,)``
  - ``hidden_out_chain``  : ``dtype``  ``(K, B, hidden)`` — per-step
    post-norm hidden, ``hidden_out_chain[k]`` is the input to step
    ``k+1`` (and the seed for the next-step chain off this last
    slot, exposed via :attr:`last_hidden`).
  - ``last_token_chain``  : int64      ``(K, B)`` — per-step
    argmax token, returned to the caller.

Real-attention paged-KV slices. The captured head
reads a per-step :class:`AttnPagedKVMeta` whose buffers point at
the persistent slices below. ``cu_seqlens_q`` / ``block_table``
are shared across steps (invariant for the chain — one query per
row, slots within already-allocated pages). The four per-step
buffers vary by step:

  - ``slot_mappings``  : list[K] of int64 ``(B,)`` — different
    slot per step; the head's K/V at step k scatters into
    ``slot_mappings[k][i]`` for row ``i``.
  - ``seq_lens``       : list[K] of int32 ``(B,)`` — request
    length post-step-k write (``pre_len[i] + k + 1``).
  - ``cu_seqlens_k``   : list[K] of int32 ``(B+1,)`` — cumsum of
    ``seq_lens[k]`` with leading 0.
  - ``positions``      : list[K] of int32 ``(B,)`` — absolute
    position of each row's new token (``pre_len[i] + k``).

For legacy stub heads (no real attention) the per-step lists may
be empty — the head's forward is invoked with no attn_meta-bearing
kwargs and the captured graph only references hidden / token
buffers. The :func:`capture_drafter_chain` lifecycle picks the
right path by reading the head's structural marker (see
:func:`_can_capture_drafter_chain`).

The captured graph references every persistent buffer above by
``data_ptr``; replay :meth:`copy_` ``s`` fresh content into the
inputs and reads ``last_token_chain`` back. Output is
:meth:`clone` ``d`` before return so the next replay does not
silently corrupt a stale view.

LRU map of ``(B, K)`` -> :class:`DrafterChainGraph`.

Bound by ``max_shapes`` (engine ``cfg.cudagraph_max_drafter_shapes``,
default 8). Mirrors :class:`CapturedGraphPool` semantics: ``put``
inserts at MRU end, evicts from LRU end on overflow; ``get``
refreshes LRU position. Eviction drops the strong reference; the
underlying CUDAGraph + persistent buffers GC after their last
reference goes out of scope, ``torch.cuda.empty_cache()`` returns
segments to the driver.

Soft pre-flight: refuse capture in unsupported scenarios.

The drafter chain capture targets the per-arch bundled MTP head
(``Qwen3_5MtpHead`` today).

Scenarios refused (engine falls through to the live chain):

  - ``head`` is ``None`` (no bundled head) or is the spec_decode
    strategy's no-op fallback (``MtpDriver`` not attached).
  - ``n_draft < 1``.

Real-attention heads (Qwen3_5MtpHead, marked via
``_mtp_real_attention = True`` on the attention block class) ARE
accepted — :func:`capture_drafter_chain` allocates
K per-step :class:`AttnPagedKVMeta` slices in the captured graph's
persistent buffer set so each chain step's
``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` / ``positions``
is :meth:`copy_` ``ed`` from the live driver's per-step metas
before replay. See :class:`DrafterChainGraph` for the per-buffer
layout.

Returns ``(True, "")`` on success or ``(False, reason)`` for the
engine to log + skip the bucket.

True iff the head's attention block declares the
``_mtp_real_attention`` class-level marker (Qwen3_5MtpHead).
Stub heads / no-op heads return False — those capture against the
simpler :func:`capture_drafter_chain` path that doesn't thread
per-step paged-attn metadata.

Capture one drafter chain at the given ``(B, K)``.

Mirrors :func:`arbi_serve.runtime.capture.decode.capture_decode`'s
lifecycle (allocate persistent inputs -> warmup -> capture under
a stream-scoped graph context).

``stochastic=False`` (default): the captured chain runs the head's
GREEDY argmax path only — feedback uses ``argmax(logits)`` per
step. Routing a sampled draft through this graph would silently
collapse to argmax (correctness bug).

``stochastic=True`` (requires ``vocab_size``): the captured chain
SAMPLES each step's carry token from the drafter's own filtered
distribution via the head's ``sampling_tensors`` path — one
counter-keyed Philox Gumbel-max draw per step
(``draw_offset = OFFSET_DRAFTER + 17 * step``, byte-matching the
live chain's draws for the same seed), with the per-row
temperature/top_k/top_p knobs and the step seed read from
persistent device buffers ``copy_()``-refreshed at replay. The
per-step ``(B, V)`` filtered distribution lands in the persistent
``q_chain`` buffer (the verify residual sampler's real ``q``).
Replay returns ``(tokens, q)``; see :meth:`DrafterChainGraph.replay`.

``vocab_size`` is the DRAFTER's row width, which is not always the
verifier's: ``--draft-vocab-prefix`` gives the drafter an lm_head over
ids ``[0, N)`` and the head then emits ``(B, N)``. It sizes ``q_chain``,
the warmup ``top_k`` and the recorded mask's host values together, and
every downstream reader takes the width off ``q_chain`` — so this one
argument is the whole contract. Pass ``MtpDriver.draft_vocab_size``, not
``full_vocab_size``; a mismatch is a shape error at the first chain step.
Verify pads the narrow ``q`` back to full width, so nothing downstream
needs the verifier's number here.

Real-attention vs stub heads. When the head's attention block
declares ``_mtp_real_attention`` (Qwen3_5MtpHead), the
caller MUST supply ``state_view`` (the engine's MTP layer KV
slab view), ``attn_op`` (the per-layer paged attention op),
``rope_cache`` (the main model's RoPE cache), ``max_pages`` (the
captured ``block_table``'s pages dimension — the engine's
``ceil(max_context / page_size)``), ``page_size`` (the KV page
size those pages are counted in — the pair sizes the captured
``max_seq_len``, and neither number means anything alone), and
optionally ``sliding_window`` (None for full-attention layers, an
int for SWA layers). The capture allocates K per-step
:class:`AttnPagedKVMeta` slices and pre-fills them with synthetic
sane values so warmup + capture both see a self-consistent batch;
replay overwrites the slices with the live driver's per-step
metas. Legacy stub heads (no ``_mtp_real_attention``) use the
no-attn forward signature and the four kwargs above are ignored.

The capture region runs under
:func:`arbi_serve.distributed.graph_capture.graph_capture` so any
row-parallel ``all_reduce`` collectives the head's ``o_proj``
issues route onto the captured side stream. TP=1 is a no-op pin.

Refresh persistent inputs + replay + return ``(K, B)`` cloned tokens.

``hidden_in`` shape ``(B, H)``, ``last_token_in`` shape ``(B,)``.
Sizes must match the captured shape exactly — non-conforming
batches must take the live :meth:`MtpHead.forward` chain.

For real-attention captures (:attr:`real_attention` True) the
caller MUST pass ``attn_metas`` (length K) and ``block_table``;
the K per-step slices ``slot_mapping`` / ``seq_lens`` /
``cu_seqlens_k`` / ``positions`` are :meth:`copy_` ``ed``
from each ``attn_metas[k]`` into the persistent slice the
captured graph reads at step k. ``block_table`` is
:meth:`copy_` ``ed`` into the shared persistent slot. The
captured ``cu_seqlens_q`` is invariant (one query per row)
and is NOT refreshed.

For legacy (no-attn) captures the ``attn_metas`` /
``block_table`` kwargs are ignored — left for caller symmetry.

Output safety. ``last_token_chain`` is a persistent buffer the
next replay will overwrite. We :meth:`clone` before returning
so the caller never aliases a stale view across replays. The
clone is ``K * B * 8`` bytes — a few hundred bytes at typical
decode shapes.

``segmented`` replays the :attr:`step_graphs` one at a time instead
of the fused graph — the same kernels, K launches instead of one.
``gate_run`` is a per-slot depth policy
(:class:`~arbi_serve.spec_decode.depth_gate.ChainGateRunner`)
consulted in that seam: it is asked before each step and shown that
step's per-row margins after it, and the walk stops where it says to. The
return is then the DRAFTED depth, which may be shorter than the
captured ``K``; the caller reads the depth off the returned shape.
Both refuse a capture with no recorded segments rather than
silently running the full depth.

TRUE-STOCHASTIC captures (:attr:`stochastic` True) additionally
REQUIRE ``sampling_params`` (the slate's per-row
``SamplingParams``, encoded host-side via
:meth:`DrafterSamplingTensors.encode_params_host` into the
pinned staging twins and copied H2D non-blocking) and
``seed_buf`` (the engine's rank-agreed int64 ``(1,)`` step
seed). The return becomes ``(tokens (K, B), q (K, B, V) fp32)``
— both cloned off the persistent buffers. ``min_p > 0`` rows
are refused loud (the routing layer must send those slates to
the live chain).

Walk the per-step graphs, consulting ``gate_run`` between them.

Returns the depth actually drafted. The policy is asked BEFORE a
step's graph is launched, so a stop here is a head forward — and a
verify row — not spent, and it is shown the step's ``q`` after the
launch, which is the same order the live chain asks in.

Encode + ``copy_()`` the slate's sampling params / step seed into
the persistent stochastic input buffers the captured kernels read.

Host encoding (:meth:`DrafterSamplingTensors.encode_params_host`)
lands in the pinned staging twins; the device copies are
``non_blocking`` H2D ordered on the current stream before
``graph.replay()`` — no pageable-H2D stall, no host readback.

``min_p > 0`` slates are capturable ONLY when the fused draw is
active (``ARBI_DRAFTER_FUSED_DRAW`` on): the captured chain's fused
tail applies min_p in-kernel. On the unfused rollback the tail never
applied min_p, so those slates are refused (the routing layer sends
them to the eager chain).

Shared capture scaffold for the drafter chains.

The bundled MTP head (:mod:`arbi_serve.runtime.capture.drafter`) and the
external draft model (:mod:`arbi_serve.runtime.capture.drafter_external`)
are both autoregressive ``K``-step chains captured into a
:class:`torch.cuda.CUDAGraph` per ``(B, K)`` shape. Everything that is
identical across the two flavors lives here:

  - :class:`PersistentBufferHost` — name-routing persistent-buffer rebind
    surface + the real-attention replay refresh.
  - :class:`_ChainGraphPool` — the LRU ``(B, K)`` -> captured graph map.
  - :func:`_alloc_real_attention_slices` / :func:`_build_per_step_metas` —
    the per-step paged-KV slice allocation + :class:`AttnPagedKVMeta` build.
  - :func:`_make_chain_runner` — the default bundled-head chain runner.
  - :func:`_capture_chain_scaffold` — the shared capture lifecycle (persistent
    buffer allocation, pre-capture stable-VA migration, warmup, TP-aware
    graph capture, capture-memory accounting, sleep-pool writeback).

See :mod:`arbi_serve.runtime.capture.drafter` for the full narrative.

Re-derive this step's page metadata into the builder's persistent
buffers (captured per chain step), reading the per-step ``seq_lens`` /
``block_table`` off the drafter meta (no ``ScheduledBatch`` here).

Backend-agnostic: delegates to the builder's no-sync
``recompute_external_meta`` (TKV writes its ``TQBufferPool``;
tkv-bypass writes its CSR triplet). Both refresh the SAME persistent
buffers the pre-capture ``decorate_external_meta`` stashed on the meta,
so the captured kernel's baked ``data_ptr``s see fresh page tables on
every replay.

Shared name-routing + replay helpers for the captured chain graphs.

:class:`DrafterChainGraph` and :class:`ExternalDrafterChainGraph`
expose the same persistent-buffer rebind surface (sleep/wake
writeback) and the same real-attention replay refresh. The accessors route scalar buffer names to plain
attributes and the per-step ``slot_mapping[k]`` / ``seq_lens[k]`` /
``cu_seqlens_k[k]`` / ``positions[k]`` names into the corresponding
list attribute (which a bare ``setattr`` cannot reach). Mirrors
:class:`arbi_serve.runtime.capture.decode_graph.CapturedGraph`'s
recurrent-dict-aware variant.

``__slots__`` is empty so the ``@dataclass(slots=True)`` subclasses
keep their slot-only layout (no per-instance ``__dict__``); the
real persistent-buffer fields live on the subclasses.

LRU map of ``(B, K)`` -> captured chain graph.

Shared base for :class:`DrafterChainGraphPool` and
:class:`ExternalDrafterChainGraphPool`. Bound by ``max_shapes``
(engine ``cfg.cudagraph_max_drafter_shapes``); ``put`` inserts at the
MRU end and evicts from the LRU end on overflow, ``get`` refreshes LRU
position. Subclasses parameterize the eviction log label and whether
eviction calls ``torch.cuda.empty_cache()`` (see :meth:`put` for what
that does and does not return).

Allocate the real-attention persistent paged-KV slices for a chain.

Returns ``(cu_seqlens_q, block_table, slot_mappings, seq_lens,
cu_seqlens_k, positions)``. ``cu_seqlens_q`` is the invariant arange
``[0, 1, ..., B]`` (one query per row); ``block_table`` is shared
across steps. The four per-step lists are K independent allocations so
the captured graph reads from K distinct ``data_ptr``s at step k.
Synthetic capture-time values mirror :func:`capture_decode`'s
convention: ``seq_lens`` picked > page_size + K so TKV's bypass-safe
gate stays False at capture (page_size assumed >= 16).

Must be called inside the persistent-buffer pool context so the
tensors land under the same allocator as the rest of the chain.

Build the K per-step :class:`AttnPagedKVMeta` instances pointing at
the (already migrated) persistent slices.

Each meta is built ONCE pre-capture so the captured kernels reference
stable ``data_ptr``s on every field; replay :meth:`copy_` ``s`` fresh
content into the slices but the meta objects never change. When
``tq_builder`` is set, each meta is decorated with the persistent
``_tq_*`` pool buffers (their CONTENTS are recomputed per step inside
the captured region; decoration only attaches the references). No-op
for tkv-bypass / GDN backends.

Build the warmup/capture chain runner over the persistent buffers.

Returns a nullary callable that runs one K-step chain reading from /
writing to ``buffers`` (post-migration). Used by warmup AND inside the
capture region so both passes stay lockstep on tensor layout.

``steps`` restricts the runner to a subset of the chain's step indices
and ``via_buffers`` takes each step's carry off the persistent chain
buffers rather than off the previous step's returned tensors. The pair
is what makes a SEGMENT capturable on its own: a segment records one
step, so its carry has to arrive through a buffer whose address the
graph can bake, and the values are the same ones the fused chain would
have handed forward (``last_token_chain[k]`` / ``hidden_out_chain[k]``
are written by step ``k`` before step ``k + 1`` reads them, in the same
stream order).

``margins`` is the ``(K, B)`` persistent buffer the per-slot depth gate
reads between depths; step ``k``'s head fills row ``k`` off the logit
row its pick already forms. Passed to the SEGMENT runners and NOT to
the fused one: the fused graph replays every depth in one launch and
leaves no seam for the gate to read in, so recording the reduction
there would be a full-row cost on the shipped path buying nothing.

Each step builds the per-step real-attention kwargs (recompute the
TKV/tkv-bypass page metadata, then the ``attn_meta`` / ``positions`` /
``state_view`` / ``attn_op`` / ``rope_cache`` dict) and hands them to
``step_fn``, which runs the head forward (plus an optional sampler) and
returns ``(token, hidden, probs_or_None)``. The per-step outputs copy
into the persistent chain buffers in the fixed order ``last_token_chain``
-> ``hidden_out_chain`` -> ``q_chain`` (the last only when the buffer set
carries one). Step-to-step chaining passes the freshly returned tensors
directly (cudagraph-pool data_ptrs are stable across replays, and the
captured stream-ordering enforces step k's writes before step k+1 reads).

Shared capture lifecycle for the bundled-head + external drafter chains.

Owns the parts that are identical across every autoregressive K-step
drafter (bundled MTP head AND the external draft model):
persistent-buffer allocation under the buffers pool,
pre-capture stable-VA migration, the per-step :class:`AttnPagedKVMeta`
build + decoration, the two warmup passes, the TP-aware graph capture
(plus failure drain), capture-memory accounting, and the sleep-pool
writeback binding.

The per-flow difference is parameterized:

  - ``alloc_buffers`` returns the path-specific scalar persistent
    buffers in allocation order (insertion order drives the migration
    order). It is invoked inside the buffers-pool context.
  - ``make_step_fn`` (bundled-head paths) receives the post-migration
    buffer dict and returns the per-step callable the DEFAULT chain
    runner (:func:`_make_chain_runner`) invokes (the head forward +
    sampler wiring). Ignored when ``make_run_chain`` is supplied.
  - ``make_run_chain`` (external-drafter path) receives the
    post-migration buffer dict + the K per-step metas and returns the
    whole nullary chain runner — used when the per-step forward is a
    full ``model.forward(input_ids, ScheduledBatch, ...)`` rather than
    the head's ``forward(last_token_ids, prev_hidden, ...)`` contract.
  - ``build_graph`` assembles the path's graph dataclass from the
    captured buffers + slices.
  - ``log_capture`` emits the path's capture log line.
  - ``compile_target`` is the module the compile trampoline gate opens
    for during warmup + capture; defaults to ``head`` (the external
    path passes its drafter model instead).
  - ``segments`` additionally records ONE graph per chain step,
    handed to ``build_graph`` as ``step_graphs``. A fused K-step graph
    replays every depth it recorded with no host in the loop, so a
    policy that decides depth from the chain's own output cannot ride
    it; K single-step graphs replay the same kernels while leaving the
    host a seam between depths. Bundled-head paths only — the external
    drafter supplies its own whole-chain runner and has no per-step
    one to record.

Mirrors :func:`arbi_serve.runtime.capture.decode.capture_decode`'s
lifecycle. The capture region runs under
:func:`arbi_serve.distributed.graph_capture.graph_capture` so a head's
row-parallel ``o_proj`` ``all_reduce`` collectives route onto the side
stream (TP=1 is a no-op pin — the external drafter is TP=1 always).

Validate + :meth:`copy_` the live per-step paged-attn metas into
the persistent slices the captured graph reads.

No-op for non-real-attention captures. For real-attention captures
the K per-step ``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` /
``positions`` slices are refreshed from each ``attn_metas[k]`` and
the shared ``block_table`` is copied into the (zeroed) persistent
slot's prefix. The invariant ``cu_seqlens_q`` is NOT refreshed.

Register ``graph`` as MRU and evict the LRU entry on overflow.

After a real eviction, subclasses that set ``_EMPTY_CACHE_ON_EVICT``
call ``torch.cuda.empty_cache()``, which returns the evicted graph's
segments only when it was captured OUTSIDE a private ``MemPool``
(``empty_cache`` does not visit a private pool's block pools,
pytorch#145168 — a graph captured into ``capture.cudagraphs`` returns
its blocks to that pool for later captures to reuse). Gated to actual
eviction so steady-state ``put`` stays free.

Drop every captured chain + release its persistent buffers.

Teardown and pool-rebuild only. The sleep path does not call
this: every chain's buffers keep their ``data_ptr`` across
release / resume behind stable VAs, so the chains stay valid.

CUDAGraph capture + replay for the external draft-model chain.

The external drafter (:class:`arbi_serve.spec_decode.external_drafter.
ExternalModelDrafter`) is a SEPARATE small model (Qwen3-0.6B drafting
Qwen3-4B) that runs its OWN K-step autoregressive decode chain against
its OWN paged-KV pool + :class:`FlatPageTable`. It is just as much an
autoregressive K-step chain as the bundled MTP head, so it reuses the
SAME capture machinery from :mod:`arbi_serve.runtime.capture.drafter_common`
— :class:`PersistentBufferHost`, :class:`_ChainGraphPool`, and the whole
:func:`_capture_chain_scaffold` lifecycle (warmup / TP-aware capture /
migration / accounting / sleep writeback). It diverges from the bundled
head ONLY where it genuinely must:

  * Per-step forward contract. The head runs ONE MTP layer via
    ``head.forward(last_token_ids, prev_hidden, attn_meta, ...,
    return_hidden=True)`` against the MAIN pool's MTP slab; the external
    drafter runs the WHOLE drafter model via
    ``model.forward(input_ids, ScheduledBatch, drafter_pool, attn_ops)``
    against its OWN pool. So this path supplies a custom ``make_run_chain``
    instead of the head ``make_step_fn`` — the rest of the scaffold is
    shared verbatim.
  * No hidden-state feedback. The drafter feeds tokens only (the next
    step's input is the argmax of the prior step's logits), so the
    persistent buffer set is ``input_ids`` + ``last_token_chain`` — no
    ``hidden_in`` / ``hidden_out_chain``.
  * Greedy only. The captured chain bakes the per-step argmax feedback;
    stochastic external drafts take the eager live chain (TP=1, so a
    capture miss is slow, never a deadlock).

Owns persistent buffers + the captured K-step external-drafter graph.

The external drafter feeds ONE token per row per step (the seed at
step 0, then the prior step's argmax) and runs the full drafter
``model.forward`` against its own paged-KV pool. The persistent set is
therefore token-only — no hidden-state chaining:

  - ``input_ids``        : int64 ``(B,)`` — step-0 seed (replay copies
    the live ``last_token_ids`` in); the chain feeds back each step's
    argmax internally inside the captured graph.
  - ``last_token_chain`` : int64 ``(K, B)`` — per-step argmax token,
    returned to the caller (``(K, B)`` matching the eager chain).

Real-attention paged-KV slices (shared with the bundled-head capture,
see :class:`DrafterChainGraph` for the per-buffer contract): the
invariant ``cu_seqlens_q`` + shared ``block_table`` and the four
per-step ``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` /
``positions`` lists. Replay :meth:`copy_`\ s the live driver's
per-step metas into these via the inherited
:meth:`PersistentBufferHost._replay_refresh_attn`.

LRU map of ``(B, K)`` -> :class:`ExternalDrafterChainGraph`.

Sister of :class:`DrafterChainGraphPool` over the SAME ``(B, K)`` key
space + ``cudagraph_max_drafter_shapes`` LRU accounting. Eviction runs
the ``_EMPTY_CACHE_ON_EVICT`` pass exactly as the bundled-head greedy
pool does.

Capture one GREEDY external-drafter chain at the given ``(B, K)``.

Reuses the shared :func:`_capture_chain_scaffold` lifecycle (allocate
persistent inputs -> warmup -> TP-aware capture). The captured chain
runs K full ``model.forward`` decode steps over the drafter's own
paged-KV ``pool``; per step it recomputes the persistent CSR page
metadata (``metadata_builder.recompute_external_meta``) into the same
buffers the captured kernel baked, runs the forward, takes the argmax,
and feeds it back to the next step — all inside the captured region.

``metadata_builder`` is the drafter's PAGED_KV metadata builder
(tkv-bypass); it must expose ``decorate_external_meta`` /
``recompute_external_meta`` (the same no-host-sync CSR refresh the
bundled-head capture uses). ``max_pages`` sizes the captured
``block_table`` (the engine's ``ceil(max_context / page_size)``).

Refresh persistent inputs + replay + return ``(K, B)`` cloned tokens.

``input_ids`` shape ``(B,)`` — the step-0 seed (the prior step's
committed bonus). ``attn_metas`` (length K) + ``block_table`` carry
the live per-step paged-attn metadata the captured graph reads (the
same contract as :meth:`DrafterChainGraph.replay`): the K per-step
``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` / ``positions``
slices are copied from each ``attn_metas[k]`` and ``block_table`` is
copied into the shared persistent slot. The captured
``cu_seqlens_q`` is invariant (one query per row) and not refreshed.

``last_token_chain`` is a persistent buffer the next replay
overwrites — :meth:`clone` before returning so the caller never
aliases a stale view across replays.

Sleepable ``cudaGraphExec_t`` ownership — de-instantiate on sleep, re-
instantiate on wake, NO recapture.

Why this exists
---------------
An instantiated ``cudaGraphExec_t`` holds DRIVER-side device memory (kernel
parameter staging, ~tens of MiB for a large whole-forward graph) that lives
OUTSIDE every torch allocator pool — cuMem sleep (``named_pools.sleep_all``)
cannot release it, so it survives ``release_memory_occupation`` and keeps the
sleep floor up. The exec is pure derived state: it can be destroyed on sleep
and re-built from the captured ``cudaGraph_t`` on wake with
``cudaGraphInstantiate`` — byte-identical replay behaviour, no recapture
(the captured graph topology + every baked ``data_ptr`` are untouched; the
stable-VA contract keeps those pointers valid across the cycle).

torch 2.9 cannot do this natively: ``CUDAGraph`` only exposes
``instantiate()`` (destroy-then-recreate — never destroy-only) and
``reset()`` (destroys the captured graph too → would force recapture). So
:class:`SleepableCUDAGraph` captures with ``keep_graph=True``, NEVER lets
torch create its internal ``graph_exec_`` (we never call
``super().instantiate()`` / ``super().replay()``), and owns the exec handle
itself via cuda-python's runtime bindings:

  * ``capture_end``       → ``cudaGraphInstantiate`` (eager, same memory
    timing as torch's default path so the post-capture VRAM gate still
    measures the true serving footprint)
  * ``replay``            → ``cudaGraphLaunch`` on the current torch stream
  * ``deinstantiate_exec`` (sleep) → ``cudaGraphExecDestroy``
  * ``reinstantiate_exec`` (wake)  → ``cudaGraphInstantiate`` again

RNG contract (load-bearing)
---------------------------
``torch.CUDAGraph.replay()`` does one thing our ``cudaGraphLaunch`` skips:
the captured-generator ``replay_prologue`` that refreshes philox seed/offset
state for RNG ops captured INSIDE the graph. arbi-serve's capture policy is
that NO captured region contains a torch RNG op — stochastic sampling
noise either enters through HOST-FILLED persistent buffers or is generated
by counter-based (Philox-in-kernel) draws keyed on a persistent seed buffer,
neither of which touches a captured torch generator. Under that policy the
prologue is a semantic no-op
and skipping it is exact. A captured RNG op would misbehave with or without
this class (frozen noise across replays); keep RNG out of capture bodies.

Threading: capture/replay/sleep/wake all run on the engine's forward thread
or under the engine critical section — no concurrent access to one graph.

Process-wide meter for MEASURED ``cudaGraphInstantiate`` driver bytes.

The exec's driver allocation lands on the default device heap: outside the
torch caching allocator, outside every cuMem pool, and therefore outside the
map-time pool caps. ``torch.cuda.mem_get_info`` is the ONLY counter that
sees it — which is why the post-capture grow (which sizes against
``mem_get_info``) has always paid this cost whether or not the budget booked
it, and why the budget's job is to predict it, not to discover it.

Every consumer of that prediction used to read a CONSTANT. The constant was
assumed ~3,300 nodes/graph. Both inputs of that model are refuted on this
driver — real graphs measure 417 nodes on average (6,258 over 15 graphs via
``cudaGraphGetNodes``), and a direct sweep at 400 / 2,000 / 8,000 nodes x 8
not a memory lever at all. A model whose inputs are both refuted should not
be re-fitted; it should be replaced by the measurement.

So: bracket the real call, in the real boot, at the one chokepoint that
makes it (:meth:`SleepableCUDAGraph.instantiate_exec`), and persist the
result under the same budget-cache identity ``graph_pool_bytes`` already
uses — cold boot on a seed, every later boot on the measurement.

Aggregate-then-divide is deliberate. A single exec's ``mem_get_info`` delta
can read 0 (the driver sub-allocates out of a suballocator arena it grew on
an earlier instantiate), so a per-graph number is only meaningful across the
budget needs.

Samples are bucketed BY FAMILY (``decode`` / ``drafter`` / ``prefill`` /
``mixed`` / ``piecewise``) via :meth:`family`, because that is the
granularity the reserve is predicted at — and because a per-family
measurement is the only thing that can settle whether a partial-forward
graph (a drafter chain bakes K steps of a ONE-layer head) really is cheaper
than a whole-model one. It replaces estimating that ratio from a layer
count, which is a model, with reading it, which is not.

Thread-safety: capture is single-threaded at boot, and the counters are
plain ints under the GIL. No lock — a lock here would be ceremony.

``torch.cuda.CUDAGraph`` whose instantiated exec can sleep.

Drop-in for every ``torch.cuda.CUDAGraph()`` construction in the
capture subsystem: ``torch.cuda.graph(...)`` (capture_begin /
capture_end), ``replay()``, ``reset()``, ``enable_debug_mode()`` /
``debug_dump()`` all keep working. The ONLY behavioural difference is
that the executable graph is owned here (``_exec_handle``) instead of
inside libtorch, which adds :meth:`deinstantiate_exec` /
:meth:`reinstantiate_exec` for the sleep path.

Attribute every instantiate inside the block to capture family
``name``. Nested-safe (restores the enclosing family on exit).

REGISTERS the family on ENTRY, before anything can be recorded. Only
``record`` used to create the key, so a family that captured graphs the
meter never saw was ABSENT from ``counts()`` — byte-identical to a
family that was never entered. That is a green meaning "did not run":
the sweep looks fully metered whether it was or not. Registering here
makes the two states distinguishable, which is what
:attr:`families_without_samples` reports on.

Costs nothing downstream: ``per_family_bytes_per_graph`` and
``resolve_instantiate_bytes_per_graph`` both filter on a positive graph
COUNT, so a registered-but-empty family can never be persisted as a
0 bytes/graph rate.

File the driver-resident growth ``family``'s own capture bracket saw.

``delta_bytes`` is the change in DRIVER-resident bytes across exactly
the span this meter's :meth:`family` scope covered, read from the boot
residency meter's phase samples — the same ``mem_get_info``-derived
quantity ``driver.modules_loaded`` is the residual of. ``None`` when
that bracket could not be read (no CUDA, a probe that raised): an
unread bracket must stay visibly unread rather than read as zero
growth, because zero growth is the load-bearing evidence below.

This is what separates a family that captured NOTHING from one that
captured through a path this meter cannot see. Both meter zero
instantiates; only the second one grew the driver.

Families ENTERED that metered no instantiate at all.

This is a statement about THIS METER and nothing else. It does not by
itself say the family captured anything: a sweep gated off by config,
or one whose duck-typed hook found nothing to sweep, enters its family
and captures zero graphs, which meters zero exactly like a capture this
meter cannot see. :attr:`blind_spot_families` and
:attr:`dormant_families` are the two halves, split on evidence.

``(family, driver bytes)`` for every family that metered zero
instantiates while its OWN capture bracket grew driver-resident memory.

Growth with nothing metered is the positive evidence that a capture
path reached the driver without passing :meth:`instantiate_exec` — a
raw ``torch.cuda.CUDAGraph()``. Those bytes are real and, being
unmetered, fall into ``driver.modules_loaded`` (the residual of the
same brackets) instead of ``driver.cudagraph_exec``, so the per-graph
rate the meter publishes does not cover them.

The test is the mechanism itself — driver bytes moved — not a size
cutoff, so the reported figure is the whole of the evidence and a
reader can see for themselves whether it is graph-sized.

Families ENTERED that metered zero instantiates AND grew no driver
memory in their own bracket — they captured nothing.

A sweep can be dormant for reasons that are entirely correct (its
config flag is off, its duck-typed binding is absent on this model),
and reporting that as a metering defect is a red that means "did not
run". Distinguishing the two is why :meth:`note_driver_bracket` exists.

Mean measured driver bytes per instantiated exec.

``family=None`` averages over every family. Returns ``None`` when
nothing was instantiated (never a fabricated 0 — an unmeasured term must
be visibly unmeasured, so the caller falls back to the labelled seed
rather than silently booking zero).

Build the ``cudaGraphExec_t`` from the kept ``cudaGraph_t``.

Idempotent. Requires a completed capture (``capture_end``).

Brackets the driver call with ``mem_get_info`` and books the delta into
:data:`instantiate_meter` — the MEASUREMENT that retires the guessed
per-graph constant (see :class:`InstantiateMeter`). Two
``cudaMemGetInfo`` calls per instantiate: boot-only work, microseconds,
and never on a replay path.

``meter=False`` for the WAKE path (:meth:`reinstantiate_exec`): waking a
parked model re-instantiates execs it already paid for, so metering it
would inflate the per-graph mean with duplicate samples of the same
graphs. Only the boot capture sweep feeds the budget.

Wake: re-build the exec from the kept captured graph. Returns
True when a new exec was created (False ⇒ already instantiated).

NOT metered — see :meth:`instantiate_exec`. A wake re-instantiates
graphs whose driver cost the boot sweep already measured and persisted;
re-booking them would double the sample count for the same population.

Launch the owned exec on the current torch CUDA stream.

Exactly what ``CUDAGraph::replay`` does minus the captured-
generator prologue — a no-op under the no-RNG-in-capture policy
(see module docstring). Raises when called asleep (exec destroyed)
— the engine's ``_memory_released`` gate must prevent that.

Whole-forward MIXED decode+prefill cudagraph capture (``ARBI_MIXED_CAPTURE``).

Captures ONE graph for the composition bucket ``(num_decode_rows,
prefill_tokens)`` — the canonical mixed-step shape at concurrency where
the scheduler co-admits one prefill chunk alongside N steady-state
decode rows. Replaces BOTH the ``split_mixed`` two-forward path and
the fully-eager per-layer dispatch with a single whole-forward replay.

Canonical row layout (fixed at capture; the runtime staging in
:mod:`arbi_serve.runtime.mixed_capture` permutes the live slate into it):

  * rows ``0 .. D-1``  — decode rows, exactly ONE query token each;
  * row  ``D`` (last)  — the prefill row, ``prefill_tokens`` query tokens.

so ``cu_seqlens_q = [0, 1, 2, ..., D, D + prefill_tokens]`` and the flat
token count is ``N_flat = D + prefill_tokens`` (≠ ``B * S`` — the
captured graph records ``flat_tokens`` so :meth:`CapturedGraph.replay`
validates against the ragged layout).

Why no new attention / GDN kernel is needed (verified):

  * The Turbo prefill paged prefill (``TkvBypassAttnOp._prefill``,
    ``backends/tkv_bypass_backend.py``) is a varlen call over
    ``cu_seqlens_q`` + per-row K lengths; its causal mask offsets the
    diagonal by ``seqlen_k - seqlen_q`` per row, so a ``q_len=1`` decode
    row and a ``q_len=N`` prefill row coexist in ONE launch.
  * The GDN prefill path (``models/_gdn_fla.py::_forward_prefill_fla``)
    is already ragged: the varlen conv kernel
    (``kernels/gdn_prefill_conv_varlen.py``) builds its chunk map
    ON DEVICE from ``cu_seqlens`` (capture-safe, composition-agnostic
    within the fixed bucket), and the chunk delta-rule kernel
    (``kernels/fla_vendored/chunk.py::chunk_gated_delta_rule_inference_fwd``)
    consumes arbitrary per-row lengths — a decode row is a length-1
    varlen segment whose recurrent state advances by one token, with
    per-sequence ``initial_state`` gathered through the SAME persistent
    ``state_indices`` buffers the decode/prefill captures already pin.

Padding-replay contract (mirrors :func:`capture_prefill`'s): the live
step may carry ``d_live <= D`` decode rows and a prefill chunk of
``real_N <= prefill_tokens`` tokens. The runtime staging pads decode
rows with benign scratch rows (page-0 null slot, ``seq_lens = 1``,
borrowed free recurrent slab rows — the :mod:`decode_pad` contract) and
pads the prefill tail with token 0 → the null slot. ``cu_seqlens_q`` is
refreshed per replay to ``[0, 1, ..., D, D + real_N]`` so the captured
last-token gather + Turbo prefill per-row Q range + the GDN varlen kernels bound
their work by the LIVE row lengths; positions past ``real_N`` are
processed for shape but masked for content (the same mechanism the
shipped single-row prefill padding-replay relies on).

Declare the fused forward's decode/prefill row split to the EXL3 leg.

A no-op context on any other backend: EXL3 is the one whose GEMM leg and
cuBLAS kernel are selected by ROW COUNT, so it is the one for which
fusing two row classes into a single call changes a decode row's output.
A bf16 / AWQ / FP8 linear computes a row the same way whatever else
travels with it, and the optional EXL3 extra may not be installed at all.

Capture one whole-forward MIXED graph at ``(num_decode_rows,
prefill_tokens)``.

Mirrors :func:`capture_prefill` with ``B = num_decode_rows + 1``
rows in the canonical layout (decode rows first, one token each;
the prefill row LAST with ``prefill_tokens`` tokens). Registered
in the shared pool under ``(B, S=prefill_tokens, lora_bucket=0,
is_prefill=True, kv_pages_bucket=0)`` — ``B >= 2`` keeps it
disjoint from the single-row prefill captures (``B == 1``) so no
lookup can confuse the two.

``context_len`` selects the capture FAMILY, exactly as it does for
:func:`~arbi_serve.runtime.capture.prefill.capture_prefill`, and the
graph is registered under ``prefill_context=bool(context_len)`` so a
lookup can never hand one family's graph to the other:

  * ``0`` — the prefill row is a FIRST chunk (``q_len == seq_len``).
    TKV's per-row route gate (``compute_row_split``) classes it
    ROW_BYPASS and the graph bakes the bypass attend, which reuses
    ``cu_seqlens_q`` as ``cu_seqlens_k`` and therefore attends over
    the chunk's own tokens only.
  * ``> 0`` — the prefill row is a CONTINUATION chunk (``q_len <
    seq_len``). The gate classes it paged, and the paged branch takes
    each row's K extent from the ``seq_lens`` buffer the replay
    refreshes, so ONE graph serves every context length.

Only the prefill row moves between the families. The decode rows are
``q_len = 1``, ``seq_len = 2`` in both — a decode row is already a
``q_len < seq_len`` row, which is why the batch-global
``compute_bypass_safe`` is False for a mixed batch either way and the
per-row split is the gate that actually decides.

The synthetic capture batch marks ``is_prefill=True`` so the
metadata builders take their prefill paths and the model routes
every layer through the varlen prefill kernels (Turbo prefill paged varlen
attention; GDN varlen conv + chunk delta-rule) — the composition-
agnostic kernels that serve BOTH row flavors in one launch.

Pre-flight is the SAME model-level gate as the prefill capture
(:func:`_can_capture_prefill`) — the mixed graph runs the same
kernel set over a ragged ``cu_seqlens_q``; any model the prefill
capture refuses is refused here for the same reasons.

CUDAGraph capture + replay for

Why this is its own (small) module instead of a call into
:func:`arbi_serve.runtime.capture.decode.capture_decode`
--------------------------------------------------------
``capture_decode`` is the engine's boot-time, shared-scheduler-pool decode
capture — built for a ladder of ``(B, S, lora_bucket, kv_pages_bucket)``
shapes captured ONCE at boot and replayed by many concurrent requests. Two
of its assumptions don't hold for
:class:`~arbi_serve.runtime.nemotron_voicechat_stt_step.EngineSttStep`:

  1. **Single vs dual logits head.** ``capture_decode``'s captured region
     calls ``collapse_dual_head_logits(model, model.forward(...))`` — by
     design (see that function's own docstring) it keeps only the FIRST
     named head and drops the rest, because its callers are the generic
     scheduler/sampler dispatch, which only knows how to consume one
     logits tensor. ``NemotronVoiceChatModel.forward`` returns
     ``(text_logits, function_logits)`` and this turn loop genuinely
     needs BOTH every step (the function channel feeds the next region-2
     fusion step, see ``EngineSttStep.__call__``) — capturing through the
     collapsing helper would silently stop producing a function head at
     all. ``NemotronVoiceChatBackboneModel.forward``'s own docstring
     already flags this: the ``batch.logits_out`` destination-buffer
     optimization ``capture_decode`` also relies on "does NOT route
     through ``LayerStackModelMixin._lm_head`` ... left as a follow-up".
  2. **Shared engine pool + scheduler-built metadata builders.**
     ``capture_decode`` takes ``metadata_builders`` built once at engine
     boot against the SHARED ``MultiStatePool``/scheduler and threads a
     synthetic batch through them before capture. This loop owns a
     PRIVATE, turn-scoped pool (never the shared engine pool — see
     ``EngineSttStep``'s own docstring on why) and is not wired into the
     scheduler at all; there is exactly one live shape for its entire
     life (B=1, S=1 — the frame-lockstep bucketing fix already pins
     ``block_table`` width).

Given a single fixed shape, no LoRA, no MTP, no TP>1 (this backbone is
TP=1-only, see the EXL3 conversion doc), a dedicated capture path scoped
to exactly this call pattern is simpler and safer than bending the
general one to fit — the DRY win of reuse would cost more in caveats
than it saves. What IS reused (the actual repeated machinery, not
reinvented): :func:`arbi_serve.distributed.graph_capture.graph_capture`
for TP-correct (here TP=1, no-op) collective pinning,
:class:`~arbi_serve.runtime.capture.graph_exec.SleepableCUDAGraph` for the
graph object, :func:`arbi_serve.compile.capture_bridge.compile_capture_ctx`
to route the captured forward through the model's already-JIT'd compiled
callable, :func:`arbi_serve.runtime.capture._common._force_end_stream_capture`
for the same failed-capture stream-drain recovery every other capture
path uses, and each per-kind metadata builder's own
``cudagraph_capture_step`` hook (the tkv-bypass PAGED_KV builder's).

What gets captured
-------------------
The steady-state single-token decode-shaped forward — region 2 (an
audio-frame-fused step, ``inputs_embeds`` given) and region 3 (an
ordinary one-token decode) are the SAME captured shape: both feed
``inputs_embeds`` (region 3's plain-token embedding is looked up via
``embed_tokens`` OUTSIDE the graph, once per step, purely so the graph
body never branches on which region it's in — see
:meth:`SttDecodeGraph.replay`). One persistent ``(1, hidden_size)``
embedding buffer, one persistent ``(1, max_pages)`` block table (already
fixed-width per ``EngineSttStep._extend``'s own bucketing fix), growing
``seq_lens``/``slot_mapping``/``positions`` refreshed every replay. The
Mamba row (``state_indices``) and ``has_initial_state`` are TRUE
constants for a turn's whole life (one request, state already primed by
the initial prompt prefill before the first captured step) and are
baked once, never refreshed.

Capture records, replay performs
--------------------------------
:func:`capture_stt_decode_step` never advances the turn: it snapshots the
Mamba row, runs its warmup forwards, records the graph, restores the row,
and returns. The caller performs the step by replaying — including the
very first time, on the step it captured on. See that function's docstring
for why the recurrent state makes this mandatory rather than tidy.

What stays out
---------------
The initial system-prompt prefill and any ``inject_tokens`` tool-response
splice (both variously-sized, one-off shapes — see ``_extend``'s own
``allow_compile`` docstring) run eager, exactly as before. This mirrors
the documented "capture the steady-state case, refuse the irregular
ones" pattern from ``docs/omni_serving.md``'s Known Limits (media-
carrying prefills there; one-off-shaped prefills here).

Capture one ``SttDecodeGraph`` at the caller's CURRENT real turn
state (not a synthetic batch — this loop captures exactly once per
turn instance, so there is no shared-pool ladder to keep generic; the
values fed here are the shape AND the addresses the first replay will
use).

Capture does NOT perform the step
---------------------------------
This function is deliberately a no-op on the turn's recurrent state:
it returns with ``pool``'s Mamba row holding exactly the bytes it
held on entry, and the caller MUST call
:meth:`SttDecodeGraph.replay` with the same ``position``/``seq_len``/
``page_ids``/``inputs_embeds`` to actually advance the step (which is
what :meth:`EngineSttStep._decode_forward` does — its capture branch
falls straight into the ordinary replay branch).

That contract is not stylistic. Two REAL warmup forwards run below
(they finish the compiled callable's lazy JIT before the capture
window opens, where a tracing-time ``cudaMalloc`` would be fatal),
and a Mamba-2 recurrent update is ``state = A*state + B*x`` — running
it twice COMPOUNDS rather than overwrites, so without an undo the
capture step would leave the turn two advances ahead of the one token
it represents, permanently, for the rest of the turn. The row is
therefore snapshotted before the warmups and restored after the
capture window closes, via the pool's own savepoint API. The single
real advance is then performed by the caller's first replay, off the
restored (correct) state.

The KV slot for ``position`` is the one thing the warmups do dirty
and this function does not restore — deliberately: unlike the
recurrent update a KV write is a plain store to a slot nothing has
read yet, and the caller's first replay overwrites it with the value
computed off the restored state. Restoring it would be dead work.

``attn_ops``/``metadata_builders`` must be the SAME persistent
objects the caller's later eager (uncaptured) calls use for
anything that still runs off-graph (the prompt prefill, tool-
response injection) — the metadata builders in particular own
per-kind scratch (e.g. the tkv-bypass CSR buffers) that the
captured region's ``cudagraph_capture_step`` hook writes into.

``recurrent_req_id`` is the pool-level request id the caller passed
to ``alloc_recurrent_state`` (``mamba_row`` is that id's resolved
slab row) — the savepoint API is keyed by request id, not by row.

Refresh the per-step buffers, replay, return (text, function)
logits clones (mirrors :meth:`CapturedGraph.replay`'s clone
contract — the persistent output buffers are overwritten by the
NEXT replay, so the caller must not hold a view across one).

CUDAGraph capture + replay for the NemotronVoiceChat perception tower's
(Fast

=== Scope decision: bucketed batch encode, not streaming chunks ===

NVIDIA's own production Triton backend (see this branch's extracted
reference, ``perception_cudagraph.py``'s ``PerceptionCudaGraph`` /
``PerceptionCudaGraphPool``) CUDA-graph-captures this tower's encoder —
but their capture target is a per-CHUNK streaming call: fixed
``max_spec_frames`` (default 561, ~5s), a ladder of BATCH-SIZE buckets
(1/4/8/16/32) for concurrent streaming sessions, decoupling the
preprocessor (variable-length, runs eager) from the encoder (fixed-shape,
captured).

Our architecture does not call the tower that way today.
``arbi_serve.runtime.nemotron_voicechat_stt_step.EngineSttStep.__init__``
calls ``binding.encode(mm_feats, device)`` — which resolves to
``NemotronVoiceChatPerceptionTower.encode`` — exactly ONCE per turn, over
the turn's WHOLE buffered utterance (see that module's docstring's
"Frame-lockstep audio-in stepping" section: audio-in is a genuine
per-frame-fused autoregressive backbone loop, not a placeholder-splice —
but the perception TOWER itself still runs the entire utterance's mel
spectrogram through in one batch/offline forward, exactly like
``att_context_style="chunked_limited"``'s own docstring in
``models/audio/nemotron_voicechat_perception.py`` describes: "the
reference has no separate 'chunk a long utterance' path for offline
inference... expressed as a banded causal attention mask inside a single
forward, not a manual split into multiple forward calls"). There is no
incremental-cache-carrying entry point on this tower, and nothing in this
codebase calls one — building streaming infrastructure (per-chunk causal-
conv state cache + attention K/V cache across calls) to literally mirror
NVIDIA's shape would be substantial new, untested code in service of a
call pattern this integration does not use.

``PerceptionStreamState``) for the duplex lane's live WS audio-in path.
This module is still correct and still what the TURN-BASED path wants —
``EngineSttStep`` continues to call ``binding.encode`` once per turn over
the whole buffered utterance, and that is the call this captures. What is
now out of date is only the "nothing in this codebase calls one" premise:
duplex connections do, via ``encode_audio_chunk``.

So this module captures what we ACTUALLY use: the batch/offline
``encoder(mels, lengths)`` + ``proj`` forward, bucketed by mel-frame
length (the tower's own docstring already establishes that zero-padding
a batch to a fixed length is numerically safe here — "the padded tail of
a shorter item... never contaminates... this module still threads exact
per-layer masks", same invariant NVIDIA's own fixed-``max_spec_frames``
padding relies on). Batch is pinned to 1 — the single real call site
always encodes exactly one turn's audio as one row (an ``extract_audios``
result CAN in principle carry more than one audio content part per turn,
but real-time voicechat turns are one continuous utterance; that case
falls back to eager, same as an over-length utterance below).

=== Bucket ladder ===

Four buckets over PRE-subsampling mel-frame count — the tower's
``_ConformerEncoder.forward(mels, lengths)`` input T axis — at
``HOP_LENGTH=160`` samples / 16 kHz = 100 mel-frames/sec (see
``multimodal/preprocess/nemotron_voicechat.py``'s ``mel_length``):
200 / 500 / 1000 / 2000 frames ~= 2s / 5s / 10s / 20s of real-time-
voicechat turn audio, spanning short utterances through a long
monologue-style turn. A turn whose audio exceeds the largest bucket
falls back to the eager ``tower.encode`` path — never raises; mirrors
:mod:`arbi_serve.runtime.capture.nemotron_voicechat_decode`'s own
"capture the steady-state case, refuse the irregular ones" convention
(that module's docstring, citing ``docs/omni_serving.md``'s Known
Limits).

What's reused (not reinvented): :func:`arbi_serve.distributed
.graph_capture.graph_capture` for TP-correct (here TP=1, no-op)
collective pinning, :class:`~arbi_serve.runtime.capture.graph_exec
.SleepableCUDAGraph` for the graph object, and
:func:`arbi_serve.runtime.capture._common._force_end_stream_capture` for
the same failed-capture stream-drain recovery every other capture path
uses. Unlike :mod:`.nemotron_voicechat_decode`, this module does NOT
route through :func:`arbi_serve.compile.capture_bridge.compile_capture_ctx`
— the perception tower has no ``torch.compile`` integration anywhere in
this codebase (only the backbone does), so there is no compiled
callable's trampoline to open during capture.

Capture one :class:`PerceptionEncodeGraph` at ``max_mel_frames``.

``tower`` must already carry its real, loaded weights on ``device``
— this captures the tower's ACTUAL forward, not a synthetic
throwaway shape (matching :func:`arbi_serve.runtime.capture
.nemotron_voicechat_decode.capture_stt_decode_step`'s own convention
of capturing against real state rather than a warmup-only stand-in).

Lazily-built bucket ladder of :class:`PerceptionEncodeGraph`,
wrapping one :class:`NemotronVoiceChatPerceptionTower`.

Exposes :meth:`encode` matching ``MediaBinding.encode``'s contract
1:1 (``(feats, device) -> (n_tokens, hidden)``), so
``NemotronVoiceChatModel``'s ``mm_bindings["audio"]`` wiring can swap
the plain ``tower.encode`` closure for this instance's bound method
without any other change — see that class's ``__init__``.

Falls back to eager ``tower.encode`` (never raises on shape alone)
when: capture is disabled (``ARBI_STT_PERCEPTION_CAPTURE_OFF=1``),
the device isn't CUDA, the request carries more than one audio item,
or the utterance's mel-frame count exceeds the largest bucket — see
module docstring's "Bucket ladder" section.

Zero-pad ``mels_real`` (``(1, feat_in, T<=max_mel_frames)``)
into the static input buffer, refresh the real valid length,
replay, and return ``(cloned output sliced to the real output
length, that length)``.

Mirrors :meth:`~arbi_serve.runtime.capture.nemotron_voicechat_decode
.SttDecodeGraph.replay`'s clone contract — the persistent
``encoded_out``/``enc_lens_out`` buffers are overwritten by the
NEXT replay, so the caller must not hold a view across one.

The captured graph for ``bucket``, capturing it if the ladder is open.

``None`` once :meth:`seal` has run and the bucket was never captured —
the caller then encodes eagerly rather than capturing into a frozen
allocator layout.

Capture the whole bucket ladder. Returns the number captured now.

Called by the boot capture sweep so every bucket a request can reach is
recorded inside the sweep's named capture pool, instead of being minted
lazily on the first request that reaches it. Idempotent — a bucket
already captured is skipped. No-op when capture is disabled, the device
is not CUDA, or the ladder is sealed.

Record + warn ONCE per bucket that this shape encodes eagerly.

A captured replay and an eager encode differ in latency for the life of
the process, so a permanent fall to eager is a named warning and a
counter rather than a debug line.

Close the ladder: no further bucket may be captured.

The Phase-2 freeze caps every cuMem pool at its live mapped bytes, so a
capture after it has no budgeted home. Sealing turns a ladder miss into
the eager encode the class already falls back to for out-of-ladder
shapes, whose transient the boot profile priced into the media reserve.

Whole-forward prefill cudagraph capture.

Mirrors :func:`arbi_serve.runtime.capture.decode.capture_decode` for
the single-sequence prefill shape ``(B=1, S=num_tokens)``: one captured
graph per ``num_tokens`` bucket replacing the per-layer piecewise
capture's hundreds of layer graphs. The pre-flight gate
(:func:`~arbi_serve.runtime.capture.preflight._can_capture_prefill`)
and the captured-graph container
(:class:`~arbi_serve.runtime.capture.decode_graph.CapturedGraph`) live in
sibling modules.

Guard the capture's EXL3 prefill accumulator, when there is one to guard.

Imported lazily and answered with a no-op when the EXL3 path is not in this
build: a bf16 or AWQ deployment must not import the EXL3 module to capture
a graph, and the guard has nothing to say about a model with no EXL3 linear
in it.

Capture one whole-forward prefill graph at ``num_tokens``.

Mirrors :func:`capture_decode` for the prefill shape ``(B=1,
S=num_tokens)``. One captured graph per ``num_tokens`` bucket
replaces the per-layer piecewise capture's hundreds of layer
graphs.

Differences vs :func:`capture_decode`:

  - ``B`` is hard-coded to 1: prefill is single-sequence by
    scheduler convention (one prefill per step, no batched
    prefill in the current run loop). Multi-sequence prefill
    falls through to live forward.
  - ``seq_len`` is the full ``num_tokens`` (no ``[1, 8]`` BLOCK_M
    gate). The prefill kernel template autotunes per-shape and
    does not require a fixed BLOCK_M; capture bakes the
    tile-shape autotune picks for this ``num_tokens``.
  - ``context_len`` selects the capture FAMILY, and with it the
    attention route the graph freezes:

      * ``context_len == 0`` (first chunk) — ``seq_len == q_len``,
        so TKV's host route gate resolves ``bypass_safe = True``
        and the graph bakes the FA-varlen BYPASS branch, which
        attends over the forward's own contiguous K/V.
      * ``context_len > 0`` (continuation chunk) — ``seq_len =
        context_len + num_tokens > q_len``, so the gate resolves
        ``bypass_safe = False`` and the graph bakes the PAGED
        branch. That branch takes each row's K extent from the
        ``seq_lens`` buffer (``mSeqUsedK``, a per-CTA device read),
        so ONE captured graph serves every context length the
        replay refreshes it to — the same property
        :func:`capture_decode` already relies on.

    A host predicate inside a captured region is frozen at capture.
    That is only a defect when the wrong branch was frozen: the
    route is constant WITHIN each family, so capturing one graph per
    family bakes exactly the branch that family needs. The live
    lookup keys on ``prefill_context`` and can never hand a
    continuation chunk the first-chunk graph.
  - No MTP synthesis: prefill never carries an mtp_meta; the K=0
    cold path runs decode-via-verify on the next step.
  - ``mtp_fill`` decides whether the captured forward RUNS the
    bundled MTP drafter-KV fill
    (:meth:`~arbi_serve.models.qwen3_5.Qwen3_5Model._mtp_kv_fill`:
    an extra MTP decoder layer + paged KV write to the MTP slab +
    vocab-parallel embed). The model gates the fill on
    ``batch.mtp_fill_enabled is not False``, so a synthetic batch
    that leaves the field at its ``None`` default silently bakes
    the fill into EVERY captured prefill graph — dead compute on
    every no-MTP replay. The flag is therefore
    REQUIRED: the caller passes the boot-constant decision
    (:func:`arbi_serve.engine.cudagraph_admin._prefill_capture_mtp_fill`)
    and the synthetic batch pins ``mtp_fill_enabled`` to it
    explicitly. ``mtp_fill=True`` additionally captures the forward
    with ``return_hidden_state=True`` and retains the LAST-SLOT
    ``(B, H)`` hidden on :attr:`CapturedGraph.hidden_out`: the
    MTP-opted FINAL prefill chunk (the seed forward,
    ``run_step._run_mtp_seed_forward``) replays these graphs and
    reads that row for the seed draft — a fill variant without
    ``hidden_out`` is refused by the lookup's ``require_hidden``
    gate and can only serve mid-prefill chunks.
    ``mtp_fill=True`` on a model that cannot run the
    fill (no bundled ``mtp_head``, or a DFlash tap owns the MTP
    slab) is a caller bug and fails loud. The resulting state is
    recorded on :attr:`CapturedGraph.mtp_fill`;
    :func:`~arbi_serve.runtime.captured_lookup.lookup_captured_graph_for_prefill`
    refuses replay when the live batch's effective fill state
    differs (falls back to eager, which honours the flag).

Pre-flight gate is :func:`_can_capture_prefill` (refused on MLA,
recurrent kinds, sliding window — see its docstring). Caller
must run the gate; we re-check inside as a defensive guard.

Capture pre-flight gates.

Soft ``(bool, reason)`` predicates that decide whether a model's shape
is safe to capture into a CUDA graph. Refusals make the engine fall
through to the live (un-captured) path.

  * :func:`_can_capture_decode` — whole-forward decode capture.
  * :func:`_can_capture_prefill` — whole-forward prefill capture
    (stricter; honours the GDN-only opt-in).
  * :func:`_can_capture_mtp_verify` — the MTP verify-shape ``(B, K+1)``
    capture.
  * :func:`_hybrid_prefill_capture_opt_in` — reads
    ``ARBI_GDN_PREFILL_CAPTURE`` at call time.
  * :func:`_capture_unsafe_component` — the component-declared
    ``capture_safe`` marker, gating every shape.

Reason from any component that declares its own dispatch unsafe.

Duck-typed on a ``capture_safe`` marker, which is how a component
whose dispatch is decided AFTER the arch picked its layer class
reports the verdict: a routed-expert set only learns whether it runs
the fixed-shape grouped GEMM or the host-syncing per-expert-module
walk once the quant swap has bound (or failed to bind) every
projection. Unlike ``capture_safe_multirow`` this is not a
row-count qualifier — a False here means no capture at any shape,
so it gates decode as well as prefill and verify.

Reason this model's ATTENTION cannot be CUDA-graph-captured, else ``None``.

``allow_over_cap_decode`` is set by the DECODE caller when the paged backend
serves head_dim>``_FA_VARLEN_MAX_HEAD_DIM`` layers through a stream-capture-
SAFE decode kernel — the tkv split-K paged decode handles any head dim and
is a plain CUDA launch (capturable). In that case an over-cap head_dim is
NOT a decode-capture blocker. It stays a blocker for prefill / piecewise
capture (default ``False``), which run the Turbo prefill kernel that is not
stream-capture-safe.

``allow_mla_decode`` is the same shape of relaxation for
``MLA_SHARED``. Decode and MTP verify read the packed slots PAGED —
``tq_mla_batch_decode`` / ``tq_mla_batch_mtp_decode`` take each row's
K extent from the ``seq_lens`` buffer the replay refreshes, so one
captured graph serves every context length. The MULTI-TOKEN branch
(:meth:`MLAAttnOp._gathered_attend`) does not: it materialises a
gathered bf16 latent workspace of ``sum(seq_lens)`` rows and launches
the gather over that count, so a captured graph freezes ONE KV extent
while the ``cu_seqlens_k`` it indexes with is recomputed from the
live ``seq_lens`` on every replay. Decode capture never records the
multi-token branch, so MLA is not a decode blocker; it stays one for
prefill / piecewise capture.

Single source of truth for every graph-capture entry point — whole-forward
decode (:func:`_can_capture_decode`), whole-forward prefill
(:func:`_can_capture_prefill`), and the piecewise prefill sweep
(``cudagraph_admin.precapture_layer_graphs``). Routing every path through
one predicate is the whole point: a new capture entry point cannot
silently forget a guard, and a new capture-hostile attention kind is
registered in ONE place instead of being copied into each path (the bug
that let Turbo-prefill head_dim>256 layers slip into the piecewise sweep and crash at
replay). Two attention kinds cannot be recorded into a CUDA graph:

  * ``MLA_SHARED`` — the multi-token attend gathers a bf16 latent
    workspace whose row count is the batch's ``sum(seq_lens)``; a
    captured graph can only hold one such extent (see the
    ``allow_mla_decode`` paragraph above).
  * ``head_dim > _FA_VARLEN_MAX_HEAD_DIM`` — these layers (Gemma 4's
    head_dim=512 full-attention) route to the Turbo prefill backend, whose
    CuTeDSL / tvm-ffi launch is not stream-capture-safe: a recorded graph
    trips ``cudaErrorIllegalAddress`` at replay.

NOTE: Inductor compile-warmup is a SEPARATE concern — every attention
op here is registered through ``torch.library.custom_op`` with a fake
impl and a ``mutates_args`` declaration, so Dynamo traces all of them
(``tests/test_compile_op_fakes.py``). Capture legality and trace
legality are independent verdicts.

Capture legality is a STRUCTURAL property of the model — layer specs
and module markers only. There is deliberately no model-level opt-out
attribute: a class-string excuse would route every capture entry
point around the model with nothing louder than a log line, serving
it eager at a fraction of its speed with correct output. An arch
whose forward genuinely cannot be captured must surface here through
a structural fact (a ``StateKind``, a layer shape, a module marker
such as ``capture_safe_multirow`` or ``capture_safe``), and the
decode sweep turns the refusal into a loud boot error rather than a
silent eager fallback.

Soft pre-flight: refuse capture in unsupported scenarios.

``over_cap_decode_splitk``: the paged backend serves head_dim>256 decode via
the tkv split-K kernel (capturable), so over-cap head dims do not block
decode capture. Default ``False`` keeps the conservative behavior.

Capture targets PAGED_KV + recurrent (Mamba / GDN / ShortConv)
decode and is TP-aware: the capture region runs under
:func:`arbi_serve.distributed.graph_capture.graph_capture` which
pins ``GroupCoordinator.all_reduce`` to a side stream the graph
records against, so row-parallel layers' all-reduces capture
cleanly under TP>1. NCCL per-stream comm buffers are pre-warmed
BEFORE capture begins so capture itself sees no ``cudaMalloc``.

Sliding-window layers are supported: each per-layer
``TkvAttnOp`` carries its own ``sliding_window``, and tkv's
``DecodeAttend`` dispatcher picks the SWA-splitk vs legacy
kernel via autotune-on-shape (`_VARIANT_SWA_SPLITK` /
`_VARIANT_LEGACY` in ``runtime/attend/decode_attend.py``). Autotune
is deterministic per ``(B, H, max_seq_len)``, so the kernel
chosen at capture time is the same one autotune would pick at
replay time for the same shape — replay runs the captured kernel
correctly.

Caveat: capture happens with the synthetic
``capture_seq_len_value`` (≈ ``page_size``). For SWA layers this
is below the sliding window, so autotune picks the
``LEGACY`` kernel rather than ``SWA_SPLITK``. Replay uses
``LEGACY`` even at long runtime contexts; we forfeit
``SWA_SPLITK``'s grid-utilisation win on small-batch large-context
SWA decode. Cost is bounded: ``LEGACY`` is correct, just less
optimal at high seq_len with low (B × H_kv). Force-pin via the
``TKV_SWA_SPLITK`` env var if needed.

Recurrent kinds (Mamba / GDN / ShortConv) read their per-row
state through the per-layer slab tensor in
:class:`arbi_serve.cache.recurrent_pool.RecurrentStatePool`. The
slab's ``data_ptr`` is stable for the engine's lifetime, and the
captured graph holds one persistent ``state_indices`` buffer per
recurrent kind in ``CapturedGraph.recurrent_bufs``; the engine
``copy_()``s the live mapping in before ``graph.replay()``.
Verified end-to-end by
``tests/test_hybrid_decode_capture_live_gpu.py`` (live-GPU smoke:
boot, replay-vs-eager parity, MTP verify-shape capture, and
sleep/resume preservation on a real Qwen3.5 hybrid checkpoint).

LoRA-aware capture is wired through the engine's
:class:`arbi_serve.adapters.lora.LoraCapturePool`: per-step
``LoraBatchState`` is ``copy_()``d into worst-case persistent
buffers, and the captured graph is keyed on a quantized active-
LoRA bucket ({0, 1, 2, 4, 8} by default). Replay finds the
matching bucket and uses zero-padded tail slots when the runtime
active count is below the bucket. No pre-flight rejection — LoRA
in flight just selects a different captured graph.

Mismatches return ``(False, reason)`` and the engine quietly
falls back to the live path.

MLA decode IS captured (``allow_mla_decode=True`` below): the only
capture-hostile step on the MLA path is the host sync in the
multi-token ``_gathered_attend`` branch, which a query-length-1
decode never reaches. Prefill / piecewise capture still refuse.

Read ``ARBI_GDN_PREFILL_CAPTURE`` at call time.

The vendored FLA path
(:mod:`arbi_serve.kernels.fla_vendored`) provides a captured-safe
forward path for GDN prefill via the vendored
``chunk_gated_delta_rule_inference_fwd`` entry.  Setting this env
var to ``1`` opts the model into whole-forward prefill cudagraph
capture even when GDN layers are present.  **Default ON**.  Set
``ARBI_GDN_PREFILL_CAPTURE=0`` to opt out.

Performant options default ON.

Companion env var ``ARBI_GDN_VENDORED_FLA`` also defaults ON.

Soft pre-flight for whole-forward prefill cudagraph capture.

Stricter than :func:`_can_capture_decode`: prefill replays one
captured graph for the ENTIRE forward (every layer, every
metadata kernel) at a fixed ``num_tokens`` bucket, so any single
layer kind that doesn't fit the persistent-buffer contract poisons
the whole capture. Refusal cases:

  - MLA — the gathered-latent workspace extent is per-batch data
    (``sum(seq_lens)``); see
    :func:`_attention_graph_capture_unsafe`.
  - Recurrent layer kinds (Mamba / GDN / ShortConv). Recurrent
    kernels advance per-token state on every prefill token; the
    per-bucket persistent buffers can't represent the full state-
    evolution trace (decode replays one row, prefill replays N).
    Hybrid Qwen 3.5 / 3.6 (interleaved GDN) hits this and falls
    through to the eager / piecewise path.

    **GDN-only opt-in**: when the vendored FLA path is in use and
    the caller sets ``ARBI_GDN_PREFILL_CAPTURE=1``, the pre-flight
    admits GDN-only hybrids (Mamba / ShortConv still refuse —
    their captured-safe path is not yet implemented).
    See :mod:`arbi_serve.kernels.fla_vendored` for the rationale.
  - Sliding-window layers (Gemma 4 hybrid). The captured
    ``block_table`` shape doesn't tolerate SWA's per-position
    windowing: at decode time SWA reads back ``min(seq_lens,
    sliding_window)`` rows per query, but a captured prefill
    graph bakes the full block-table read pattern at capture-
    time positions, which doesn't match the runtime per-token
    positions a real prefill walks.
  - Models with an MTP-marked layer still pass. The main decoder
    loop skips ``is_mtp_layer=True`` specs, but the captured
    forward MAY still run the bundled MTP drafter-KV fill after the
    loop (``Qwen3_5Model._mtp_kv_fill``) — ``capture_prefill``'s
    ``mtp_fill`` flag pins whether it does, and the graph records
    the outcome in :attr:`CapturedGraph.mtp_fill`.

Mismatches return ``(False, reason)`` and the caller logs +
falls through to the live (un-captured) path.

Pre-flight for the MTP verify-shape capture at draft count ``K``.

The verify forward shape is ``(B, K + 1)`` (one ``last_committed``
slot + ``K`` draft slots per row). Capture targets the served verify
dispatch at ``BLOCK_M = K + 1``, so ``K + 1`` must fit under the
SERVED verify ceiling (``tkv_max_verify_block_m()`` — the split-K
register kernel plus the capture-safe Turbo prefill verify above it; see
the body's rationale for why the split-K register ceiling alone
would wrongly refuse served high-K shapes).

Refusal cases:
  - ``backend_supports_capturable_verify`` is ``False`` — the active
    PAGED_KV backend has no capturable verify kernel. Both tkv codec
    modes (compressed codec AND tkv-bypass raw-bf16) DO have one: they
    route the (B, K+1) verify shape to the split-K paged MTP kernel
    (``tq_splitk_batch_mtp_decode``; head_dim up to 512), a plain CUDA
    launch with no host sync — so they report ``True`` and the verify
    graph captures (including Gemma-4 head_dim=512 full-attention
    layers). This False branch is for a hypothetical backend whose
    verify would run a capture-unsafe kernel; it then runs eager by
    design and refusing here keeps the sweep from recording it.
  - ``K < 1`` — no draft tokens, no MTP shape (the K=0 cold-path
    step takes the plain S=1 decode capture if any).
  - ``K + 1 > tkv_max_verify_block_m()`` — above the served verify
    BLOCK_M ceiling.
  - ``K > max_k`` — driver cap (``MtpDriver.max_k`` chosen at
    attach time from ``cfg.mtp.n_draft``); a per-step
    ``step_k > max_k`` is impossible by construction so capturing
    beyond the cap is wasted work.
  - Underlying decode pre-flight fails — same surface as
    :func:`_can_capture_decode` (today this rejects MLA models;
    recurrent state kinds — GDN / Mamba / ShortConv — pass
    because the captured graph carries per-kind
    ``state_indices`` buffers via
    :attr:`CapturedGraph.recurrent_bufs`).

LoRA in flight is not a rejection here either: the LoRA bucket
sweep in :func:`precapture_decode_graphs` captures additional
``(B, K + 1, lora_bucket)`` graphs and the lookup picks the
matching bucket. Same contract as plain decode.

Captured-graph bucket / lookup gates for :class:`EagerModelRunner`.

These are the cold-ish "which captured graph, if any, serves this
batch" decision functions, implemented as module-level free functions
taking the runner explicitly. The runner keeps thin delegating method
wrappers so every ``runner._lookup_*`` / ``runner._captured_lookup``
call site (engine path + tests) is preserved byte-for-byte.

The hot per-step ``execute`` / ``forward`` methods call these via the
runner's wrappers; the lookup itself is a few dict gets and integer
comparisons — no kernels, no allocation — so the extra function hop is
in the noise.

This module is a thin public facade. The implementations live in
three sibling modules, each kept under 1000 lines and co-locating its
MUST_FIRE flag-truth counters with the fire sites:

  * :mod:`arbi_serve.runtime.captured_lookup_buckets` — LoRA / KV-page
    bucket-key quantizers and the core ``captured_lookup`` pool accessor.
  * :mod:`arbi_serve.runtime.captured_lookup_gates` — the decode /
    prefill / mixed / verify / pad-up decision gates (``_PREFILL_LOOKUP``,
    ``_PREFILL_LOOKUP_CONT``, ``_MIXED_LOOKUP``).
  * :mod:`arbi_serve.runtime.captured_lookup_recurrent` — recurrent
    slab-row resolvers (``_RECIDX_REUSE``).

Every name below is re-exported here so both the engine path
(``arbi_serve.runtime.captured_lookup.<name>`` attribute access) and
tests (``from arbi_serve.runtime.captured_lookup import <name>``)
resolve unchanged.

Bucket-key helpers + core pool ``get`` for captured-graph lookup.

The LoRA / KV-page bucket quantizers and the single ``captured_lookup``
pool accessor. They take the runner explicitly and are re-exported from
:mod:`arbi_serve.runtime.captured_lookup` so every call site
(engine path + tests) is preserved byte-for-byte.

The ONE pool :func:`captured_lookup` resolves against, for callers
that must ENUMERATE captured shapes rather than ask for one key.

Enumeration and resolution must read the SAME object. Returns
``eng.captured_graphs`` — the active member's own pool (stable-VA model
cycling repoints ``cudagraph_pools`` wholesale per resident member via
``ENGINE_MODEL_ATTRS``).

Pick the smallest configured KV-page bucket that covers
``ceil(max_seq_len / page_size)`` pages.

Returns the bucket size (an entry from ``buckets``), or ``None``
when ``buckets`` is empty (legacy single-shape capture: lookup
falls back to the ``kv_pages_bucket = 0`` sentinel — the same
pool key pre-bucketing capture paths register under) or when
every configured bucket is too small (live request exceeds the
largest bucket → captured-graph miss → run loop falls through
to the live forward, no regression).

NOTE (#1): the boot sweep (``precapture_decode_graphs``) PINS the
top captured bucket to ``ceil(max_context / page_size)``, so for
the captured sweep this ``None``-from-too-small branch is
unreachable for any ADMISSIBLE sequence (admission caps prompts at
``max_context``). The branch remains a correctness backstop for the
out-of-contract case and for callers passing an arbitrary
``buckets`` tuple (e.g. unit tests).

The captured-graph pool's persistent ``block_table`` for a given
bucket is sized to ``(B, bucket)`` columns; the lookup picks the
smallest bucket so the captured kernel's gather reads the
fewest pages while still covering every valid token. The
eager-path block-table bounding short-circuits on
``cap_pages == max_pages_alloc`` so the bucket-bound capture
runs zero allocator pressure inside the captured region.

``True`` when the operator pinned a non-empty bucket sweep.

Empty sweep ⇒ legacy single-shape capture (every captured
graph is keyed at ``kv_pages_bucket = 0`` — the pre-bucketing
sentinel) and lookups skip the bucket dim entirely.

Resolve the KV-page bucket key for ``batch`` against the
operator-pinned ``cfg.cudagraph_kv_pages_buckets`` sweep.

Returns ``None`` when bucketing is disabled (legacy single-
shape capture) OR when the live ``max_seq_len`` exceeds every
configured bucket. Callers gate the latter as a captured-
graph miss (fall through to the live forward, no regression).

The candidate bucket set MUST mirror exactly what the boot sweep
(:func:`arbi_serve.engine.cudagraph_admin.precapture_decode_graphs`)
captured under, or the lookup asks for a bucket key the pool never
registered and every step misses → fully-eager decode. The boot
sweep does TWO transforms to the raw ``cfg.cudagraph_kv_pages_buckets``
before capturing: it CLAMPS each entry to ``max_context_pages``
(``capture_decode`` registers under ``min(bucket, max_context_pages)``)
and PINS ``max_context_pages`` as the top rung. We apply the same two
transforms here so a live ``max_seq_len`` whose ``needed_pages`` lands
between a raw rung and the clamp (e.g. raw ``(1, 4, 16, 64)`` with
``max_context_pages=8``: a 5-page request) resolves to the captured
bucket ``8`` instead of the uncaptured raw bucket ``16`` — avoiding
the long-context decode cudagraph-miss cliff once the sequence
passes the largest raw bucket below ``max_context_pages``.

Return the captured graph for
``(batch_size, seq_len, lora_bucket, is_prefill,
kv_pages_bucket, prefill_context)`` from the active member's pool.

``prefill_context`` selects the whole-forward prefill capture
FAMILY: ``False`` = first-chunk (bypass branch frozen),
``True`` = continuation-chunk (paged branch frozen). Defaults
``False`` so decode / MTP-verify / first-chunk callers are
unchanged.

Captured-graph decision gates (decode / prefill / mixed / verify).

The "which captured graph, if any, serves this batch" decision
functions. Each takes the runner explicitly and is re-exported from
:mod:`arbi_serve.runtime.captured_lookup` so every call site (engine
path + tests) is preserved byte-for-byte.

The flag-truth MUST_FIRE counters (``_PREFILL_LOOKUP``,
``_PREFILL_LOOKUP_CONT``, ``_MIXED_LOOKUP``) live here, co-located with
their fire / refuse sites in the gates below.

Return the whole-forward prefill graph for ``batch``, or ``None``.

``require_hidden``: the caller will replay with
``return_hidden_state=True`` (the MTP-opted final-chunk seed forward),
so a graph without a retained ``hidden_out`` buffer cannot serve it —
refuse to eager instead of tripping the replay-time guard. Host
attribute read only.

Refusal cases (any one ⇒ live path):

  - ``cuda_graphs`` disabled or no prefill buckets configured.
  - Batch is not prefill (decode / verify routes through
    :func:`lookup_captured_graph`).
  - More than one batch row. Whole-forward prefill capture
    is single-sequence by construction (one captured graph per
    ``num_tokens`` bucket at ``B=1``); multi-row prefill —
    rare under the current scheduler — falls through to live.
    Batch-keyed on purpose (like the pad-up lookup): at TP every
    rank resolves this rung from the same derived / mirrored
    batch, so all ranks replay the same bucket graph or all run
    eager — never a mixed replay/eager split across ranks.
  - LoRA in flight: prefill capture only covers LoRA capture pool wired yet).
  - Multimodal prefill (``batch.mm`` / ``batch.mrope_positions``
    set): the captured graph carries no media scatter, so a
    replay would silently drop the media embeddings.
  - ``real_N`` exceeds every CAPTURED rung. The rung set comes from the
    pool (``CapturedGraphPool.prefill_rungs``), not from
    ``cfg.prefill_cudagraph_buckets`` — see the selection site below for
    why the configured tuple is the wrong source.
  - **MTP-fill variant mismatch**: the graph's baked drafter-KV-fill
    state (:attr:`CapturedGraph.mtp_fill`) differs from what the eager
    forward would do for this batch (``Qwen3_5Model._mtp_kv_fill``'s
    executable gate). A fill-baked graph replayed on a no-MTP slate
    pays a dead MTP decoder layer + paged KV write every chunk; a
    fill-less graph replayed on an MTP slate starves the drafter of
    prefill KV and collapses accept rate. Host-int compare, both
    sides fall back to eager (which honours ``mtp_fill_enabled``).
  - **CONTINUATION CHUNK with no continuation family, or one whose
    continuation gates trip** (``max_seq_len > real_N``). See below.

Two-family dispatch for chunked prefill
---------------------------------------
A continuation chunk carries ``context_len = max_seq_len - real_N > 0``
tokens of KV from earlier chunks; its attention must gather that prefix.
Whole-forward prefill capture records TWO families per bucket
(``precapture_prefill_graphs``):

  * ``prefill_context=False`` — first chunk (``q_len == seq_len``). The
    host route gate resolves the BYPASS branch at capture; it attends the
    forward's own contiguous K/V and reuses ``cu_seqlens_q`` as
    ``cu_seqlens_k``.
  * ``prefill_context=True`` — continuation chunk (captured at
    ``seq_len = ctx + N > q_len``). The gate resolves the PAGED branch,
    whose per-row K extent is a device read of ``seq_lens``
    (``mSeqUsedK``), so ONE captured graph attends the correct window at
    every context length the replay refreshes it to.

A host branch frozen in a captured region is a defect only when the WRONG
branch was frozen. The route is constant WITHIN a family, so keying the
lookup on continuation-ness hands each family exactly the branch it
needs and never gives a first-chunk graph to a continuation chunk.
Proven bit-exact end-to-end on real weights: on a dense model a needle
planted in a continuation chunk is recalled byte-identical to a
pure-eager reference, with the continuation replay confirmed fired
(``test_prefill_capture_chunked_live.py``).

Continuation correctness gates (host-int, no D2H) — fall back to eager:

  * full-sequence K extent past the captured block-table width: the
    in-graph paged gather would silently truncate the K window.

Recurrent (GDN-hybrid) captures are NOT refused: the continuation replay
refreshes the captured persistent ``recurrent_bufs`` / int64 alias /
block-table buffers per replay (``_replay_prefill_padded`` →
:meth:`CapturedGraph.replay`), the captured ``has_initial_state`` is
all-True (fresh rows read the zero-cleared slab — byte-identical to "no
initial state"), and the in-graph GDN gather + final-state writeback go
through ``meta.state_indices_long`` = the refreshed persistent buffer.
Validated on Qwen3.5-0.8B (hybrid GDN), both backends, with a
planted-needle continuation run: continuation-chunk logits
AND post-chunk recurrent/conv slab rows match the eager forward within the
eager null-control floor, with the continuation family confirmed fired.

Padding-replay contract. We pick the SMALLEST configured bucket
``bucket_N`` such that ``real_N <= bucket_N`` and replay the
captured graph at that bucket with the live request padded to
``bucket_N`` tokens. The caller (``execute``) refreshes
``cu_seqlens_q = [0, real_N]`` (the query span, always this chunk's
tokens) and ``seq_lens = cu_seqlens_k[1] = [max_seq_len]`` (the full K
extent) per replay, so the model's captured last-token gather
(``hidden[cu_seqlens_q[1:] - 1]``) returns the correct ``real_N - 1``
position and the Turbo prefill kernel honours per-row Q range + ``seqused_k`` for the actual
K length. On a first chunk ``max_seq_len == real_N``; on a continuation
chunk ``max_seq_len == context_len + real_N`` and the paged branch reads
that extent from ``seq_lens``.

Padded positions ``[real_N..bucket_N)`` are processed for shape but
ignored for content; padded ``slot_mapping`` entries scatter K/V into
the page-0 null sentinel so valid cache rows are not corrupted.

``real_N == bucket_N`` is the exact-match path and still hits the
same code; the cu_seqlens_q refresh is a no-op copy in that case.

Mixed decode+prefill whole-forward lookup (``ARBI_MIXED_CAPTURE``).

Returns ``(captured, d_cap, bucket_n)`` — the smallest captured MIXED
graph whose canonical layout (``d_cap`` decode-row slots + one
prefill row of ``bucket_n`` tokens) covers the live composition — or
``(None, 0, 0)`` to fall through to the split / eager paths.

Mixed graphs are registered by :func:`~arbi_serve.runtime.capture.
mixed.capture_mixed` under ``(B = d_cap + 1, S = bucket_n,
lora_bucket = 0, is_prefill = True, kv_pages_bucket = 0)``. ``B >= 2``
keeps them disjoint from the single-row prefill captures (``B == 1``),
and ``is_prefill=True`` from every decode / verify capture, so this
enumeration can never pick a non-mixed graph.

Refusals (any ⇒ ``(None, 0, 0)``; every edge is attributed on the
``mixed_graph_lookup`` flag-truth counter with a reason, mirroring
the single-row prefill lookup above):

  * cuda graphs disabled / batch not flagged prefill (a pure-decode
    step routes through :func:`lookup_captured_graph`).
  * slate is not "N single-token decode rows + exactly ONE prefill
    row" (multi-prefill co-admits and verify rows keep today's
    paths — correctness over coverage).
  * no graph in the prefill row's FAMILY (first chunk vs
    continuation chunk, keyed on ``prompt_consumed > 0``) covers the
    live shape. The two families freeze different attention branches;
    see the inline note below.
  * LoRA in flight (mixed capture only covers ``lora_bucket = 0``).
  * no captured mixed shape with ``d_cap >= d_live`` AND
    ``bucket_n >= real_N`` within ``max_batch``.

Return the captured graph for an MTP-verify-shape forward.

Matches ``(batch_size, max_query_len, lora_bucket)`` against
the captured pool. Returns ``None`` when capture is disabled,
the batch is prefill, ``B * S != num_actual_tokens`` (a
defensive guard — under R1 this never trips because the
engine coerces verify slates to uniform K), or no graph was
captured for this shape. ``max_query_len`` may be ``1``
(plain batched decode through ``forward``) or ``> 1`` (MTP
verify with K+1 tokens per row).

LoRA in flight is NOT a rejection: the lookup picks the
captured graph keyed at the matching active-LoRA bucket
(see :func:`arbi_serve.adapters.lora.quantize_active_loras`). When no
bucket matches (e.g. operator only captured the 0-bucket),
the lookup misses and the engine falls through to the live
path.

Partial-accept replay forwards pin ``batch.state_indices`` to
the verify-batch row indices for those rows. The captured
graph's ``recurrent_bufs[kind]`` buffer accepts an arbitrary
per-replay mapping (``CapturedGraph.replay`` ``copy_()``s the
live tensor in via ``recurrent_state_indices``), so partial-
accept replays can route through the captured graph too.

The captured graph at (B, S) bakes ``cu_seqlens_q = [0, S,
2S, ..., B*S]`` at capture time — that's the ONE input not
re-copied on replay (see :meth:`CapturedGraph.replay`).
``positions``, ``slot_mapping``, ``seq_lens``,
``cu_seqlens_k``, ``block_table``, ``input_ids`` are all
copy_()d in fresh on every replay. The
``B * S == num_actual_tokens`` gate is a defensive double-
check: under R1 the engine coerces every verify slate to
uniform K so ``cu_seqlens_q == [0, S, 2S, ..., B*S]`` already
matches the captured shape; the gate exists to fail loud on
a regression rather than silently produce wrong results.

Pad-up decode / verify lookup: return ``(captured, pad_to_B)``
for the smallest captured shape ``B' >= B`` at the SAME
``(S, lora_bucket, kv_pages_bucket)`` key, or ``(None, 0)``.

Covers BOTH decode flavors with one batch-keyed decision:

  * ``S == 1`` — plain batched decode. The generated
    ``cudagraph_shapes`` B-ladder is sparse; a live ``B`` between
    rungs misses the exact lookup and would run eager.
  * ``S > 1`` — MTP verify. The verify pass coerces every slate to
    a uniform per-step ``K`` (R1's uniform-K invariant), so the
    live verify batch is always flat ``B × S`` with ``S = K + 1``
    and ``cu_seqlens_q = [0, S, 2S, ..., B*S]``; the captured
    ``(B, K+1)`` ladder only covers the plain-decode rungs, so an
    off-rung ``B`` (5/6/7 under continuous arrivals) misses.

Batch-keyed on purpose: at TP every rank resolves this rung from
the same derived / mirrored batch, so both ranks pick the same
``B'`` graph (or both miss) and the in-graph collectives pair
1-to-1 — the decision reads nothing a peer rank could disagree on.

Used only when the exact-``B``
:func:`lookup_captured_graph_for_forward` missed AND
``cfg.decode_pad_cudagraph`` is on. The caller pads the live
batch from ``B`` to ``B'`` with benign scratch rows (each
contributing exactly ``S`` flat tokens so the padded
``cu_seqlens_q`` still matches the captured graph's baked
``[0, S, ..., B'*S]``) and replays the captured ``B'`` graph,
discarding the padding rows' logits.

Refusals (any ⇒ ``(None, 0)`` → caller takes the eager forward):
  * cuda graphs disabled / prefill batch (the prefill rung owns
    ``is_prefill`` batches).
  * ``S > 1`` without ``mtp_meta`` (non-verify S>1 shape).
  * ragged flat (``input_ids != B * S`` — would break the baked
    ``cu_seqlens_q``).
  * KV-page bucket uncovered (same gate as the exact lookup).
  * no captured shape ``B' >= B`` at this ``(S, lora, kv)`` key,
    or ``B' > cfg.batch.max_batch`` (cannot pad past the
    recurrent-slab / scheduler capacity — the free-list borrow
    needs ``B' <= max_batch``).

Enumeration wiring: enumerate the pool :func:`captured_lookup` will
RESOLVE against, not a different one.
:func:`~arbi_serve.runtime.captured_lookup_buckets.active_captured_pool`
is that single resolution (``eng.captured_graphs`` — the active
member's pool, which stable-VA model cycling repoints wholesale per
resident). Enumerating a different store would let ``candidates[0]``
pick a rung the resolve pool lacks, miss, and return ``(None, 0)``
without ever trying a larger ``B'`` the active pool DOES have — pad-up
would read as wired and silently not fire.

Recurrent slab-row resolution for captured-graph replays.

The per-row ``req_id → recurrent slab row`` resolvers. Each takes the
runner explicitly and is re-exported from
:mod:`arbi_serve.runtime.captured_lookup` so every call site (engine
path + tests) is preserved byte-for-byte.

The persistent-index MUST_FIRE counter (``_RECIDX_REUSE``) lives here,
co-located with its fire site in :func:`live_recurrent_state_indices`.

Return per-row recurrent slab-row indices for ``slate``.

Used by :meth:`_build_batch` to write a stable mapping into the
persistent ``pb.recurrent_state_indices`` slice when piecewise
prefill capture is in play. Returns ``None`` when no recurrent
pool is registered (the model is pure-attention; no recurrent
meta to build). On a partially-allocated pool (the caller's
contract violation — every slate request MUST have its
recurrent row pre-allocated by
:meth:`MultiStatePool.alloc_recurrent_state` at admission)
raises :class:`RuntimeError` to surface the bug
loudly instead of silently routing the request to slab row 0.

A naive fallback of ``arange(len(slate))`` would use slate
position, NOT a per-request row. Combined with B=1 decode
that routes every request to live row 0, leaking the prior
request's GDN/Mamba state. The sub-builder in
:mod:`arbi_serve.backends.sub_builders.state_indices` applies
the same RAISE-on-unallocated treatment; this
site is the model-runner-level mirror of that policy so a
future regression hits the loud error here BEFORE reaching
the sub-builder.

Build live ``state_indices`` per recurrent kind for a captured replay.

Returns ``{StateKind: int32 (B,)}`` or ``None`` when the
captured graph has no recurrent buffers (pure-PAGED_KV).

Memory accounting: per-step state-indices tensors
are routed through the ``activation_arena`` named pool by the
caller in :meth:`execute`. This helper does not enter the
pool ctx itself — :class:`NamedMemPool.use` is re-entrant
(refcounted), so nesting would be safe, but keeping the
single-entry pattern at the caller avoids the double-entry
no-op on the hot path.

Persistent-buffer fast path (``ARBI_PERSISTENT_RECURRENT_INDEX``,
default ON): all recurrent kinds in a single arch share the SAME
per-row ``req_id → slab row`` mapping (one slate row consumes the
same slab row across every recurrent layer — see
:func:`resolve_recurrent_rows`). The mapping only changes on
admit / evict / preempt, NOT on a steady-state decode step, so this
helper resolves the host row list ONCE and caches it alongside a
persistent device int32 buffer on the runner. When the next step's
slate membership matches the cached signature (same req_ids, same
order, same recurrent-pool identity) the cached device slice is
returned for every kind UNTOUCHED — no Python ``row_for`` loop, no
fresh ``torch.tensor`` H2D. On any membership / pool change the
buffer is patched in place (one host loop + one H2D) and the
signature refreshed. ``CapturedGraph.replay`` ``copy_()``s the
returned tensor into its OWN persistent ``recurrent_bufs[kind]``
before launching, so handing the same source slice to every kind
(and reusing it across steps) is safe and bit-identical to the
per-kind ``torch.tensor`` rebuild.

Like :func:`live_recurrent_state_indices` but for ``forward(batch)``.

Resolution order for the per-row slab-row mapping:

  1. ``batch.state_indices`` when set — the verify pass / partial-
     accept replay pin it explicitly.
  2. ``batch.req_ids`` via the recurrent pool view's ``row_for`` —
     the SPMD plain-decode batch (:meth:`DistributedEngineDriver.
     _spmd_batch_from_derived`) leaves ``state_indices=None`` and
     carries ``req_ids`` so the metadata builders resolve each row's
     REAL slab row through ``RecurrentStatePool.row_for``. The
     captured-replay refresh MUST use the SAME route — otherwise it
     would fall back to ``arange(B)`` (the capture-time default) and
     the GDN/Mamba conv + recurrent kernels would read the WRONG
     (unallocated, zeroed) slab row, degenerating hybrid TP2 decode
     under cudagraphs (the eager SPMD path is correct because the
     metadata builder already resolves via ``row_for``; only the
     captured replay bypassed it).
  3. ``arange(B)`` from the seq count — last-resort fallback for
     callers that supply neither (CPU tests / pre-admission shapes).

Memory accounting: per-step ``.to()`` / ``arange()`` transients
are routed through the named pool by the caller (the ``forward``
method's captured-replay branch).

cuMem-backed pluggable allocator for sleepable ``torch.cuda.MemPool``.

Phase-2 sleep currently leaves the big torch ``NamedMemPool``s resident
because those pools are backed by the ordinary CUDA caching allocator,
which has no "unmap physical pages but keep the VA" primitive. This
module provides the missing primitive as a ``torch`` *pluggable*
allocator: every allocation routed through it is individually backed by
a CUDA virtual-memory reservation (``cuMemAddressReserve`` +
``cuMemCreate`` + ``cuMemMap`` + ``cuMemSetAccess``). On
:meth:`CuMemPoolAllocator.sleep` the physical pages are unmapped (and
optionally dumped to pinned host RAM first) while the VA reservation —
and therefore every live tensor's ``data_ptr()`` — stays put. On
:meth:`CuMemPoolAllocator.wake` fresh physical pages are mapped back at
the SAME VA, so captured cudagraphs and any held tensor views keep
working.

Two properties distinguish the shim:

  - the C shim is built at import via ``cpp_extension.load_inline``
    (no pre-compiled ``.so`` shipped in the wheel), and
  - allocations are tagged so ``sleep`` can offload some tags (copy to
    host, restore on wake) while discarding others (no restore).

This module is gated behind a ``cuMem*`` driver probe. On CPU-only
hosts, or where the driver lacks the virtual-memory API, the module
imports cleanly and :func:`driver_available` returns ``False`` — every
public entry point then raises a clear error rather than crashing at
import.

Scope: single-device / single-rank, mechanism only. No engine wiring.

``torch.cuda.memory_reserved(device)``, clamped to what's physically
possible right now.

``memory_reserved()`` can go STALE (over-report) once a cuMem-tagged
pool's physical is released through the raw driver-level unmap+release
path (:meth:`CuMemPoolAllocator.sleep`, and therefore
``config_variant.adrop_variant``'s namespace drop): that path frees real
physical — confirmed against ``nvidia-smi`` and this same device's OWN
``mem_get_info`` — without going through torch's block/segment free()
hook, so torch's internal ledger never learns the bytes came back and
keeps counting them as "reserved" forever. GPU-measured on a repeated
config-variant ``drop_previous`` walk: the raw reading climbs roughly the
size of ONE dropped member's pools per drop and, left unclamped,
eventually exceeds the card's own total physical while real usage stays
flat. Left unfixed this compounds into a real ``BudgetExceeded`` refusal
wherever a caller persists this reading as a future sizing prediction
(:func:`arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.
torch_caching_reserved_overhang_bytes`) or reports it as a live "leak"
signal (:mod:`arbi_serve.engine.phase2_freeze`) — GPU-reproduced this
session after ~15 drops.

``mem_get_info`` is a direct driver query (always accurate); reserved
bytes can never legitimately exceed real in-use physical
(``total - free``), so clamp to that hard fact rather than trusting a
torch-internal counter that can only drift upward. Returns 0 on any
driver read failure (no CUDA / dead context) — same contract as a bare
``memory_reserved()`` call would raise into, made non-fatal for callers
that already wrap this in a best-effort block.

``external_mapped_bytes`` tightens that clamp to the bound that is actually
physical: the mapped bytes of every cuMem region this process drives
DIRECTLY — the growable KV region, the state pools' sentinel-alias arenas —
rather than through the torch caching allocator. Those pages are resident
physical torch does not and cannot hold, so the most torch can legitimately
have reserved is ``(total - free) - external``. Bounding at ``total - free``
alone leaves a stale reading room to grow across the ENTIRE external
region's mapped bytes are known; 0 keeps the looser device bound for
callers that legitimately have no KV region.

cuMem bytes the allocator tracks but has UNMAPPED (physical released).

``tracked_bytes − mapped_bytes`` over the process-singleton pluggable
allocator: the discard-slept allocations whose physical is gone but which
``torch.cuda.memory_reserved`` still counts. Uses :meth:`CuMemPoolAllocator.
peek` so a read-only accounting call never constructs the native allocator on
a boot that does not otherwise use it. Returns 0 when no allocator exists.

Report ONCE that ``memory_reserved`` exceeded what the device can hold.

The clamp keeps the arithmetic honest; this says the counter is wrong, so
the condition is visible rather than silently absorbed.

The UNQUALIFIED pool name of a (possibly namespace-qualified) tag.

Caps and the per-tag byte counters are keyed on this — see
:meth:`CuMemPoolAllocator.set_tag_cap` for why. Pool names live in the flat
``category.role`` taxonomy (``capture.cudagraphs``, ``model.weights``,
``scratch.forward_arena``, …) and never contain ``/``; a tag NAMESPACE is a
model key (typically a filesystem path such as ``/models/Qwen3.5-0.8B``) and
is joined with ``/``. Taking the segment after the LAST ``/`` therefore
recovers the pool name from either form, and is idempotent on an already-
unqualified name.

One SUB-ALLOCATION carved out of another tag's already-mapped bytes.

A slab fold maps ONE physical segment under a host pool's tag and hands
out ``narrow()`` views of it to several logical consumers. The views map
no physical of their own, so without this record the host tag would carry
every consumer's bytes and the consumers would report zero — one opaque
ledger row and one cap where there were several.

A charge PARTITIONS the host tag's mapped bytes: the sub-tag reports the
carved bytes, the host reports what it maps minus everything carved out of
it. The sum over tags is unchanged, so the VRAM identity still closes.

One armed per-tag map-time budget.

``enforce=False`` is the OBSERVE posture: the cap is NOT applied (the
allocation proceeds) but every breach is recorded and reported LOUDLY. It
exists for cap VALUES that have never been validated against a real boot —
arming an unvalidated value as a hard deny can refuse a boot that would
otherwise succeed, and an over-tight cap costs KV exactly as surely as an
absent cap costs safety. OBSERVE is never a silent default: every observed
cap is logged at WARNING on every report, and a breached observed cap at
ERROR (see :func:`arbi_serve.engine.pool_caps.log_cap_report`).

Process-singleton cuMem-backed pluggable allocator.

Usage::

    alloc = CuMemPoolAllocator.get()
    pool = alloc.make_pool()  # the registered-at-birth MemPool chokepoint
    with torch.cuda.use_mem_pool(pool), alloc.tag("weights"):
        w = torch.empty(..., device="cuda")
    alloc.sleep(offload_tags={"weights"}, discard_tags=set())
    ...
    alloc.wake()

Every allocation made while a tag is active is tracked; ``sleep``
unmaps physical pages (offloading tagged contents to pinned host RAM
first, or discarding them) while keeping the VA, and ``wake`` binds
fresh pages back at the SAME VA.

Copy ``nbytes`` at device ``ptr`` D→a fresh pinned-host tensor.

The bytes are page-locked and whole-VRAM-scale, so they are booked through
the host ledger BEFORE the allocation and given back when the tensor is
collected — a sleep the host cannot hold refuses here rather than putting
the kernel's OOM killer in charge of which process dies. The release rides
a ``weakref.finalize`` on the tensor because ownership passes to the
caller: booking bytes with no release path leaves the ledger over-reporting
and shrinks every later reservation for bytes nobody holds.

A cuMem region whose VA is reserved at a MAX size but whose physical
backing covers only a growing prefix (Option B / D1 primitive).

Lifecycle::

    region = GrowableRegion(max_bytes=..., device=0)  # reserves VA only
    region.map_to(K)  # back first K bytes
    region.byte_view(0, K).fill_(1)  # write the mapped part
    region.map_to(K + N)  # grow at the SAME VA
    region.byte_view(K, N).zero_()  # write the grown part
    region.unmap_all()
    region.close()  # reclaim

Only ``mapped_bytes`` of physical is backed at any time — the rest of the
reserved VA costs nothing physical (``torch.cuda.mem_get_info`` reflects
only the mapped pages). All sizes/offsets are rounded UP to the device
granularity. The KV pool builds these to back its slabs
(:mod:`arbi_serve.cache.paged_kv_pool`, Option B).

True when the cap has provably never bound and was expected to.

A cap configured for a tag that is never allocated (a typo, a renamed
pool, or — the bug this machinery exists to make impossible — a
key-space disagreement between writer and reader) is silent no-op
protection. It must be surfaced, never assumed to be holding.

A forward-looking budget (``expect_bound=False``, e.g. the LoRA
admission cap) is exempt: zero allocations is its normal state.

Initialise the pure-Python bookkeeping state (no driver, no C shim).

BINDS — is constructible and
directly testable on a CPU-only host, where the driver probe fails and
the C shim cannot be built. The native wiring stays in
:meth:`_init_native`.

The singleton if it already exists, else ``None`` — never creates.

:meth:`get` is get-or-CREATE, and construction calls
:meth:`_init_native` (ctypes-loads the C shim, registers the malloc/free
callbacks, builds the torch pluggable-allocator handle). A read-only
observer — accounting, logging, triage — must never bring the native
allocator into existence on a boot that does not otherwise use it.

Stop the C shim calling back into Python. Idempotent.

After this the free hook returns the VA with ``cuMemAddressFree`` and
touches no Python object. Physical handles outstanding at this point
are left to process exit, which is the only moment this runs.

Disable ``expandable_segments`` for the duration of a cuMem-pool op.

``torch.cuda.MemPool`` construction AND ``use_mem_pool`` / allocation
into it both refuse to run while ``expandable_segments:True`` is active
(pytorch#147851).

The engine pins the setting OFF process-wide before first CUDA init —
:func:`arbi_serve._force_expandable_segments_off` on import, and the
image's ``PYTORCH_ALLOC_CONF`` / ``PYTORCH_CUDA_ALLOC_CONF`` ENV — so on
a shipped boot this scope has nothing to strip and its restore-to-True
arm never fires. It stays because the pin is a runtime fact, not a type:
an embedder that turns the setting back on mid-process still gets a
working cuMem op here rather than a ``RuntimeError`` deep in the loader.
cuMem pools route through THIS pluggable allocator, not torch's
expandable segments, so toggling it off for the cuMem op is a no-op for
them — and whatever the caching allocator had is restored on exit. Wrap
the WHOLE cuMem pool lifecycle (construct + ``use_mem_pool`` +
allocations), not just the constructor.

Construct a cuMem-backed ``torch.cuda.MemPool`` for this allocator,
with ``expandable_segments`` scoped off. The pool must
still be USED inside :meth:`expandable_segments_off` too — prefer
wrapping the whole cuMem lifecycle in that context.

Routes through :func:`arbi_serve.runtime.named_pool.create_mempool` —
the single MemPool construction chokepoint — so every cuMem-backed pool
is registered in the canonical registry at birth (no untracked pool can
exist). Never construct a raw MemPool directly here.

Move the mapped-bytes counter for ``base_tag`` by ``delta``.

Called on every physical map (malloc, wake) and unmap (free, sleep) so
``_tag_mapped`` tracks REAL mapped bytes for every tag, capped or not.
Caller holds ``_lock`` (or is the precheck, which does not mutate).

Hard-cap gate (called from C BEFORE reserve+map, under the GIL).

Return nonzero to DENY mapping ``size`` more bytes into the active
tag's pool — its budget would be exceeded. Only reached when a cap
is registered (the C ``g_any_cap`` gate skips it otherwise), so this
is never on the serving hot path. Minimal work (dict lookups + a small
record), no logging and no CUDA, to avoid re-entering the allocator
under the GIL — :meth:`cap_report` drains the breach records later.

Keyed on the UNQUALIFIED pool name (``_current_tag_base``), which is
the SAME key :meth:`set_tag_cap` writes. Keying this on the qualified
``_current_tag`` while the writer keyed the pool name is what made every
by-construction cap INERT (a cap is only a guarantee if writer and
reader agree on the key space).

Return ``(handle, size)`` so the C side can unmap+release+free.

If the region was already unmapped (discard-sleep) the handle is
0 but the size is still returned so the C side frees the VA.

Set (or clear, with ``None``) the mapped-bytes cap for ``tag``.

``tag`` is the UNQUALIFIED pool name (``capture.cudagraphs``,
``model.weights``, …). Caps are keyed on that name — deliberately NOT
on the namespace-qualified tag that :meth:`tag` stamps on allocations —
for three reasons:

  * every writer names a POOL, not a model instance;
  * a pool's allocations are qualified during BOOT (the build runs
    inside ``tag_namespace(model_key)``) but UNQUALIFIED during
    SERVING (the namespace context has exited), so a qualified key
    could not bind across the boot→serve handoff that the Phase-2
    freeze caps exist to guard; and
  * stable-VA residency keeps only ONE model VRAM-resident at a time,
    so a per-pool-name budget is the physically meaningful one.

A qualified name is accepted and normalised via :func:`base_tag_of`.

While ANY cap is registered, the C precheck gate is armed
(``set_any_cap(1)``) so an allocation that would push a tag past its
cap is denied at map time — a contained in-pool OOM, never a card OOM.


``enforce=False`` arms the cap in OBSERVE mode: breaches are recorded
and reported LOUDLY but the allocation proceeds. Use it for a cap VALUE
that has never been validated against a real boot — see :class:`TagCap`.
``origin`` names the arming site in the report. ``expect_bound=False``
marks a purely forward-looking budget whose tag may legitimately have no
allocations yet (it is then exempt from the INERT-cap alarm).

Lift every cap and disarm the C precheck gate (zero hot-path cost).

Breach records are DELIBERATELY kept — the report that explains what a
cap did (or failed to do) usually runs after the disarm. Use
:meth:`reset_cap_breaches` to clear them.

Book ``nbytes`` of ``parent``'s mapped bytes to the sub-tag ``subtag``.

For a slab fold: one physical segment is mapped under ``parent`` and
sliced into ``narrow()`` views, each owned by a logical consumer that
used to own a pool of its own. Charging keeps that consumer's ledger row
and its cap key alive without mapping physical for it.

Re-charging the same sub-tag REPLACES its charge. ``parent`` may be
given qualified or unqualified; the charge is keyed on the unqualified
names, the key space caps and counters live in.

``raw`` bytes for tag ``base``, net of the fold partition.

A sub-tag reports exactly its charge; a host reports what it maps minus
everything carved out of it.

The partition applies only while the host actually maps at least what
has been carved from it. A charge names a slice of MAPPED physical, so
once the host has less than that — asleep, or its slab released — the
slice is not backed by anything and the host's raw bytes are the whole
truth again. Caller holds ``_lock``.

Bytes currently mapped under the pool named ``tag``.

Maintained for EVERY tag (capped or not) and kept in step with every
map/unmap, so this equals ``mapped_bytes(tag)`` at all times. Accepts a
qualified tag (normalised via :func:`base_tag_of`).

Partitioned by any live :meth:`charge_subtag`: a sub-tag reads its
carved bytes and its host reads the rest.

Mapped bytes per QUALIFIED tag, derived from the live allocations.

The honest per-tag picture (one row per ``<namespace>/<pool>``) for the
boot log. Independent of the ``_tag_mapped`` running counter, so a
disagreement between the two is itself a detectable defect — see
:meth:`counter_drift`.

Sub-allocation charges (:meth:`charge_subtag`) are applied as a
PARTITION of their host row, so a slab fold keeps one row per logical
consumer and the column still sums to the mapped physical.

Base tags where the running counter disagrees with the live scan.

Returns ``{base_tag: (counter_bytes, scanned_bytes)}`` for every tag
whose ``_tag_mapped`` counter differs from the sum over live mapped
allocations. MUST be empty: the counter is what the cap gate consults,
so drift means the gate is budgeting against a fiction.

RAW on both sides — sub-allocation charges (:meth:`charge_subtag`) are a
read-side partition of bytes that are already mapped and already counted
under their host tag, so applying them here would manufacture drift.

One row per armed cap: limit, live usage, posture, and breaches.

The anti-inertness surface. ``CapRow.inert`` is True when the capped
tag has never been allocated — the cap is protecting nothing.

Reserve ``size`` bytes of virtual address space (NO physical).

Returns the base VA. ``size`` should be a :meth:`granularity`
multiple. The range is unbacked until :meth:`map_range`; reading /
writing it before mapping faults. Release with :meth:`free_va` after
unmapping every mapped sub-range.

Reserve ``size`` bytes of VA at the FIXED base ``addr`` (no physical).

Multi-model stable-VA residency primitive: each model's arena is placed
at a deterministic, disjoint base so its captured graphs' baked
addresses never collide with another model's. Returns the granted base
which is GUARANTEED to equal ``addr`` (the primitive frees + fails the
reservation if the driver grants a different base, so an arena is always
exactly where it was placed). Returns 0 on failure so the caller can
retry at another base. ``addr=0`` is identical to :meth:`reserve_va`
(driver picks). ``align`` defaults to 0 (driver default); pass the device
:meth:`granularity` for clean arena boundaries.

Map physical for the sub-range ``[va+offset, va+offset+nbytes)``.

Creates a handle, maps it at the offset VA, sets RW access. Returns
the opaque allocation handle (pass it back to :meth:`unmap_range`).
``offset`` and ``nbytes`` must be :meth:`granularity` multiples.

Stamp every allocation made in this context with ``name``.

When a tag NAMESPACE is active (see :meth:`tag_namespace`), the
effective tag is ``"<namespace>/<name>"`` so two models that both
allocate into a pool called ``weights_pool`` get DISTINCT tags
(``modelA/weights_pool`` vs ``modelB/weights_pool``). That is the
primitive multi-model stable-VA residency needs: :meth:`sleep` /
:meth:`wake` can then target ONE model's allocations by its tag
prefix while leaving the other model's allocations untouched.

The UNQUALIFIED name is tracked alongside (``_current_tag_base``): it is
the key that caps and the per-tag byte counters use, so a cap armed by
pool name binds whether or not a namespace happens to be active.

Prefix every :meth:`tag` made inside this context with ``namespace/``.

Multi-model residency wraps a model's whole boot (pool fills +
capture sweep) in ``with alloc.tag_namespace(model_key):`` so every
allocation that model makes is tagged ``"<model_key>/<pool_name>"``.
A later :meth:`sleep_namespace` / :meth:`wake_namespace` then
offloads / restores exactly that model's physical pages, keeping
every OTHER model's VA reservations alive (so the driver hands the
next model DISJOINT VAs — the disjoint-arena guarantee, achieved by
construction without a fixed-base sub-allocator). ``None`` clears any
active namespace (nested-safe; restores the prior namespace on exit).

Public form of :meth:`_qualify` — the tag a RAW pool name is stamped
with right now.

Any accounting read that looks a pool up BY NAME must go through this.
:meth:`tag` stamps ``_current_tag = _qualify(name)``, so under an active
:meth:`tag_namespace` the live allocations carry ``"<ns>/<name>"`` while
a caller holding only the raw ``"capture.cudagraphs"`` looks up a tag
that no allocation ever had — and silently reads 0. That exact miss made
:meth:`arbi_serve.engine.engine_boot._EngineBootMixin.total_cudagraph_bytes`
fall through to a double-counted fallback for the whole life of the
stable-VA residency path.

Pure string work — safe to call on any path.

The namespace currently prefixing tags, or ``None``.

Read by ``build._stable_va_namespace`` so a runtime pool-member
build that ALREADY runs under its member key's namespace (the
``build_member_into_engine`` wrapper) is not silently re-tagged
under the boot model's served-name — a same-served-name config
variant would otherwise collide with the boot member's tags and a
targeted ``sleep_namespace(member_key)`` would match nothing.

Unmap physical for every alloc tagged under ``namespace`` (keep VA).

Mirror of :meth:`sleep` scoped to one model's tag prefix. ``offload``
copies each region D->pinned-host first (restored on
:meth:`wake_namespace`); ``offload=False`` discards. Returns bytes
released. The VA reservations stay put, so the next model the driver
serves gets DISJOINT addresses and this model's captured graphs'
baked ``data_ptr``s remain valid (just unbacked) until wake.

Unmap + release physical for every alloc under ``tags``, for good.

The same driver work as a discard :meth:`sleep` and the opposite
lifecycle intent. A park keeps the pool: its VAs stay meaningful, a
wake owes them physical back. A TEARDOWN destroys the pool — a boot
reclaim (``NamedPoolRegistry.release_empty_pool``), a rolled-back
member — and its VAs are dead. The allocation records outlive the
teardown (torch's segment records for a destroyed ``MemPool`` do too,
and the free hook is what finally drops them), so without this
distinction the next wake finds them unmapped, cannot tell a corpse
from a sleeper, and maps fresh physical into every one — bytes that
nothing will ever read and that the rest of the wake needs.

Returns the physical bytes released. The VA reservations stay put and
free through the allocator's free hook when torch drops the segment.

``(tag, bytes, count)`` for allocations under ``tags`` that a wake
declines to map.

A CORPSE is ``mapped=False`` and ``asleep=False``: :meth:`release_tags`
marks a destroyed pool's records that way and :meth:`wake` skips them
by contract, since a dead pool's VA must not draw fresh physical. The
VA stays reserved either way, so a tensor that still holds one of these
addresses faults on its first touch and the allocator raises nothing.

A woken member asks this about its OWN tags. A corpse under a tag it
still owns means a release reached a pool this member needs, and the
member is about to serve on unbacked VA.

:meth:`release_tags` for every alloc under one model's tag prefix.

The teardown counterpart of :meth:`sleep_namespace` — for an orphan
whose snapshot is being dropped, where ``sleep_namespace(offload=False)``
would leave records a later wake resurrects.

Unmap physical pages for every alloc whose tag is targeted.

Offloaded tags: D→pinned-host copy into ``cpu_backup`` first, then
unmap+release (restored on :meth:`wake`). Discarded tags:
unmap+release only (contents lost). Untargeted tags are left
mapped. Returns the total physical bytes released.

VA reservations are kept for every alloc, so ``data_ptr()`` of
any held tensor is unchanged across the cycle.

``park=False`` marks the unmapped allocs as RELEASED rather than
asleep, so :meth:`wake` leaves them alone — see :meth:`release_tags`,
which is the name teardown callers should use.

Re-map physical pages at the SAME VA for unmapped allocs.

``tags=None`` wakes every unmapped alloc; otherwise only allocs
whose tag is in ``tags``. Offloaded contents are restored H→D;
discarded contents come back zero/garbage (mapped but not
refilled). Asserts the remapped VA equals the original. Returns
total physical bytes mapped back.

Whether ``a`` belongs to ``tag``, given either tag form.

An UNQUALIFIED pool name matches the pool across every namespace (the
pool-name key space caps and counters use); a QUALIFIED tag matches
exactly one model's instance of that pool. Callers that name a pool
(the freeze, the serving pool table, the boot ledger) get the pool;
the residency controller, which holds ``"<model>/<pool>"``, still gets
its one model. Before this, an unqualified name matched NOTHING once a
namespace was active — which silently reported every pool as 0 bytes
and froze them at a 0-byte cap.

Physical bytes currently mapped (optionally for one tag).

``tag`` may be an unqualified pool name or a namespace-qualified tag —
see :meth:`_tag_matches`. A named tag is partitioned by any live
:meth:`charge_subtag`; the device-wide total (``tag=None``) is not,
since a charge moves bytes between rows and never adds or removes any.

Ensure at least ``target_bytes`` (rounded up) of the prefix is
physically backed, mapping one new chunk for the shortfall. Idempotent
when already mapped to/past the target. Returns the new mapped prefix
length. Raises if the target exceeds the reserved max.

Back the sub-range ``[offset, offset+nbytes)`` with fresh physical.

``offset`` and ``nbytes`` are device-granularity multiples and the
range must not overlap a chunk this region already mapped — the
driver refuses a second mapping at the same VA. Returns the bytes
mapped. The prefix form is :meth:`map_to`; this one is what a slab
whose backing advances at several independent offsets uses.

Zero-copy uint8 view that SPANS the reserved VA up to ``nbytes``
(default: the full ``max_bytes``), regardless of how much physical is
currently mapped (Option B / B2: the KV slab tensor must span the full
max-pages VA so block_table indices are valid addresses, while only the
mapped prefix is physically backed).

CAUTION: reading / writing an offset beyond ``mapped_bytes`` FAULTS the
device. The caller is responsible for ensuring every offset it actually
touches lies within the mapped prefix (the capture sweep indexes only
low pages; serving only touches a page after :meth:`map_to` has backed
it). Used to back a KV slab tensor whose VA is stable across grows.

Drop physical pages for every mapped chunk, KEEPING the VA + the
chunk layout, so a captured graph's baked ``base_va + offset`` pointers
stay valid and :meth:`wake` can remap the IDENTICAL prefix.

``offload=True`` copies each chunk D→pinned-host first (restored on
:meth:`wake`); ``offload=False`` discards (woken pages come back
zero/garbage). Returns the physical bytes released. Idempotent — a
no-op (returns 0) when already asleep.

This is the GrowableRegion analogue of
:meth:`CuMemPoolAllocator.sleep`: the growable KV slab lives in its OWN
VA reservation (NOT under a pluggable-allocator tag), so the pool-tag
sleep path can't see it — this method registers it explicitly into the
sleep/wake lifecycle.

Remap fresh physical at the SAME VA for every slept chunk and
restore offloaded contents. Mirror of :meth:`sleep`. Returns bytes
mapped back. No-op (returns 0) when not asleep.

Always a FULL remap of the identical chunk layout — a
content-preserving wake must restore every offloaded chunk at its
original offset. A DISCARDED region that must come back under a
tighter budget is NOT a wake (there is no content to restore): use
:meth:`reprovision`, which maps a fresh exact-sized prefix instead of
replaying the prior chunk layout.

Re-provision a DISCARDED, slept region with fresh physical for the
prefix ``[0, target_bytes)`` (granularity-rounded up), dropping the
stale chunk records. Returns the mapped byte count.

A discard-park keeps the chunk RECORDS but frees their physical;
re-activating the member must actually re-provision the backing,
sized to what fits NOW, not replay the prior chunk layout. Content
semantics are explicit: the region comes back UNINITIALIZED (the
caller zeroes what it exposes); there is nothing to restore because
the park discarded it.

Fail-loud contract:
  * refuses a region that is not fully asleep (a mapped chunk means
    this is not a discarded-park wake — programming error),
  * refuses if ANY chunk still carries a pinned-host backup (content-
    bearing regions must go through :meth:`wake`, never lose data),
  * a driver refusal in ``map_range`` raises — the caller must never
    be told "mapped" when the driver said no.

Low-level cuMem virtual-memory primitive (single owner of the driver calls).

One place owns the CUDA virtual-memory driver API: reserve a VA range,
back a range with fresh physical pages (create + map + set access), drop
the physical pages (unmap + release) while keeping the VA, free a VA
range, and query the device allocation granularity. Both higher-level
Python allocators build on this:

  - :class:`arbi_serve.runtime.sleep_allocator.StableVAReserver`
    (model-param / KV-slab sleep-wake), and
  - :class:`arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator` /
    :class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion`
    (the reserve-max-VA + partial-map + grow path).

The driver calls go through ``cuda.bindings.driver`` (cuda-python). The
``CuMemPoolAllocator`` torch-pluggable callback path (``my_malloc`` /
``my_free``) stays in its C shim because the CUDA caching allocator calls
that hook from C; every Python-initiated reserve/map/unmap routes here.

A physical-pages handle is the opaque ``cuda-python`` allocation handle.
``None`` means "no physical backing" (a reserved-but-unmapped VA).

This module imports cleanly on CPU-only hosts: :func:`driver_available`
returns ``False`` and every driver call raises a clear error rather than
crashing at import.

Device cuMem allocation granularity (minimum size / alignment).

Every VA size, map offset and map length must be a multiple of this.
Raises on driver failure or a zero result rather than guessing a
page size.

Bind ``device``'s primary context if no context is current.

cuMem* needs a current context; a callback or worker thread may have
none. Retains the primary context only when one is missing (so the
retain refcount is not bumped on every call), mirroring the C shim.

Reserve ``size`` bytes of virtual address space (no physical backing).

Returns the base VA. ``size`` must be a granularity multiple. The
range is unbacked until :func:`create_and_map`; touching it before
mapping faults.

``addr`` requests a FIXED base (multi-model residency places each
arena at a deterministic, disjoint base so captured graphs' baked
addresses never collide). The driver treats ``addr`` as a hint; when
a non-zero ``addr`` is requested but the driver grants a different
base the reservation is freed and ``0`` is returned so the caller can
retry elsewhere. ``addr=0`` lets the driver pick (any grant accepted).
``align`` is the reservation alignment (0 = driver default; pass the
device granularity for clean arena boundaries).

Back ``[va, va+size)`` with fresh physical pages and grant RW access.

Runs the cuMemCreate -> cuMemMap -> cuMemSetAccess sequence and
returns the opaque allocation handle (pass it back to
:func:`unmap_and_release`). ``size`` must be a granularity multiple.
Binds ``device``'s primary context first if none is current.

Drop the physical pages backing ``[va, va+size)``; KEEP the VA.

Unmaps the range and releases ``handle`` (the value returned by
:func:`create_and_map`). The VA reservation lives on, so any tensor
whose ``data_ptr`` points into the range keeps a stable address.

Map an EXISTING physical handle at ``va`` and grant RW access.

The aliasing primitive: the same handle mapped at several addresses
makes those addresses resolve to the same physical pages. ``offset``
selects the byte offset INTO the handle's allocation. ``va``, ``size``
and ``offset`` are granularity multiples.

The counterpart is :func:`unmap` — the handle outlives every mapping
of it and is released once, separately, by :func:`release`.

Unmap ``[va, va+size)``, keeping the VA and the physical handle.

The range must be exactly one previous map's range — the driver
refuses to unmap a SUB-range of a mapping, so an arena that means to
drop pages one granule at a time must have mapped them that way.

Decode-cudagraph pad-up replay.

When a plain-decode batch of size ``B`` has no exact captured decode
cudagraph (``B`` falls between the sparse ``cudagraph_shapes`` ladder
rungs — e.g. 5/6/7 between 4 and 8), the engine would otherwise run the
eager / piecewise forward (~hundreds of kernel launches/step,
host-bound). This module pads the live batch UP to the smallest captured
shape ``B'`` and replays that graph, discarding the ``B'-B`` padding
rows' logits.

Correctness contract for the real ``[:B]`` rows:

  Decode is per-row independent (no cross-row reduction in the
  attention / recurrent / norm kernels), so the real rows' MATH is
  faithful — a padding row cannot change what row ``i < B`` computes.
  The one coupling is the GEMM batch dimension ``M`` (``B`` ⇒ ``B'``),
  which changes the cuBLAS split-K reduction ORDER; in bf16 that
  perturbs the last bits and can flip an argmax on a near-tie. This is
  the SAME intrinsic nondeterminism the engine already exhibits between
  any two captured batch sizes (e.g. genuine ``B=6`` eager vs genuine
  ``B=8`` cudagraph diverge on near-ties with the flag OFF) — it is a
  property of bf16 GEMM, not a defect of padding, and is unavoidable
  for any pad-up scheme. Byte-identical OFF-vs-ON is therefore not a
  reachable bar in bf16; faithful per-row math is.

  * Padding rows borrow FREE recurrent slab rows (rows on the pool's
    ``_free_rows`` stack — not owned by any in-flight request). The FLA
    fused-recurrent kernel reads/writes ``slab[state_indices[i]]``
    per-row independently, so a padding row writing garbage into a free
    slab row cannot touch a real row's state. When that free row is
    later claimed by a real request, ``alloc_for_request`` queues a
    zero-clear flushed before any kernel reads it — so no stale state
    leaks. ``B' <= max_batch`` guarantees ``max_batch - B >= B' - B``
    free rows exist.
  * Padding rows' ``slot_mapping`` points at the page-0 null slot, so
    the KV scatter writes harmlessly into the reserved null page.
  * Padding rows' ``block_table`` is all-zero (page 0); their attention
    reads the null page. Padding ``seq_lens = 1`` so the per-row K range
    is a single (null) token. Padding logits are never read.
  * The captured graph's per-row ops are independent across the batch
    dimension (no cross-row reduction in decode), so the real rows'
    logits are bit-identical to the exact-``B`` replay.

The whole path is gated behind ``cfg.decode_pad_cudagraph`` (env
``ARBI_DECODE_PAD_CUDAGRAPH``); default ON. Padded rows read the null page
and their logits are never read, so the real rows' logits are bit-identical
to the exact-``B`` replay. Set ``ARBI_DECODE_PAD_CUDAGRAPH=0`` to force the
legacy eager-on-stride-miss path.

Return ``n_pad`` slab rows for padding that do NOT collide with
``real_rows`` (the in-flight requests' rows for this kind).

Prefers the pool's ``_free_rows`` (rows owned by no request). Falls
back to the permanently-zero sentinel row 0 if the free list is
somehow short (defensive; ``B' <= max_batch`` makes this
unreachable). We do NOT pop the free list — a borrow is transient
for this single replay; the row stays free for real admission.

Replay ``captured`` (a ``B'``-row decode graph) for a live
``B``-row decode ``batch`` padded to ``B'``; return ``logits[:B]``.

``b_prime`` is the captured graph's batch size; ``B = batch.num_seqs``
and ``n_pad = b_prime - B``. All replay-input tensors are extended to
``B'`` rows with benign scratch values (see module docstring).

Replay a ``(B', S)`` verify cudagraph for a live ``(B, S)`` verify
``batch`` padded to ``B'``; return ``logits[:B*S]`` (and, when
requested, ``hidden[:B*S]``).

The verify pass is flat ``B × S`` with ``S = K + 1`` tokens per row
and the captured graph bakes ``cu_seqlens_q = [0, S, ..., B'*S]``
(the one input ``CapturedGraph.replay`` does NOT re-copy). So each
of the ``n_pad = B' - B`` padding rows must contribute EXACTLY ``S``
flat tokens — the padded flat length is ``B' * S`` and the implicit
``cu_seqlens_q`` lines up with the captured graph automatically.

Correctness contract (identical reasoning to :func:`replay_padded`,
one token-per-row generalised to ``S`` tokens-per-row):

  * The verify forward is per-row independent — attention is
    per-sequence (each row's ``S`` queries attend only within that
    row's K range via ``seqused_k`` / ``cu_seqlens_k``); there is
    no cross-row reduction. So a padding row cannot change what a
    real row ``i < B`` computes. The only coupling is the GEMM
    batch dim ``M`` (``B*S ⇒ B'*S``), whose cuBLAS split-K
    reduction order shifts the last bf16 bits — the SAME intrinsic
    nondeterminism :func:`replay_padded` documents for plain decode,
    not corruption. Gated behind ``cfg.decode_pad_cudagraph``.
  * Padding rows' ``slot_mapping`` ⇒ page-0 null slot; their KV
    scatter writes harmlessly into the reserved null page.
  * Padding rows' ``block_table`` all-zero (page 0); ``seq_lens =
    1`` so each reads a single null token. Padding logits/hidden
    are sliced off before return.
  * Padding rows borrow FREE recurrent slab rows (when the arch is
    hybrid) so a padding row's recurrent write cannot touch a real
    row's state — same borrow logic as :func:`replay_padded`.

Whether a device-scoped CUDA API may be called for a given device.

``torch.cuda.is_available()`` answers a PROCESS-level question — "was torch
built with CUDA and is a driver present". The device-scoped APIs
(``synchronize``, ``reset_peak_memory_stats``, ``memory_allocated``,
``get_device_capability``) answer a DEVICE-level one, and reject a non-CUDA
device by raising. Gating the second on the first is therefore wrong on any
box that has a GPU while the code holds a CPU device: the guard passes and the
call raises. :func:`cuda_device_active` asks both halves.

Shared config-fingerprint hashing for the boot-manifest.

Caches key on "the same engine configuration produces the same artifact":
the :mod:`~arbi_serve.runtime.boot_manifest` activation profile needs a
STABLE, ORDER-INDEPENDENT key derived from a dict of JSON-serializable
config inputs. This module owns that hashing primitive so callers can't
drift on the canonical form (key-sorted, compact-separator JSON → SHA256).

SHA256 of the key-sorted JSON dump of ``fingerprint``, truncated.

Key-sorted + compact separators so two boots at the same configuration
produce a byte-identical canonical form and hence an identical key. The
fingerprint values must be JSON-serializable.

``length`` is the number of leading hex chars to keep (16 → 64 bits,
collision-safe for the handful of configs a single host ever caches).

The ONE place a ``model.forward`` declares what kind of step it is.

Three independent channels ride on a forward — the EXL3 prefill numerics
class, the large-M accumulator's PHASE, and the step CLASS (with the slate
width) the int8 GEMM routes on — and each one selects a different KERNEL
for the same linear. That makes the declaration part of what the forward
COMPUTES, not telemetry about it.

Why this module exists rather than the declaration living at the serving
call site: it used to live there, inline in
:func:`arbi_serve.runtime.forward_exec.run_model_forward`, which is one of
twenty-six ``model.forward`` sites in the tree. The other twenty-five —
every cudagraph capture site, the piecewise sweep, the compile warmup, the
boot readiness gate, the drafter — called ``model.forward`` directly and so
declared NOTHING. Two consequences, and the second is the worse one:

* a captured graph RECORDS the kernel that was launched while it was being
  captured, so an undeclared capture bakes the wrong leg into the graph and
  every replay serves it — the flag reads armed and production runs the
  other kernel, permanently;
* :data:`~arbi_serve.weight_quant.exl3.custom_op.INT8_PREFILL_GEMM` could
  not tell "the leg declined" from "the leg was never asked", because the
  op short-circuits on the undeclared module global BEFORE reaching the
  refusal ladder. A counter that reads zero for two different reasons
  cannot report a negative.

So the declaration is a seam, and this module is the seam. Call
:func:`forward_declared` instead of ``model.forward``;
``tests/test_forward_declaration_seam.py`` asserts no site reintroduces the
bare call. Adding a fourth channel means editing
:func:`declare_forward` once.

COST WHEN NOTHING IS ARMED, which is every deployment today and every
non-EXL3 model always: three ``runtime_flags()`` attribute reads and no
import of :mod:`arbi_serve.weight_quant.exl3` at all. That package's
``__init__`` registers the EXL3 backend, so importing it from a hot path
would change what a non-EXL3 boot loads. Each channel resolves to a
``nullcontext`` and the op reads one module global that is ``None``.

Declare the enclosed forward's numerics class, phase and step class.

``batch`` is the :class:`~arbi_serve.engine.batch.ScheduledBatch` the
forward will run — the SAME object, not a description of it, so a
capture site and a served step that build the same batch declare the
same thing. Every channel is derived from it, never from a caller's
own opinion about what kind of step this is: an argument would be a
second source of truth and would drift from what the forward does.

Re-entrant and exception-safe: each channel restores its previous
value on exit, so a failed capture cannot leave the process computing
decode with a prefill accumulator.

Declare a DRAFTER's own forward. Always DECODE, and never a batch.

A speculative drafter runs quantized linears of its own — the DFlash
draft model's, and the target ``lm_head`` it borrows — but it is not a
``model.forward`` and it has no :class:`~arbi_serve.engine.batch.
ScheduledBatch` to derive a class from, so :func:`declare_forward`
cannot serve it. Without a declaration of its own the drafter reached
the EXL3 op with nothing declared, which is the one state
:func:`~arbi_serve.weight_quant.exl3.custom_op._undeclared_forward_reason`
reserves for a call site that skipped the seam. A detector that fires on
the healthy path cannot report the unhealthy one, so the drafter says
what it is instead.

DECODE, and that is a property of the drafter rather than a default: a
draft block proposes tokens one position at a time off a committed
prefix, which is the regime the trellis leg's accept-invariance rests on.
It is not prefill (no prompt is being consumed) and it is not a verify
slate (nothing is being re-checked against a target). So the class is not
a placeholder that a later reading might refine — it is the answer, and
``int8_serves_class`` refuses it by name.

ROUTING-INERT BY CONSTRUCTION, which is what makes it safe to add to a
served path: DECODE is off-class for the int8 leg, and an undeclared
forward already resolved to DECODE
(:func:`~arbi_serve.weight_quant.exl3.custom_op.active_step_class`). The
kernel every drafter linear launches is the one it launched before —
including the kernels a cudagraph RECORDS, which is the half that
mattered: a drafter capture used to record while the flag read armed and
nothing said which leg was baked in.

``rows=0`` because a drafter slate is not a verify slate and no rule
reads a drafter's width; passing its real row count would offer a number
to a rule that must not consult it.

Costs one flag read on an unarmed deployment, the same as
:func:`declare_forward`, and imports the EXL3 module on no boot that does
not already load it.

``model.forward(*args, **kwargs)`` with the step declared.

The batch is ``args[1]`` — every ``forward`` in the tree takes
``(input_ids, batch, pool, attn_ops, ...)``, so the seam can find it
without each caller naming it, and a model whose signature ever
diverges fails here rather than declaring the wrong thing.

Prefer this over entering :func:`declare_forward` by hand: the reason
the channels were missed for twenty-five call sites is that entering a
context manager is something a new call site must REMEMBER to do,
while calling the wrapper is something it cannot forget without the
seam test noticing.

Forward-call / metadata-build / sample execution helpers for
:class:`EagerModelRunner`.

These three collaborators are the inner per-step execution primitives
the runner's ``execute`` / ``forward`` orchestration calls into:
``run_model_forward`` (the single canonical ``model.forward`` call site
with its inference_mode + activation-arena wrapping), ``build_metadata``
(the per-kind metadata-builder loop), and ``sample`` (the sampler-chain
dispatch). They are module-level free functions taking the runner
explicitly; :class:`EagerModelRunner` keeps thin delegating method
wrappers (``runner._run_model_forward`` / ``runner._build_metadata`` /
``runner._sample``) for every call site (engine path, worker bridge,
MTP seed, tests).

Each function preserves the AST scan walks this module alongside ``model_runner.py``.

Publish ``batch.attn_meta`` onto every attn op's side-channel.

The TKV attention runs through the ``arbi_serve::tkv_attention``
custom op, whose Tensor-only signature cannot carry the rich
:class:`AttnPagedKVMeta` (TQRunState + the ``_tq_*`` mirrors the
AttendKernel wrappers introspect). Under ``torch.compile`` the
model's Python ``forward`` (which used to stash the meta on the op)
is traced away and the custom op is emitted directly into the
Inductor graph, so the op real-impl reads the meta off a
side-channel (``TkvAttnOp._call_batch_meta``) that the engine must
publish EAGERLY, OUTSIDE the compiled region.

This is the SINGLE canonical publish. ``run_model_forward`` (the live
hot path) calls it; so MUST every synthetic-forward site that drives
``model.forward`` directly without going through ``run_model_forward``
— the boot ``precapture_compile_warmup`` and the cudagraph
capture sweeps (prefill + decode). Skipping it there leaves the op
reading the PREVIOUS forward's stale meta: e.g. a 1-token warmup
forward inheriting the activation-profile's 17-seq / 528-token
``cu_seqlens_q`` makes the Turbo prefill kernel index the 1-row q
buffer as if it held 528 rows → ``cudaErrorIllegalAddress``.

No-op for ops that don't define ``publish_call_meta`` (bf16 carries
everything in the op signature) and when ``batch.attn_meta`` is None.
All PAGED_KV layers in a step share the one meta object.

Collapse a dual-logit-head model's ``forward()`` tuple down to the
single-tensor (or ``(logits, hidden)``) shape every generic post-
forward consumer — the cudagraph capture builders
(``runtime/capture/{decode,prefill,mixed}.py``) and the sampler
dispatch downstream of :func:`run_model_forward` — is written
against.

Almost every model's ``forward()`` returns ONE tensor (or, with
``return_hidden_state=True``, ``(logits, hidden)``) per
:class:`arbi_serve.models.base.ModelBase`'s contract. A model with
genuinely independent, untied logits heads sharing one hidden state
(currently only :class:`~arbi_serve.models.nemotron_voicechat.
NemotronVoiceChatBackboneModel` / ``NemotronVoiceChatModel`` —
``(text_logits, function_logits[, hidden])``) instead returns a
tuple of PER-HEAD logits, which the generic dispatch has no
business interpreting: it doesn't know which head to sample, and
blindly forwarding the raw tuple crashes downstream with
``AttributeError`` the first time anything calls ``.dim()`` /
``.shape`` / ``.data_ptr()`` on it (the actual bug this function
fixes — see its call sites).

Capability-checked via ``model.LOGITS_HEAD_NAMES`` (present + len>1
only on dual-head models) — never a hardcoded ``isinstance`` check,
mirroring the ``output_modalities``/``mm_bindings`` capability-check
convention elsewhere in the engine
(:func:`arbi_serve.multimodal.output.resolve_output_spec` /
:func:`arbi_serve.multimodal.registry.resolve_mm_bindings`). A
single-head model (the common case) has no such attribute, so this
is a no-op ``getattr`` check and an immediate pass-through — zero
behavior change, zero added cost on the hot path.

Picks ``result[0]`` (the FIRST-named head — ``"text"`` for
NemotronVoiceChat) as THE logits the generic engine samples;
every other head (e.g. the function-call head) is intentionally
dropped HERE, not silently further downstream — see
``NemotronVoiceChatBackboneModel.LOGITS_HEAD_NAMES``'s docstring for
why that's the right behavior for a generic caller (the real
multi-head consumer, the realtime turn loop's ``EngineSttStep``,
drives ``model.forward`` directly and never routes through here).

Give every row that opted in its OWN per-head logits, before
:func:`collapse_dual_head_logits` drops all but the first head.

The multi-head decode primitive (design doc §7.42). A request sets
:attr:`~arbi_serve.engine.request.Request.retain_logits_heads` to say
"I need the heads the generic dispatch is about to discard"; this
function writes them to that request's
:attr:`~arbi_serve.engine.request.Request.logits_heads` as a tuple of
``(vocab_size,)`` tensors ALREADY sliced to its own row. What
``run_model_forward`` returns to every generic consumer (sampler,
commit, capture) is untouched, so only the opted-in request observes
the extra head.

``reqs`` is the batch's per-row request list in slate order — the SAME
ordering :func:`resolve_embed_override` consumes and the same one the
logits tensor's rows carry, which is what makes ``enumerate`` the
correct row index here rather than a value a later consumer has to
re-derive.

Three no-op gates, cheapest first, so this costs a single ``if`` on
every ordinary forward: ``reqs is None`` (every call site that has not
threaded the slate through — all of them except the embed-override
path a duplex tick is already forced onto), a model with no
``LOGITS_HEAD_NAMES``, and a result that is not the expected per-head
tuple. The ``isinstance`` check on ``head_names`` is exact for the
same reason :func:`collapse_dual_head_logits`'s is — see its comment.

This replaces ``eng._dual_head_sink``, an engine-global dict armed and
disarmed around exactly one step, which in turn replaced a
process-global rebinding of ``run_model_forward`` itself. Both earlier
forms held ONE slot for what is per-request state; this one cannot be
read by the wrong request or left armed by a step that never ran.

This step's row-boundary offsets, host-side, without a device sync.

``batch.cu_seqlens_q`` lives on the device, so ``.tolist()`` on it is a
blocking D2H that stalls the host until the GPU drains every queued
kernel — issued here, between metadata-build and the forward launch,
which is exactly the stall
:meth:`~arbi_serve.backends.tkv_metadata_builder.TkvMetadataBuilder.build`
already refuses to pay for the same tensor (see its ``host_mirror``
branch and the reasoning quoted there). ``batch.host_mirror`` is the
pinned host ring the persistent build path H2D-copied FROM, so it is
bit-identical by construction, not a re-derivation.

Falls back to the device read on the fresh-alloc / CPU-stub path, where
no mirror exists (``cuda_graphs=False``, or a batch built without
:class:`PiecewiseBuffers`) — same values either way.

Collect + CLEAR this step's ``pending_embed_override`` tensors.

``reqs`` is the batch's per-row request list, in the SAME order as
``batch``'s rows (``batch.cu_seqlens_q``) — i.e. slate order. Returns
``None`` (a pure no-op, zero cost beyond one cheap ``any()`` scan)
unless at least one request in ``reqs`` has
:attr:`~arbi_serve.engine.request.Request.pending_embed_override` set —
which nothing except an opted-in NemotronVoiceChat-family caller ever
does. ``reqs is None`` (every call site that hasn't threaded the slate
through — the default) is the SAME no-op, so this function is a
complete no-op for every existing call site until one is explicitly
updated to pass ``reqs``.

When at least one row DOES have an override set, this function
enforces the Phase-1 contract (see
:attr:`Request.pending_embed_override`'s own docstring) and fails
LOUD, never silently wrong, on any violation:

  * the served model must declare ``SUPPORTS_EMBED_OVERRIDE = True``;
  * each override tensor's row count must exactly match the number of
    tokens that request contributes to this batch step.

A MIXED batch — some rows overridden, some not — is served when the
model declares an ``embed_input_ids`` hook: the un-overridden rows are
embedded from their own slice of ``batch.input_ids`` through that
hook, which is by construction the same lookup the model's ``forward``
would have done for them. Without the hook a mixed batch is still
refused loudly. The duplex lane relies on this: its seed-prefill and
context-injection rows deliberately carry no override while
steady-state frame rows in the SAME slate do (design doc §7.9.14). A
mixed batch whose un-overridden rows carry ``batch.mm`` is ALSO
refused loudly, because ``embed_input_ids`` does not merge media —
see the raise below.

On success, concatenates the per-row overrides into one
``(N_tokens, hidden_size)`` tensor in row order (matching
``batch.input_ids``'s own flat layout) and CLEARS every consumed
request's ``pending_embed_override`` back to ``None`` — the field
never silently persists across steps (see its own docstring).

Run ``model.forward`` with the canonical inference_mode +
activation-arena wrapping.

Single source of truth for the forward call. ``execute``,
``forward``, and the ``run_step.step`` MTP-seed branch
all call through here. The wrapping is REQUIRED for memory
correctness and for hidden-state lineage. Centralizing it means
future changes — notably CUDA-graph capture for the verify
path — touch one site.

Caller built the batch + ran ``_build_metadata``. This helper
does NOT touch the captured-graph pool; ``execute`` runs
captured-graph replay inline. ``forward`` paths (verify pass,
MTP-seed prefill) skip captured graphs because their shapes
aren't in the captured pool.

TP>1 hand-off: this helper does NOT issue the worker-bridge
broadcast itself. The three call sites (``execute``,
``forward``, ``run_step.step`` MTP-seed branch) issue their
own broadcast BEFORE calling here so workers rendezvous with
rank 0 on the same batch shape.

Per-pool routing: the body runs inside the engine's
``activation_arena`` named-pool ctx so per-step transient
``torch.empty`` calls inside the model (mamba conv buffers,
short-conv intermediates, attention scratch when the bump-
pointer arena is off) route into the named pool instead of
the default heap. The bump-pointer arena's allocations are
already in the same named pool (the arena's underlying buffer
was allocated inside ``with named_pool.use()`` at construct
time); double-wrapping is harmless — both end up in the same
bucket.

``reqs`` (optional, default ``None``): the batch's per-row request
list in slate order, ONLY needed by a caller that wants to honor a
request's :attr:`~arbi_serve.engine.request.Request
.pending_embed_override` — see :func:`resolve_embed_override`, which
this function delegates to. ``None`` (every call site that hasn't
been updated to pass it) is a complete no-op: the ``inputs_embeds``
kwarg is never added to the ``model.forward`` call, so this is
byte-identical to before ``reqs`` existed. Passing ``reqs`` is safe
even when no row has an override set — ``resolve_embed_override``'s
own no-op path is a single cheap ``any()`` scan.

``collapse_heads`` (default ``True``): whether to run the result
through :func:`collapse_dual_head_logits` before returning. Every
existing caller leaves this at the default, so this parameter is a
complete no-op for them (byte-identical return shape). A caller that
needs BOTH heads of a dual-head model (e.g. a NemotronVoiceChat
duplex-tick stepper reading ``function_logits`` for its own
frame-lockstep feedback — see
:mod:`arbi_serve.runtime.nemotron_voicechat_duplex_step`) passes
``collapse_heads=False`` to get the model's raw ``forward()`` return
(the untouched per-head tuple; a single-head model's result is
unaffected either way, since :func:`collapse_dual_head_logits` is
skipped entirely rather than replaced with different logic).

:attr:`~arbi_serve.engine.request.Request.retain_logits_heads`
(design doc §7.42): the OTHER way to reach both heads, for a caller
that does NOT own the forward call and therefore cannot pass
``collapse_heads=False`` — notably the duplex lane, whose step is
issued by ``run_forever``'s own loop, not by the duplex code. A
request that sets that field gets its own row's per-head logits on
:attr:`~arbi_serve.engine.request.Request.logits_heads` (see
:func:`stash_retained_logits_heads`) while this function still returns
the SAME collapsed value it always would, so every downstream consumer
is bit-identically unaffected. Requires ``reqs``, which the only path
a multi-head decode row can take — the embed-override eager forward —
already threads.

Run every active per-kind metadata builder and route the
result onto the correct ``*_meta`` slot via
:meth:`ScheduledBatch.set_meta_for`.

First-call cost detection (TKV autotune table miss + per-shape
buffer-pool ensure() growth) is attributed to the latest timed
request when ``capture_first_call=True`` — only the regular
step path enables this; the verify-pass forward never does
because its shapes are not part of the canonical graph pool.

Drain-before-build contract: ``GdnMetadataBuilder._finalize``
refuses a metadata build with recurrent zero-clears still pending,
so every build must sit downstream of a drain. Draining
HERE — the single seam every ``_build_metadata`` caller routes
through (plain eager fall-through, split_mixed sub-batches, verify
forward, MTP seed) — makes that hold for every entry path instead
of at N call sites. Zero-clears FIRST, savepoint resumes SECOND
(a later clear would wipe a restore). Eager-side only: no
capture/replay region runs this seam — replay skips metadata build
entirely and the capture sweep drains via
``cudagraph_admin.drain_recurrent_flushes_before_capture`` — so the
flush's H2D + indexed zero-write are never baked into a graph.

Score each row's sampled token onto ``req._pending_logprob``.

Reads the SAME ``(rows, vocab)`` logits the sampler just consumed, so
there is no second head pass. ``_logprobs_armed`` is False on every
deployment that did not arm the surface, which is where this returns.

Price ONE fused forward for a mixed MTP verify + prefill step
(``ARBI_MTP_FUSED_MIXED_PROBE``) — a MEASUREMENT arm, never a serving path.

WHAT IS BEING PRICED, and why it needed its own instrument
----------------------------------------------------------
Under MTP a step that carries both a prefill chunk and MTP-opted rows runs
TWO forwards: ``MtpStrategy.run_step`` partitions the slate and runs the
prefill half and the verify half separately, so the model weights are read
twice. Whether admitting both into ONE forward would pay is a subtraction of
two quantities, and neither is a constant of the engine:

  * what fusing DELETES — the whole verify forward, a full read of the
    weights at ``B x (K + 1)`` flat tokens;
  * what fusing COSTS — the marginal price of those ``B x (K + 1)`` tokens
    riding inside the prefill forward instead.

The A/B recorded at
:data:`~arbi_serve.runtime_flags.SPLIT_MIXED_MIN_DECODE_ROWS_DEFAULT` measures
the same subtraction for PLAIN single-token decode rows. Verify rows are a
different shape and a different kernel ladder, so that pair establishes the
ORDER of the prize and not the prize. This module measures the MTP-shaped
version directly.

WHY A PROBE AND NOT THE FUSED PATH
----------------------------------
Serving the fused step is blocked on three per-row-class dispatches that a
single forward cannot express as it stands, and each is a correctness
condition rather than a tuning choice:

  * GDN. :meth:`GatedDeltaNetBlock.forward` picks its kernel for the WHOLE
    batch. The verify rows need ``_forward_verify_fla``, which writes the
    per-token recurrent/conv SNAPSHOTS that ``pool.rollback_partial_accept``
    restores the accepted prefix from; the chunk kernel a fused batch selects
    writes no snapshot, so a partial accept has nothing to roll back to.
  * Attention. The tkv dispatch tests ``num_actual_tokens == num_seqs *
    mtp_block_m`` for the whole call, so a ragged fused batch routes the
    verify rows onto the varlen prefill kernel rather than the split-K verify
    kernel.
  * EXL3. The fused row count crosses ``auto_reconstruct_threshold``, moving
    the verify rows off the trellis leg onto reconstruct+hgemm.

The probe deliberately takes all three as they fall, because that makes its
number a LOWER BOUND on what a correct fused forward costs: a real one has to
add back a per-class GDN dispatch that this does not run. A lower bound is the
useful direction — a prize that is already absent at the lower bound is absent.

WHAT IT DOES TO THE ENGINE
--------------------------
Nothing that outlives the call, by construction rather than by care:

  * every fused row's ``slot_mapping`` is the page-0 NULL slot, so the
    probe's KV writes land where :mod:`arbi_serve.runtime.decode_pad`'s pad
    rows already write;
  * every fused row's recurrent ``state_indices`` is slab row 0, the pool's
    zero sentinel — no request owns it, and an admitted row is zero-cleared
    rather than seeded from it.

It reads the real ``block_table`` / ``seq_lens``, so the attention work is the
real work; only the writes are diverted.

It is still not free: the arm runs an EXTRA whole forward on every mixed step,
so its own TTFT/TPOT are not a serving measurement and must not be read as
one. What the arm produces is the per-shape device time in
:class:`~arbi_serve.runtime.step_mix_stats.StepMixStats`
(``/v1/admin/capture_hist``), beside the ``eager_verify_*`` and
``eager_prefill_*`` shapes of the same run — three device times measured by
one instrument on one boot, which is the only form in which the subtraction
above means anything.

TOKEN BUDGET
------------
The fused batch is TRIMMED to the step width the boot activation profile
measured (:func:`~arbi_serve.runtime.activation_profile.reachable_step_tokens`)
by dropping the tail of the prefill row's span. That is not tidiness: the
profile's mixed shape is ``max_batch`` rows carrying that width with
SINGLE-TOKEN decode rows, so a fused MTP step at ``(B - P) x (K + 1) + P x
chunk`` tokens sits ABOVE the widest shape the serving floor is sized for, and
the recorded failure of running a mixed forward wider than its profile is a
step OOM inside the forward that surfaces as an illegal memory access and
poisons the CUDA context. The trim is reported in the shape key, so the
comparison arithmetic can add the dropped tokens back at the marginal
per-token rate the same run's prefill shapes give.

Host twin of ``batch.seq_lens``, or ``None`` when there is none.

Never a ``.cpu()`` of the device tensor: the probe runs between two
enqueued forwards, and a device read there parks the host for both.

Build the fused probe batch, or return the reason it cannot be built.

Layout is ``[verify rows | prefill rows]`` — the verify rows FIRST, the
same ordering :mod:`arbi_serve.runtime.mixed_capture` fixes for its
canonical mixed layout, so the one split index the EXL3 leg is told about
separates the two numerics classes.

Returns ``(batch, n_verify_flat, total_flat)``. The prefill span is
trimmed from its TAIL when the concatenation would exceed
``max_flat_tokens`` — see the module docstring.

A trimmed row keeps its full ``seq_lens`` entry, so it presents as a
chunked-prefill CONTINUATION: fewer queries than KV, which is a shape the
varlen kernels serve every step and which bounds the attention work at or
above the untrimmed row's. The queries then sit at the wrong offsets
within that KV extent — which is why this is a cost probe and its output
is discarded, not a forward whose values mean anything.

Time ONE fused-shape forward for this step. Output discarded.

Called only from :meth:`MtpStrategy.run_step`, only when the step really
was the mixed composition, and only under ``ARBI_MTP_FUSED_MIXED_PROBE``.
Every decline is attributed on the ``mtp_fused_mixed_probe`` counter: an
armed instrument that silently produced nothing is indistinguishable from
a workload that never presented the shape, and those are opposite
findings.

ONE forward for a mixed MTP verify + prefill step (``ARBI_MTP_FUSED_MIXED``).

WHAT THIS IS
------------
Under MTP a step that carries MTP-opted decode rows beside a prefill chunk
runs two forwards today: :meth:`MtpStrategy.run_step` partitions the slate,
the prefill half takes the seed forward and the verify half the verify
forward, and the model weights are read twice. This module serves that step
with ONE forward over the batch ``[verify rows | prefill rows]`` and then
hands each half's outputs to the post-forward tail it would have taken
anyway: the verify rows' per-token logits and hidden state go to the verify
accept / rollback / drafter-seed tail
(:mod:`arbi_serve.spec_decode.mtp_verify_async` or ``_sync``), the prefill
rows' last-token logits and hidden state go to the seed tail
(:func:`arbi_serve.engine.run_step._mtp_seed_post_forward`). Nothing after
the forward is new; what is new is that both tails read one forward.

It is the serving twin of :mod:`arbi_serve.runtime.fused_mixed_probe`, which
prices the same shape with every write diverted. Here every write is real:
the verify rows' KV lands in the slots the verify plan allocated, the
prefill rows' in the slots the batch build allocated, and each row's
recurrent state is addressed by its own slab row.

THE THREE PER-ROW-CLASS DISPATCHES, and where each is resolved
-------------------------------------------------------------
* GDN — the batch carries :attr:`ScheduledBatch.row_class_split`, the
  metadata builder copies it onto :class:`GDNMeta`, and the block runs one
  forward with one recurrence launch per arithmetic class
  (:meth:`GDNBlock._forward_mixed_rows_fla`): the verify rows take the
  masked-replay launch (raw inputs saved in-kernel, no state commit — what
  the accept path replays from), the prefill rows the chunk chain with its
  commit. Bit-identical to the two class-pure forwards below the
  projections; asserted by ``tools/gdn_row_class/bench_gdn_row_class.py``.
* EXL3 — :func:`~arbi_serve.runtime.capture.mixed._mixed_row_split_ctx`
  declares the split so leg B runs one ``hgemm`` per class at a fixed
  ``M``. The int8 prefill leg and the dense projections are NOT split by
  this: above the GDN block the verify rows' arithmetic is the prefill
  class's, which is why the bar for arming this is MTP accept rate and not
  ``torch.equal`` (see the flag's contract).
* Attention — the fused batch is a prefill batch, so the verify rows take
  the varlen paged kernel over their real KV instead of the split-K verify
  kernel. Correct (the causal offset is per row), a different kernel.

WHAT IS REFUSED, BY NAME
------------------------
Every decline lands on the ``mtp_fused_mixed_step`` counter with its
reason. The decisions that can be made from host facts alone are made
BEFORE either half is built, because neither build is idempotent: the
verify plan takes each row's tail slot through ``allocate_slots`` and the
prefill build takes the chunk's, and nothing gives a slot back. Once a
half is built the step is served here whatever happens — fused, or over
the two built halves through the same tails the two-forward path runs.

WIDTH
-----
The fused step must fit the step width the boot activation profile
measured; it is never widened (a wider step is an unprofiled shape, and
the recorded failure of running one is a step OOM inside the forward that
poisons the CUDA context). The scheduler makes room instead: with the flag
armed every MTP-opted decode row debits ``1 + K`` of the step's token
allowance rather than ``1`` (:meth:`Scheduler._decode_row_step_tokens`),
so the co-admitted chunk is narrower by the verify rows' width. A step
that would still not fit — a boot whose scheduler did not debit — is
refused here before anything is built.

What :func:`run_fused_mixed_step` did with the step.

``served`` — the step ran here (fused, or over the two built halves) and
the caller must not run it again. ``fused`` — it ran as one forward.
``served=False`` means nothing was built and the caller runs the
two-forward path as if this module did not exist.

The prefill half of the fused batch, as the seed tail's observer sees it.

:meth:`DFlashDrafter.observe_seed_forward` reads a batch's host
``cu_seqlens_q`` / ``positions`` twin to slice the tap slab, and the
tap holds the WHOLE fused forward — verify rows first. So the prefill
rows' boundaries are the fused boundaries from the split onward, NOT
rebased: ``feats[lo:hi]`` indexes the slab, and the slab starts at the
verify rows. Positions are the fused positions for the same reason.

The served ``[verify rows | prefill rows]`` batch, or the reason it cannot be.

Every write is REAL — this is the difference from
:func:`~arbi_serve.runtime.fused_mixed_probe.fused_probe_batch`, whose
layout this shares. ``slot_mapping`` is the two halves' own, so the
verify rows write the slots the verify plan allocated and the prefill
rows the slots the batch build did; ``state_indices`` is the two halves'
own slab rows, with ``prefill_state_rows`` supplying the prefill rows'
when their batch does not carry them.

Returns ``(batch, n_verify_flat, total_flat, host_cu_seqlens_q)`` — the
last is the host boundaries the batch was built from, returned beside it
because the batch's own twin is withheld on a capture-capable boot and
the seed observer needs the host values either way. Never trims: a step
that does not fit is refused, and the scheduler is what makes it fit.

Upper bound on the verify half's flat tokens, from host facts.

``K`` may still collapse lower at plan time (drafter miss); it can never
go higher than the row's own opted depth.

Why :meth:`GDNBlock._forward_mixed_rows_fla` would refuse, or ``None``.

Host-side reads of the same module state the block checks, made here so
the decision lands on the counter before any slot is allocated rather
than as an exception inside the forward.

Serve a mixed step with one forward when every precondition holds.

Called from :meth:`MtpStrategy.run_step` on a step that IS the mixed
composition, after the drains. Returns :data:`_DECLINED` — nothing
built, nothing run — on every host-fact refusal, so the caller's
two-forward path runs exactly as before. Once a half is built the step
is served here in every case (see the module docstring).

LayerStack iterator.

A single generic iterator over per-architecture forward loops
(``qwen3.py``, ``qwen3_5.py``, ``nemotron_h.py``, ``lfm2.py``,
``deepseek_v3.py`` would each otherwise iterate layers their own way;
hybrid models would hardcode their ``if spec.kind == LayerKind.GDN``
dispatch). Models declare *the specs*, not the loop. The single
``if kind ==`` switch lives in exactly one place — the per-model
dispatcher that picks each block class — so the runtime iterator is
data-driven from there on. Each model file declares ``layer_specs``
only and hands the matching ``blocks`` list (one block instance per
spec) to :class:`LayerStack`. Under ``torch.compile`` this iterator is
the FX split boundary.

The dispatcher concern lives here as the optional
:paramref:`dispatch_wrapper` parameter — each per-arch block is a thin
shape-converter that pulls its kind-specific extras out of ``**extra``
and calls the inner decoder layer with its native argv, while
LayerStack threads every block call through one shared
``model_dispatch`` wrapper supplied by the per-model ``forward`` body.

Generic iterator over a list of :class:`LayerSpec` + matching
``nn.Module`` blocks.

The constructor pairs spec ``i`` with block ``i`` and stashes the
pair. :meth:`forward` walks the pairs and, per the spec's
:class:`StateKind`, looks up the per-state-kind metadata + view and
threads them into the block.

The dispatch is data-driven — there is **no** ``if kind == ...``
switch inside :meth:`forward`. Each block is responsible for its
own internals (PAGED_KV blocks know how to consume PAGED_KV
metadata, GDN blocks know how to consume GDN metadata, etc.). The
one ``if kind ==`` switch lives in the per-model dispatcher
that constructs the ``blocks`` list (e.g.
``_make_qwen3_5_decoder_layer`` in ``models/qwen3_5.py``); after
that, the runtime is uniform.

Block call signature (uniform across all kinds):

    block(
        x,                                  # (N_tokens, hidden_size)
        *,
        positions=...,                      # (N_tokens,) int
        batch_meta=...,                     # this kind's metadata
        state_view=...,                     # this kind's per-layer view
        **extra,                            # rope_cache, attn_op,
                                            # lora_state, arena, ...
    ) -> Tensor

Per-kind block signatures differ today (e.g. PAGED_KV
:class:`AttentionBlock` takes ``rope_cache + attn_op`` while
:class:`GDNBlock` takes neither). Each block adapter wraps its
decoder layer in a thin shape-converter that pulls its kind-specific
extras out of ``**extra`` and threads them into the inner argv.

Optional per-call dispatch wrapping
-----------------------------------
The piecewise CUDAGraph dispatcher (:func:`arbi_serve.runtime.capture.dispatch.model_dispatch`)
needs to wrap every per-layer call so the per-layer captured-graph
fast path replays for the live hot path AND the precapture sweep
can install graphs on miss. The per-model ``forward`` constructs a
single closure that captures the per-step context (``model``,
``num_seqs``, ``is_prefill``) and hands it to :class:`LayerStack`
via :paramref:`dispatch_wrapper`. The iterator
runs every block call through that wrapper, so every model gets the
piecewise fast path uniformly without per-arch boilerplate.

The wrapper signature is:

    dispatch_wrapper(
        layer_call: Callable[[Tensor], Tensor],
        *,
        hidden: Tensor,
        block: nn.Module,
        layer_idx: int,
        kind_name: str,
    ) -> Tensor

where ``layer_call(h)`` is the thunk that runs ``block(h, ...)``
with the resolved per-step kwargs and returns the new hidden. Most
callers forward directly into :func:`piecewise.model_dispatch`:

    dispatch_wrapper=lambda call, *, hidden, block, layer_idx, kind_name: (
        piecewise.model_dispatch(
            self,
            layer_idx=layer_idx,
            layer=block,
            hidden=hidden,
            num_seqs=batch.num_seqs,
            is_prefill=batch.is_prefill,
            layer_call=call,
        )
    )

The wrapper is optional: ``None`` means "call every block eagerly";
pure-CPU tests and CPU smoke paths leave it unset.

Args:
    layer_specs: per-layer architecture spec (length L). Models
        that carry MTP-marker specs (``is_mtp_layer=True``) MUST
        filter those out before constructing :class:`LayerStack`;
        the iterator does not know about MTP semantics.
    blocks: matching ``nn.Module`` blocks (length L). ``blocks[i]``
        is whatever block class the model defines for
        ``layer_specs[i].kind`` (``AttentionBlock`` / ``MlaBlock``
        / ``Mamba2Block`` / ``GdnBlock`` / ``ShortConvBlock`` /
        ``MlpOnlyBlock``). The per-model dispatcher decides the
        class; :class:`LayerStack` only sees ``nn.Module``.
    dispatch_wrapper: optional default wrapper threaded around
        every per-layer block call. ``None`` (default) runs every
        block eagerly. :meth:`forward` accepts a per-call override
        via the same-named kwarg.

Positional, closure-free fast path used by per-arch model
forwards on the compile-on hot path.

Block forward signature: ``block(hidden, positions, meta, view, *extras)``.
Each ``extras`` element is forwarded positionally; recurrent
blocks that ignore them must accept ``*_unused``.

``dispatch`` is an optional bound method on the model with
signature ``dispatch(layer_idx, block, hidden, positions, meta,
view, extras) -> hidden``. ``None`` runs every block eagerly.
Both branches are Dynamo-traceable: no closures constructed
inside the loop, no kwargs unpacking, no method-call on the
spec object.

Cross-layer (residual + input_layernorm) fusion. Per-arch
models that opt into the cross-layer fusion (qwen3_5 today)
thread the
model-owned persistent ``residual_buf`` through ``extras`` as
an additional positional. Captured layer-graphs bake the
residual_buf's ``data_ptr`` at capture time; the model
guarantees the slab is allocated in ``graph_buffers_pool`` and
never freed/reallocated, so every replay references the same
address. Per-step content is reset by an eager
``residual_buf.zero_()`` at the top of the model forward; the
``hidden`` (dispatched by the LayerStack) is the layer-N →
layer-N+1 ``pending_add``, not the running residual.

Map each layer's ``StateKind`` to its pre-resolved per-layer view.

``state_views`` is the per-:class:`StateKind` dict the per-arch
``forward`` builds (typically ``{StateKind.PAGED_KV: cache,
StateKind.GDN: cache, ...}`` where ``cache`` is the engine's
:class:`MultiStatePool`). The pool's pre-bound
:meth:`MultiStatePool.per_layer_views` is consulted lazily —
when present, the per-layer entry replaces the per-kind handle
in ``block_args[2]``; otherwise the kind handle passes through
unchanged.

The ``getattr`` lookup is the compatibility seam: tests that
build a :class:`LayerStack` against mock state-view objects
(``MagicMock``s, raw dicts) won't expose ``per_layer_views`` and
get the "per-kind handle" semantics. Production engines route
through :class:`MultiStatePool` and get the de-specialized fast
path.

Iterate every (spec, block) pair, threading per-kind
metadata + view to the block.

Args:
    x: ``(N_tokens, hidden_size)`` flat hidden tensor (the
        running residual).
    positions: ``(N_tokens,)`` int — per-token absolute
        position. Threaded uniformly because every existing
        attention block takes it; pure-recurrent blocks (GDN /
        Mamba / ShortConv) accept and ignore it via the
        uniform-signature adapter.
    batch_meta: per-:class:`StateKind` metadata produced by
        the :class:`MetadataBuilder`. The iterator looks
        up ``batch_meta[spec.state_kind()]`` for each layer.
        Missing kinds are passed as ``None`` — it's the
        block's job to validate it got what it needs.
    state_views: per-:class:`StateKind` slab view (or
        per-layer view; the per-model adapter decides). The
        iterator looks up ``state_views[spec.state_kind()]``
        for each layer; the block's adapter calls
        ``view.layer_view(spec.layer_idx)`` if it wants the
        per-layer slab.
    dispatch_wrapper: optional per-call override of the
        constructor-time wrapper. ``None`` (default) falls
        back to ``self.dispatch_wrapper``; if both are
        ``None`` the iterator calls every block eagerly. See
        the class docstring for the wrapper's signature.
    **extra: any additional per-step kwargs (``rope_cache``,
        ``lora_state``, ``arena``, etc.) threaded through to
        every block. Block adapters pull what they need; what
        they don't recognize is silently accepted via
        ``**kwargs`` (the per-model adapter convention).

Returns:
    ``(N_tokens, hidden_size)`` post-final-block hidden state.
    The model's final RMSNorm and lm_head live in per-model
    code, not in the iterator.

CUDA memory snapshot tooling.

Wraps PyTorch's memory-history APIs so admin endpoints (and tests)
can:

  - flip allocation-history recording on / off
    (``torch.cuda.memory._record_memory_history``);
  - dump the resulting pickled snapshot to a path on disk
    (``torch.cuda.memory._dump_snapshot``) for offline analysis with
    https://pytorch.org/memory_viz;
  - read a parsed live snapshot grouped by named pool
    (``torch.cuda.memory._snapshot``);
  - report the top-N largest live allocations with best-effort
    call-site information for triage when a request OOMs in
    production.

The PyTorch API names are leading-underscore-prefixed because the
team officially treats them as "private" — but they are the
documented memory-debugging interface and ship with every released
wheel. We wrap them here so every other module
(metrics, admin endpoints, tests) imports a stable wrapper instead
of touching ``torch.cuda.memory._*`` directly.

Turn on PyTorch's allocation-history recording.

The history is stored in a fixed-size ring buffer of
``max_entries`` events; older events fall off the back. 100k is
enough to capture a decent slice of a multi-second prefill while
keeping the JSON dump tractable.

Args:
    max_entries: ring-buffer size for the history.
    enabled: ``"all"`` records both alloc + free events.
        ``"alloc"`` skips frees. ``True`` is an alias for
        ``"all"``; ``False`` disables (caller should prefer
        :func:`stop_recording`).
    context: ``"all"`` includes a Python stack snapshot per
        event. ``"alloc"`` only stacks on alloc. ``"state"``
        disables stacks (smallest dump).
    stacks: ``"python"`` is the default; ``"all"`` captures C++
        frames too (heavier).

Write the recorded snapshot to ``out_path``; returns the path.

Default: ``/cache/memory-snapshots/<timestamp>.pkl``. The parent
directory is created if missing. The file format is whatever
``torch.cuda.memory._dump_snapshot`` writes (a pickled dict;
loadable by https://pytorch.org/memory_viz).

Read a live snapshot tolerating PT version drift.

Returns whatever ``torch.cuda.memory._snapshot`` returned (newer
PT: a dict with ``segments`` / ``device_traces``; older PT: a
flat list of segments). Callers should not rely on the shape;
use :func:`live_snapshot` for the parsed view.

Pick a stable label for a snapshot segment.

Four outcomes, in the order they are tested:

  * the segment's ``segment_pool_id`` matches a registered named pool →
    that pool's name;
  * the ``(0, 0)`` sentinel (or no id at all) → :data:`DEFAULT_POOL_LABEL`,
    the caching allocator's own pool. This is NOT unattributed: it is a
    distinct, named ledger row, and folding it in is what made the metric
    report live default-pool tensors as untagged bytes;
  * an id this process released (:meth:`~arbi_serve.runtime.
    named_pool_registry.NamedPoolRegistry.released_pool_ids`) →
    :data:`RELEASED_VA_LABEL`: VA whose physical is already back with the
    driver;
  * anything else → :data:`UNATTRIBUTED_LABEL`, the residual — resident
    bytes of ours in a pool nobody registered.

Group live snapshot segments by pool name.

Args:
    pool_id_to_name: map of ``MemPool.id`` (a tuple) to the
        named-pool label. The metrics callback builds this from
        the engine's :class:`NamedPoolRegistry`.
    released_pool_ids: ids of pools this process destroyed, from
        :meth:`~arbi_serve.runtime.named_pool_registry.NamedPoolRegistry.
        released_pool_ids`. Their surviving segment records are bucketed
        under ``address_space.released_pool_va`` rather than counted as residual.

Returns:
    ``{pool_name: {allocated_bytes, reserved_bytes, num_blocks,
    num_segments}}``. The three :data:`NON_POOL_LABELS` buckets are always
    present, even at zero, so the label set is stable across scrapes and a
    consumer never has to derive one of them by subtraction.
    segments: an allocator walk the caller already took. The walk IS the
        cost here; a caller that needs both this grouping and per-pool sums
        off the same instant takes one walk and passes it to both, which is
        also the only way the two views are guaranteed to describe the same
        allocator state rather than two states a few milliseconds apart.

Aggregate LIVE out-of-arena blocks by call-site.

The park co-residency assert names a byte TOTAL; this names WHO. Groups
every ``active_allocated`` block sitting in a segment that resolves to one
of the :data:`NON_POOL_LABELS` (i.e. NOT any registered named pool) by its nearest
allocation frame signature, so a GiB-scale out-of-arena residual is
attributed to the exact build/capture site that stranded it. Frames are
present only when history recording is on; without it entries group under
``"(no frames)"`` (still useful: count + total by size). Returns the top
``n`` sites by total bytes.

Return the top ``n`` live allocations by size.

Each entry: ``{size_bytes, address, pool_id, frames}`` where
``frames`` is the best-effort call-site (top-N Python frames the
history-recording option captured). When recording is OFF the
``frames`` list is empty — we still report sizes + addresses.

Flag-gated: start allocation-history recording for the SERVING phase.

NO-OP unless ``ARBI_SERVE_CUDA_MEMORY_HISTORY`` is set. This is the
telemetry-must-be-free contract in force: PyTorch's history recorder hooks
EVERY device alloc/free, so it is NOT free — it stays OFF by default and in
every benchmark arm, and is enabled only to pin a transient serving-tail
OOM. Called once per rank at the END of boot (post-capture, at the serving
transition — deliberately NOT during the compile sweep, whose host-side
recorder footprint OOMs the box), so the ring buffer fills with SERVING
allocations rather than boot allocations. Idempotent per process.

Best-effort SIGUSR2 → on-demand snapshot (sample a long flood, no OOM).

Only wired when serving-history recording is enabled, so it has ZERO
footprint on the default (flag-off) path. ``signal.signal`` works only on
the main thread of the main interpreter; if boot ran off-thread this
degrades to a warning (the dump-on-OOM path is unaffected). The handler
itself just calls :func:`dump_snapshot` — a plain file write, no allocation
on the hot path (it fires only when the operator sends the signal).

Best-effort dump-on-OOM snapshot. Returns the path, or ``None``.

Fires only when ``ARBI_SERVE_CUDA_MEMORY_HISTORY`` is on (otherwise the
recorder was never started, so a dump would carry no allocation frames —
a no-op keeps the OFF path truly zero-work). Writes a ``_dump_snapshot``
pickle carrying the recorded allocation history (with per-event Python
stacks, incl. the escaping alloc's call-site) to a DURABLE path under
:data:`DEFAULT_SNAPSHOT_DIR` (the ``/cache`` mount survives the container),
and logs the path LOUD so the operator finds it.

NEVER raises: the whole body is wrapped so the dump can never mask the
ORIGINAL OOM — the caller's fail-loud teardown proceeds unchanged whether
the dump succeeds or not. The filename is made unique WITHOUT depending on
the wall clock (rank + pid + a monotonic per-process seq); a human
timestamp is added only for findability.

MIXED decode+prefill whole-forward captured-graph replay
(``ARBI_MIXED_CAPTURE``).

Runtime half of the mixed-composition capture: takes the ALREADY
materialized mixed :class:`ScheduledBatch` (page-table advance, slot
allocation and recurrent-row resolution happened exactly once in
``build_batch`` — none of that is re-run here), permutes the live rows
into the captured graph's canonical layout, pads to the capture bucket,
replays ONE graph, and scatters the per-row logits back into slate
order.

Canonical layout (fixed by :func:`~arbi_serve.runtime.capture.mixed.
capture_mixed`): decode rows first (one token each, slate order), the
single prefill row LAST. Live steps with ``d_live < d_cap`` decode rows
pad with benign scratch rows — the :mod:`arbi_serve.runtime.decode_pad`
contract (page-0 null slot, ``seq_lens = 1``, borrowed free recurrent
slab rows) — and pad the prefill chunk's tail to ``bucket_n`` with
token 0 → the null slot. ``cu_seqlens_q`` is refreshed per replay to
``[0, 1, ..., d_cap, d_cap + real_n]`` so every varlen kernel bounds
its work by the LIVE lengths; pad logits are never read.

Numerics: greedy-argmax-identical to the FUSED eager forward — the
decode rows run the varlen prefill kernels (Turbo prefill varlen / GDN chunk)
instead of the decode kernels (split-K / fused-recurrent), and the GEMM
batch dim changes, both of which shift the last bf16 bits: the SAME
reduction-order noise class :mod:`split_mixed` and :mod:`decode_pad`
already document and ship default-on. On any lookup miss / unsupported
composition the caller falls through to the split / eager paths —
loud, never wrong.

THE FUSED FORWARD IS NOT THE SHIPPED BASELINE, though, and on a backend
whose GEMM leg is chosen by ROW COUNT that distinction is the whole
design problem. The sentence above compares this rung to a path
``split_mixed_decode_prefill`` (default ON) exists to replace; measured
against what actually runs, one fused forward moves the decode rows
from a decode-sized ``M`` to ``d_cap + bucket_n``. On EXL3 that crosses
``auto_reconstruct_threshold`` (280, or 64 for a narrow
``out_features``), so the decode rows switch from the trellis kernel to
reconstruct+cuBLAS-hgemm — a different ALGORITHM, not a reduction-order
shift, and precisely the "classes fuse and decode rows ride the prefill
row count" condition :mod:`arbi_serve.weight_quant.exl3.custom_op`
names as voiding the prefill-numerics pin.

WHAT CLOSES IT: a fixed-M row split, not a refusal.
:func:`~arbi_serve.runtime.capture.mixed.capture_mixed` declares the
canonical layout's split point to the EXL3 op
(``custom_op.mixed_row_split``), and leg B then runs TWO hgemms — at
``M = d_cap`` and ``M = bucket_n`` — against ONE reconstructed weight.
The reconstruct is the expensive half and the reason a second forward
costs anything, so the single weight pass survives; and because each
block's M is a constant of its own class rather than a function of what
else is in the slate, a decode row's output stops depending on the
chunk that travelled with it. That rests on leg B's measured property
that at fixed M an output row is invariant to its position in the block
and to its neighbours' values.

Pinning the hgemm M is NECESSARY BUT NOT SUFFICIENT, which measurement
established rather than review: leg B carries a second row-keyed
choice, ``_use_fused_reconstruct``'s variant gate, and with the M
already fixed a decode row still moved between a 512-token and a
2048-token chunk (10036/15360 elements, max delta 1.95e-3, measured
against a gate at 1024 rows, while 1024 vs 2048 agreed — the signature
of a step at that gate). The split therefore pins the variant too. Both
are needed; either alone leaves a decode row a function of the slate.

WHERE THAT PIN ACTUALLY BITES is narrower than the sentence above
suggests, because the gate is a function of the output width. On the
WIDE single-slice class its threshold is 0 — no row-keyed branch at all
— so the pin is a no-op on the linears that carry most of a chunk, and
what it buys is confined to the narrow and N-sliced classes, which keep
a real crossing.

WHAT IT DOES NOT BUY, and must not be read as buying: the decode rows
are on leg B here, where under ``split_mixed`` they are on leg A. The
split is STABLE, not identical — a decode row's value differs from its
``split_mixed`` value by leg B's deviation on the affected linears. On
the served 27B those include the 34 ``k_proj``/``v_proj`` that write the
KV a drafter reads, so the acceptance bar for this rung is MTP accept
rate, not output KL. The
:func:`~arbi_serve.engine.capture_admin.mixed.precapture_mixed_graphs`
sweep refuses only the one composition the split cannot express —
grouped MoE experts, whose rows arrive permuted into expert-major order.

One consequence worth stating because it is easy to assume otherwise:
``capture_mixed`` and :func:`try_mixed_replay` both call
``model.forward`` DIRECTLY rather than through
:func:`~arbi_serve.runtime.forward_exec.run_model_forward`, the single
site that marks the EXL3 numerics class and the prefill phase. So this
rung neither arms the prefill accumulator nor honours
``exl3_prefill_row_invariant``; the row split above is its own channel
precisely so it does not depend on that site.

Build the canonical-layout replay tensors for a mixed step.

Pure tensor staging (no engine state) so the layout contract is
CPU-testable. Returns the kwargs bundle for
:meth:`CapturedGraph.replay` MINUS ``recurrent_state_indices``
(which needs the live slab-row resolution — see
:func:`try_mixed_replay`).

Layout of the returned flat tensors (``flat_cap = d_cap + bucket_n``):

  ``[decode tokens (d_live) | decode pad (d_cap - d_live) |
    prefill tokens (real_n) | prefill pad (bucket_n - real_n)]``

and of the per-row tensors (``B = d_cap + 1`` rows):

  ``[decode rows (d_live) | pad rows (d_cap - d_live) | prefill row]``

``cu_seqlens_q = [0, 1, ..., d_cap, d_cap + real_n]`` — pad decode
rows still own exactly one flat token each (the layout is position-
fixed); only the LAST entry tracks the live chunk length, so the
captured last-token gather reads position ``d_cap + real_n - 1``
(the prefill row's true last token), and pad-tail positions past it
belong to no row.

Per-kind slab-row mapping in canonical order, pads borrowed.

Resolves each live row's slab row ONCE via
:func:`captured_lookup.resolve_recurrent_rows` (all recurrent kinds
share the mapping), reorders to the canonical layout, and fills the
``d_cap - d_live`` pad slots with borrowed FREE rows (the
:mod:`decode_pad` borrow contract — a pad row's recurrent write can
only dirty a row no request owns). ``None`` for pure-PAGED_KV
captures.

Replay the captured MIXED graph for this step, or ``None`` on miss.

Returns a ``(num_seqs, vocab)`` logits tensor row-aligned to
``slate`` (same contract as :func:`split_mixed.split_mixed_forward`)
so the sampler path is unchanged. ``None`` ⇒ the caller falls
through to the split / eager paths — the miss is never wrong, just
slower.

The caller has already materialized ``batch`` (page-table / slot /
recurrent-row resolution done — exactly once) and flushed the
per-step zero-clears / savepoint resumes.

Model execution runner — the "execute one slate" seam.

The run loop (:mod:`arbi_serve.engine.run_step`) owns *orchestration*
(timeout sweeping, MTP routing, commit / post-token, exception
handling); this runner owns *execution* (build flat tensors, run
per-kind metadata builders, captured-graph replay-or-live forward,
sample).

The forward-only seam :meth:`ModelRunner.forward` is what the MTP
verify pass calls — it builds nothing, it just runs metadata builders
+ ``model.forward`` against a caller-built :class:`ScheduledBatch` and
returns either the last-token logits or the full hidden state.

The runner does NOT touch :mod:`arbi_serve.engine.run_step`'s
post-step concerns: scheduler.commit, stop-token / length / context
checks, the per-token streaming consumer signal, or finish-state
metrics. Those stay in the run loop.

Output of :meth:`ModelRunner.execute`.

``sampled`` has one entry per slate row in slate order: the sampled
token id for "ready" requests (post-prompt or last-prompt-chunk),
``None`` for mid-prefill rows. ``timed_records`` carries the
``(request, StepRecord)`` pairs the runner opened so the run loop
can stamp commit-done after :meth:`Scheduler.commit` returns.

``plan`` is the :class:`StepPlan` the runner consumed, with the
flat tensor fields (``input_ids`` / ``positions`` / ``cu_seqlens``)
populated from the materialized :class:`ScheduledBatch`. Tests
consume this to verify the tensor-population invariant; production
code drops it.

Execute one scheduled slate.

Implementations differ on the executor surface (eager-mode,
captured-graph-only, distributed); the contract is the same:
consume a slate (or a pre-built :class:`StepPlan` carrying one),
return per-row sampled tokens.

Default single-process eager-mode runner.

Owns the per-step batch materialization + the captured-graph
eligibility check. Holds an ``Engine`` reference so it can read
the model / pool / attn_ops / metadata builders / captured-graph
pool / sampler — same access pattern as the rest of the run-loop
helpers. The engine remains the single source of those fields;
the runner is a *function* over them, not a parallel owner.

Materialize batch, run forward, sample. Pure compute path.

Accepts either a legacy slate ``[(req, n_tokens), ...]`` or a
:class:`StepPlan`. When given a plan, the runner populates the
plan's flat tensor fields (``input_ids`` / ``positions`` /
``cu_seqlens``) from the materialized batch so the populated
plan becomes the canonical input. The populated plan is
returned on :attr:`ExecuteResult.plan`.

``on_inputs_built`` (host-ahead pipeline, ARBI_ASYNC_SCHEDULE):
called once, after every request- / scheduler- / page-table READ
of this step is complete — see :meth:`EagerModelRunner.execute`.

Run per-kind metadata builders + model.forward on a caller-
built batch. Used by the MTP verify pass (it builds its own
batch shape ``[last_committed, draft_1, ..., draft_K]`` per
row and just needs the forward).

StepPlan-aware verify-forward entry point.

The MTP verify path emits a :class:`StepPlan` whose
``state_meta`` carries the full per-step routing header (see
:func:`spec_decode.mtp.build_verify_plan`). This method is the
public engine seam the verify driver dispatches through; it
reconstructs the per-step :class:`ScheduledBatch` from the plan
and invokes the inner ``forward(batch)`` transparently — there
is no staged ``ScheduledBatch`` payload on the plan.

Return a context manager that routes torch allocations into
the engine's ``activation_arena`` named pool.

Used as the per-step transient sink for forward intermediates
(mamba / short_conv / attention scratch / MLP), the sampler's
top-k / top-p masks, and the xgrammar bitmask H2D copy. These
allocations are stack-disciplined per step but not all of them
flow through the bump-pointer :class:`ActivationArena` — the
recurrent kernels in particular call ``torch.empty`` directly.
Wrapping the LIVE forward / sample / metadata-build paths in
the named pool's :meth:`use` ctx ensures every escaping
``torch.empty`` lands in ``activation_arena`` instead of the
default heap, which keeps the ``unattributed`` residual sub-

Returns :class:`contextlib.nullcontext` on CPU stub paths or
when the pool was never registered (the engine constructor
registers it unconditionally on CUDA, so this only fires in
unit-test fixtures that build the runner without a live
engine).

Arm (or clear, with ``None``) a per-step raw-logits tap.

``sink`` is called from :meth:`_finish_execute` with the exact
``(rows, vocab)`` logits tensor that step's forward produced —
BEFORE sampling — once per :meth:`execute` call, whether that
step replayed a captured cudagraph or ran eager. This is the
cudagraph-safe replacement for a ``lm_head`` forward hook: a
hook is a Python ``nn.Module.__call__`` callback, which fires
during EAGER forward and during graph CAPTURE (both real Python
calls) but never during graph REPLAY (only the recorded CUDA
kernels launch — no Python runs). ``logits`` at this chokepoint
is unaffected by that distinction: eager returns the model's
fresh tensor, and a captured replay already reads back its
persistent output buffer via a plain ``.clone()`` inside
:meth:`~arbi_serve.runtime.capture.decode_graph.CapturedGraph.replay`
before the value ever reaches here.

Calibration (:mod:`arbi_serve.calibration.engine_capture` /
:mod:`arbi_serve.calibration.engine_drift`) is the one consumer.
Every production deployment leaves this ``None`` — the tap in
:meth:`_finish_execute` is a single ``is not None`` check, at
the same cost class as the existing debug-only counters there.

Arm (or clear, with ``None``) a per-step raw-KV tap.

``sink`` is called from :meth:`_finish_execute` with the
:class:`~arbi_serve.engine.batch.ScheduledBatch` that step's
forward just ran — in particular its ``slot_mapping`` (device
tensor, one page/offset slot per token) — once per :meth:`execute`
call, whether that step replayed a captured cudagraph or ran
eager. The sink is expected to read the paged KV pool itself
(``eng.pool.layer_view(layer_idx)`` for each layer it cares
about, indexed at ``batch.slot_mapping``) to recover the K/V this
step just wrote.

This is the cudagraph-safe replacement for
ONLY way calibration's stage-1 capture recorded
K/V: a live Python method called from ``attn.py``'s eager
per-layer forward (via ``kv_seam.maybe_apply``). That call never
happens during a captured-graph REPLAY (only the recorded CUDA
kernels launch — no Python runs), so a chunk that happened to hit
a captured bucket silently contributed NO tokens to the capture.
Real-verified: this dropped exactly half the tokens fed to
centroid fitting under ``cuda_graphs=True`` on a real basket
(1024 of 2048), reproducibly — not noise, a real gap.

The paged KV pool is not a hook, unlike the old seam: it is a
persistent, stable-address buffer every attention kernel
scatters into as a normal side effect of running (baked into a
captured graph's kernel sequence the same as any other op that
writes to a persistent buffer, e.g. ``logits_out`` for
:meth:`set_raw_logits_sink`). Reading it back via
``batch.slot_mapping`` right after the step, on the same CUDA
stream the write issued on, sees the real data whether that step
replayed or ran eager — stream ordering (not an explicit sync)
is what makes this correct, exactly like the logits tap.

Calibration (:mod:`arbi_serve.calibration.engine_capture`) is the
one consumer. Every production deployment leaves this ``None`` —
the tap in :meth:`_finish_execute` is a single ``is not None``
check, same cost class as :meth:`set_raw_logits_sink`.

Fire both calibration taps for a step that has just forwarded.

THE ONE PLACE either tap is read. Every forward path that produces a
step's logits must call this, because a path that does not is
invisible: the sinks stay armed, nothing raises, and calibration
simply receives less data than it asked for -- or none.

That is not hypothetical. ``_finish_execute`` used to read the two
sinks inline, so the MTP seed-draft path
(``run_step._run_mtp_seed_forward``, which forwards and samples
WITHOUT going through :meth:`execute` because it needs the per-token
hidden) captured nothing at all. In-process calibration on an engine
with a DFlash drafter attached routes every one of its prefills down
that path, and refused with "only 0/16 KV layers captured" -- after
running all 48 prefills correctly, with the sinks armed on this very
object.

Cheap by construction: two ``is not None`` tests, both False in every
production deployment, in the same cost class as the step's existing
debug counters.

Materialize the batch, run forward, sample, and return the
per-row sampled tokens plus the populated :class:`StepPlan`.

Accepts either a legacy slate ``[(req, n_tokens), ...]`` or a
:class:`StepPlan` (with the slate carried through on
``plan.slate``). When given a plan, the runner extracts the
slate from it, materializes the batch, then writes the flat
``input_ids`` / ``positions`` / ``cu_seqlens`` tensors back
onto a freshly-frozen plan via :meth:`_populate_tensors`.
That populated plan is the canonical input shape; the legacy
:class:`ScheduledBatch` survives as an internal staging
structure.

``on_inputs_built`` (host-ahead pipeline, ARBI_ASYNC_SCHEDULE):
invoked exactly once, after the batch is materialized, the plan
populated, the batch mirrored to workers, and the admission-time
flush queues drained — i.e. after every request-STATE-DEPENDENT
read (lengths, ``prompt_consumed``, page chains) and every
page-table mutation (``allocate_slots``) of this step is DONE.
From that point to the return, this thread reads only fields
that are STABLE during the caller's overlap window (``is_prefill``
/ ``state`` / ``state_handles`` / ``sampling`` — the engine
thread defers finishes and mutates only output-token bookkeeping,
see run_step.pipelined_step_async) plus the materialized batch,
the model, and the GPU; the per-row ``last_sampled_gpu`` stamp
and this step's own timing records are the only writes. The
caller uses the signal to start the overlapped drain/commit of
the PREVIOUS step on the engine thread while this thread
dispatches the forward. ``None`` (the default) preserves the
byte-identical synchronous path.

Pick the captured-graph replay or the live forward; return logits.

Four-way dispatch: exact-shape decode-graph replay, decode
pad-up replay, whole-forward prefill replay, and the eager
live forward fall-through. The eager-fallback branch needs
``timed_records`` for :meth:`_build_metadata`'s first-call gate.

Eager forward that also scores this chunk's prompt positions.

Returns the same last-token logits every other rung returns, so
the caller is unchanged; the per-token hidden states are consumed
here and never leave.

Whole-forward prefill replay with padding-to-bucket; return logits.

``return_hidden_state=True`` requires a graph captured with
``hidden_out`` retained (the mtp-fill variant; the lookup's
``require_hidden`` gate enforces this) and returns
``(logits, hidden_last)`` clones, where ``hidden_last`` is the
LAST-SLOT ``(B, H)`` hidden — the in-graph gather reads the
refreshed ``cu_seqlens_q``, so it is the live request's
``real_N - 1`` row, not the captured bucket's.

The lookup picks the smallest ``bucket_N >= real_N`` so we may
need to pad the live tensors to ``bucket_N`` before replay.
``cu_seqlens_q`` is refreshed per replay so the captured-graph
last-token gather + Turbo prefill per-row Q range see ``real_N`` instead of
the captured ``bucket_N``.

``seq_lens`` / ``cu_seqlens_k`` carry the FULL K extent
(``batch.max_seq_len == context_len + real_N``); ``cu_seqlens_q``
carries only this chunk's ``real_N`` query span. On a first chunk
the two coincide (``context_len == 0``). On a continuation chunk the
lookup resolved the ``prefill_context=True`` family, whose PAGED
branch derives each row's K window from ``seq_lens`` at replay
(device read), so the full-K refresh here attends the correct
window — see
:func:`~arbi_serve.runtime.captured_lookup.lookup_captured_graph_for_prefill`.

The pad targets are PERSISTENT, reusable staging buffers
(allocated once, sized to ``max(prefill_cudagraph_buckets)`` / the
captured ``max_pages_in_table``) — see
:class:`_PrefillReplayStaging`. ``real_N < bucket_N`` is the
common case (you rarely hit a bucket exactly). Safe because
``CapturedGraph.replay`` ``copy_()``s every kwarg into its OWN
persistent buffers before launching (it does not retain our
``data_ptr``), so the next replay can freely overwrite the
staging slices.

Stamp forward-done, sample (sync or async), stamp sample-done.

The shared sample/timing/return epilogue of :meth:`execute`.
Async output: when ON, sample without the host sync and
return the deferred ``PendingSample`` on the result; the run
loop materializes + commits it one tick later, and the
synchronous ``sampled`` int list stays ``None`` (the caller
MUST go through the deferred drain).

Run one forward pass over ``batch`` and return its logits.

``compute_logits=False`` (only with ``return_hidden_state=True``)
skips the wasted last-token ``lm_head`` GEMV on the LIVE verify
forward — the verify pass recomputes ``lm_head`` over all K+1 rows
itself. The captured-replay path ignores it (the head is baked into
the graph); correctness holds either way since the extra logits are
simply discarded.

Replays a captured CUDA graph when the batch shape is in the
pool (zero-allocation hot path) and otherwise falls through to
a live forward. On TP>1, rank 0 broadcasts the batch to workers
first so every rank runs the row-parallel collectives in
lockstep. When ``return_hidden_state`` is set, also returns the
per-token hidden state (requires a capture that retained it).

One histogram observation per VERIFY forward: the rows every linear saw.

``rows`` is the M the verify GEMM ran at — the slate's flat token
count, or the captured graph's padded row count when a pad-up replay
served it — so the series reads in the currency of
``exl3_int8_verify_min_rows`` and the int8 verify tile ladder, and a
reader can count steps per rung. Always on and host-side only:
``max_query_len`` is already materialised on the batch, the metrics
bundle is one attribute read, and the histogram observe is the whole
cost. Which LEG served the step is decided per linear inside the op
and is not known here; ``exl3_int8_verify_gemm`` and its band census
carry that half.

Classify this eager fall-through ONCE and raise the alarm for the
classes that are anti-invariants, returning the class label.

Two classes are alarms, one per axis of the capture ladder, and both
mean the same thing: a shape the scheduler CAN present has no captured
route and runs the live per-layer forward from now on.

  * ``S = K + 1 > 1`` with ``mtp_meta`` and no captured graph — the
    verify width axis — fires ``eager_verify_offladder``.
  * ``S == 1`` with no captured graph — the batch-size axis — fires
    ``eager_plain_offladder``.

Each fires its always-on counter plus a WARNING bounded to one line per
distinct shape. Every other class is returned for the debug step-mix
histogram only. One classification feeds both consumers, so the
production alarm and the debug accounting can never disagree about why
a forward went live.

Per-(model, params_hash) captured-graph cache.

Wraps :class:`arbi_serve.runtime.capture.decode.CapturedGraphPool` with an
extra dimension so the multi-group scheduler can keep already-captured
graphs valid across swaps.

API contract:

  - ``put(model_path, params_hash, graph)`` — register a captured
    graph under its ``(model, params_hash)`` group.
  - ``get(model_path, params_hash, batch_size, seq_len)`` — return the
    matching graph or ``None``.
  - ``drop_for(model_path, params_hash)`` — forget every graph for
    one group; called by the swap orchestrator on backend rebuilds and
    model swaps.

The engine's hot-path lookup branches on ``len(group_pool) > 0`` then
calls ``group_pool.get(B, S)`` — same semantics as the single-group
path, just one extra hash-lookup deeper.

Return the captured graph for
``(model, params, B, S, lora_bucket, is_prefill,
kv_pages_bucket, prefill_context)`` or ``None`` on miss.

``is_prefill`` defaults to ``False`` so existing decode +
MTP-verify call sites retain their semantics; whole-forward
prefill lookups must pass ``is_prefill=True`` explicitly. The
underlying :class:`CapturedGraphPool` keys on the same flag,
so prefill captures never collide with decode captures at the
same ``(B, S, lora_bucket)`` flat-token count.

``kv_pages_bucket`` defaults to ``0`` (the legacy "no bucket"
sentinel) so pre-bucketing call sites continue to hit pre-
bucketing captures. Bucket-aware lookups pass the resolved
bucket; the runner picks the smallest configured bucket that
covers the live ``max_seq_len``.

Multimodal merge + M-RoPE position plumbing for
:class:`EagerModelRunner` — modality-agnostic over the model's
``mm_bindings`` (see :mod:`arbi_serve.multimodal.registry`).

The two collaborators (``attach_multimodal`` / ``ensure_media_encoded``)
are module-level free functions taking the runner explicitly. The
runner keeps thin delegating method wrappers so the
``runner._attach_multimodal`` / ``runner._ensure_media_encoded`` call
sites (the ``_build_batch`` hot path's text-only-skip guard + tests)
are preserved.

Text-only steps never enter here — ``_build_batch`` gates the
``_attach_multimodal`` call behind a cheap ``any(... multimodal ...)``
check.

Build ``batch.mm`` (+ M-RoPE positions for M-RoPE models).

Every tower runs **once per request** — on the first prefill step
that touches it, ``ensure_media_encoded`` encodes all of the
request's media (every modality) and caches the position-ordered
embeddings (``_media_embeds``) on the Request; M-RoPE models also
cache the full 3-D position grid (``_mrope_full``) and the
decode-continuation offset (``_mrope_delta``).

Three regimes:

  * **Prefill media chunk.** Scatter only the cached media tokens
    whose placeholders fall in *this* chunk
    (``[prompt_consumed, prompt_consumed+n)``). Because the towers
    already ran, a single media item may span multiple prefill
    chunks — the encoder never needs the chunk to hold the whole
    item.
  * **Decode step.** Media tokens are gone from the query; on
    M-RoPE models the stored delta is added to the row's flat
    positions IN PLACE (shapes stay 1-D so the captured decode
    graph is untouched). Flat-position modalities need nothing.
  * **Text rows** in a mixed batch get their flat positions
    broadcast across all three M-RoPE channels (M-RoPE models
    only).

Order per-modality embeddings by placeholder position in the prompt.

Single-modality requests (every current model) short-circuit to the
one tensor. A mixed request walks the prompt once: each contiguous
run of a modality's placeholder id consumes the next rows of that
modality's queue, yielding one tensor aligned with the mask's
left-to-right ``True`` positions.

Encode a request's media (all modalities) once and cache on it.

Idempotent — subsequent prefill chunks of the same request reuse
the cache. Runs each modality's tower eagerly in the runner's CUDA
context; the position-ordered embeddings land in
``req._media_embeds``. M-RoPE models additionally cache the full
3-D position grid and the decode-continuation delta.

Named per-bucket :class:`torch.cuda.MemPool` wrapper.

PyTorch's caching allocator uses a single default pool for every
allocation by default — KV slabs, CUDAGraph capture, LoRA adapters,
sampler scratch, model activations all share one heap. Once the
caching allocator fragments, large transient allocations (e.g. a
14K-prefill RoPE-cache batch) can OOM even when the headline ``free``
counter looks healthy.

The fix: name your buckets. ``torch.cuda.MemPool()`` (PT 2.4+) gives
you a private heap; the
``with torch.cuda.use_mem_pool(pool):`` context routes every alloc
inside the block into that pool. Different lifecycle classes
(static-for-engine-lifetime

This module defines :class:`NamedMemPool` — a thin wrapper that:

  - owns one ``torch.cuda.MemPool`` per logical bucket named in the
    ``category.role`` taxonomy (``state.attn_kv``, ``capture.cudagraphs``,
    ``model.lora``, ``state.mla_kv``, ``state.gdn_recurrent``,
    ``state.conv``, ``scratch.attn_codec``; see
    :mod:`arbi_serve.runtime.pool_taxonomy`);
  - exposes ``with named_pool.use():`` so call sites are one line;
  - tracks ``size_target_bytes`` (operator hint for what's expected to
    land here) and the high-water peak;
  - reports ``snapshot()`` with the metric names the metrics gauges
    consume (``caching_allocator_allocated_bytes``,
    ``caching_allocator_reserved_bytes``, ``fragmentation_ratio``).

Pools live on the engine and survive for the engine's lifetime.
Engine shutdown drops them; the underlying ``torch.cuda.MemPool`` is
freed when its last allocation goes out of scope.

Hard-cap enforcement
--------------------
``torch.cuda.MemPool`` itself doesn't expose a hard limit (PyTorch
provides no per-allocation interception hook for caching-allocator
sub-pools). When ``enforce_size_target=True`` we snapshot the device
allocator's ``memory_allocated`` on :meth:`use` enter and re-check on
:meth:`use` exit; if the delta plus the prior peak crosses
``size_target_bytes`` we raise :class:`PoolSizeExceededError` so a
misconfigured pool fails LOUDLY rather than silently growing past
its budget.

We also emit a one-shot WARNING at :meth:`snapshot` (used by the
metrics callback) when a pool's live allocated bytes cross 80% of the
configured ``size_target_bytes`` so operators get advance notice
before a hard-cap raise on the next ``.use()`` cycle.

THE single chokepoint for constructing a ``torch.cuda.MemPool``.

Every arbi MemPool — the plain caching-allocator pool a
:class:`NamedMemPool` wraps and the cuMem-backed pool
:meth:`CuMemPoolAllocator.make_pool` builds — is born here and registered
in :data:`_MEMPOOL_REGISTRY` at construction, so no untracked
("wrapper-less") pool can exist. The capture keep-alive pins from this
complete registry; the demoted ``gc.get_objects`` scan only proves it is
complete (see :func:`begin_capture_keepalive`).

NO arbi code path may call ``torch.cuda.MemPool(...)`` directly — a CI lint
guard (``tests/test_mempool_chokepoint_guard.py``) fails the build if the
raw constructor appears anywhere outside this module.

``allocator``: the optional pluggable allocator passed straight to
``torch.cuda.MemPool`` (a cuMem ``torch_allocator`` object for cuMem-backed
pools; ``None`` for the default caching-allocator pool).

Snapshot the live MemPools in the canonical registry (strong refs).

Materialises the :data:`_MEMPOOL_REGISTRY` ``WeakSet`` into a list of
strong refs for the caller. Used by :func:`begin_capture_keepalive` as the
primary keep-alive source and by the completeness assert.

Snapshot the live :class:`NamedMemPool` wrappers (strong refs).

The deterministic replacement for a ``gc.get_objects()`` scan when a
caller needs every live wrapper — e.g. the test-suite teardown that
balances ``c10::cuda::MemPool`` ref-counts. Cost is O(live pools), not
O(heap).

Return True if a CUDA graph capture is recording on the current stream.

Used by :meth:`NamedMemPool.use` to no-op begin/end-allocate-to-pool
when entering a different named pool's ctx while
``torch.cuda.graph(pool=...)`` is already routing every allocation
to a captured pool. Calling
``_cuda_beginAllocateCurrentThreadToPool`` for a foreign pool
mid-capture corrupts the allocator's internal capture state —
PyTorch will core-dump on engine teardown with
``c10::Error: captures_underway.empty() INTERNAL ASSERT FAILED``.

Wrapped so unit tests on CPU-only boxes (no CUDA context) don't
hit ``RuntimeError`` from poking the CUDA API. Returns False when
CUDA isn't initialised — the cudagraph code path can't run there
anyway.

Raised when a :class:`NamedMemPool` allocation would push the pool
past its configured ``size_target_bytes``.

Carries the pool name + observed bytes + configured target so logs
and admin endpoints can show actionable numbers.

Routing state threaded from :meth:`NamedMemPool._enter_pool_routing`
to :meth:`NamedMemPool._exit_pool_routing` and
:meth:`NamedMemPool._check_size_target_on_exit`.

``capturing``: a cudagraph capture was active on the stream at enter.
``suppressed``: begin/end-allocate were skipped (nested same-pool entry
or active capture). ``device_index``: resolved CUDA device ordinal.
``tag_ctx``: the cuMem per-alloc tag scope the caller wraps the block in.
``es``: the open ``expandable_segments_off`` scope to close on exit
(``None`` for non-cuMem pools or when suppressed).

True while any ``NamedMemPool.use()`` block is open on any thread.

ADVISORY ONLY — a bare read, so by the time a caller acts on it the
answer can already be stale. Never guard an allocator query with this;
use :func:`allocator_query_lease`, which makes the check and the query
one indivisible step. Kept for callers that only want to describe the
state (logs, diagnostics) rather than act on it.

Yield True only while it is SAFE to query the CUDA caching allocator.

The one correct way for a non-engine thread to reach
``torch.cuda.memory_reserved`` / ``memory_stats`` / ``memory._snapshot``.
See :data:`_ALLOC_GUARD` for the deadlock this prevents::

    with allocator_query_lease() as ok:
        if not ok:
            return ()  # skip the sample
        ...  # query freely; the window is held open

Yielding False costs one gap in a metric series. Ignoring the lease
costs the process.

The query runs while the lease holds :data:`_ALLOC_GUARD`, so an engine
thread entering ``use()`` blocks at the counter bump — which releases the
GIL, letting the query finish. Do NOT do slow or blocking work inside a
True lease: an engine thread is waiting on it.

Declare a build in progress: allocator queries stand down for it.

See :data:`_ENGINE_BUILD_DEPTH`. Re-entrant across threads by counting,
so a nested / re-entered build (pool-member prep, admin reload) closes
the window only when the outermost one does.

Yield True only while a serving-time pool REPAIR may run.

The counterpart of :func:`allocator_query_lease` for a caller that does not
merely READ the allocator but destroys and re-maps a pool through it — the
idle forward-arena re-tightening
(:func:`~arbi_serve.engine.arena_watch.note_engine_idle`). It runs on the
engine's step loop, which keeps ticking while a member build replaces that
same engine's pools on another thread (the backend swap and every live
config override build one), and there it is wrong twice: it prices the
release off ``boot_state`` arena rows the build is part-way through
rewriting, and it enters ``use()`` against the build's own allocation, which
is the GIL-vs-allocator-mutex inversion :data:`_ALLOC_GUARD` describes —
both threads parked, 0% CPU, nothing recovers it.

A query lease will not do: the repair itself opens ``use()`` blocks, and
those take :data:`_ALLOC_GUARD`, so a repair holding it would wait on
itself. This publishes a window instead — the same shape ``use()`` does —
which :func:`engine_build_window` waits out, so the exclusion holds in both
directions from ONE lock.

Yielding False costs one skipped repair attempt; the caller must leave its
"already tried at this pool size" memo unset so a later idle retries.

One named ``torch.cuda.MemPool`` per logical bucket.

Args:
    name: stable taxonomy label (``state.attn_kv``,
        ``capture.cudagraphs``, …; see
        :mod:`arbi_serve.runtime.pool_taxonomy`). Carried
        on the ``pool`` metric label so a Grafana panel can split
        ``caching_allocator_*_bytes`` per bucket.
    device: CUDA device the pool's allocations live on.
    size_target_bytes: operator hint for the bucket's expected
        steady-state size. Surfaced on :meth:`snapshot`. NOT
        enforced unless ``enforce_size_target=True``.
    enforce_size_target: if True, :meth:`use` raises
        :class:`PoolSizeExceededError` on context exit when the
        delta from the snapshot taken on enter would push live
        allocated bytes past ``size_target_bytes``. Production
        opt-in for buckets with deterministic budgets (e.g.
        ``activation_arena``, ``lora_pool``); leave OFF for buckets
        whose footprint legitimately varies (KV slabs grow with
        occupancy).

Every pool-id form this segment could be matched on, as tuples.

The mirror image of :func:`_pool_id_match`'s key walk, hoisted so a
caller reading MANY pools can index the walk once instead of re-running
that walk per pool. Both must stay in step: the index is only sound
while it is built from the same key set the predicate tests.

``{pool-id tuple: [segments]}`` — ONE pass over the walk.

Matching used to be quadratic in the two things a deployment grows:
every pool tested every segment, so a registry read ran
``len(pools) * len(segments)`` id comparisons. MEASURED on a live boot
after the walk itself was already shared: 18,060 ``_pool_id_match``
calls across ten ``/metrics`` scrapes, 2.0 ms of the 7.8 ms — the
largest single term left in the exposition once the walks were gone.

A segment lands under every key form it carries, because a caller
holding one form must find it; the per-pool lookup then tries the
pool's own id and takes the first hit.

Does a snapshot segment belong to the given ``MemPool``?

PT exposes ``mempool.id`` as a ``(driver_handle, pool_id)`` tuple.
Snapshot segments carry the same tuple under ``segment_pool_id``
(older builds) or ``pool_id`` / ``allocator_pool_id``. Walk
every plausible key — the worst case is ``False``.

Sum ``fn(seg)`` over every allocator-snapshot segment owned by ``mempool``.

Shared spine for :func:`_bytes_for_pool` and :func:`_reserved_for_pool`:
accumulate ``fn`` over the segments whose pool id matches.

``segments`` is a walk the CALLER already took. The walk is the whole cost
of this read -- ``torch.cuda.memory._snapshot`` materialises every segment
and every block of the caching allocator as Python dicts, and it takes the
allocator mutex to do it -- while the matching below is a few dict lookups
per segment. A caller reading many pools therefore takes ONE walk and hands
it to all of them; passing ``None`` takes a private walk, which is correct
for a single ad-hoc read and quadratic for a registry-wide one.

Both sums from ONE pass over the segments owned by ``mempool``.

:meth:`NamedMemPool.snapshot` needs both numbers and the pool-id test is
the same test for both, so asking for them separately walked the segment
list twice and ran that test twice per segment.

``index`` is :func:`index_segments_by_pool` over the caller's walk. With
it this pool touches only its OWN segments; without it every pool scans
the whole list, which is quadratic in the two things a deployment grows.

Best-effort wrapper around ``torch.cuda.memory._snapshot``.

Snapshot is a private API but stable across PT 2.x — we fall back
cleanly to an empty list if anything goes sideways so a metric
callback never raises into the OTEL exporter thread.

Whether this pool hosts cudagraph capture allocations.

Read by teardown (``_release_mempool_registrations`` /
``reap_leaked_pools``) to skip the manual ``releasePool`` /
destruct that would double-release a CUDAGraph-owned registration.

Setting it ``True`` ALSO permanently pins this pool's underlying
``torch.cuda.MemPool`` in :data:`_CUDAGRAPH_POOL_PINS`. That pin is
the deterministic guarantee that a cudagraph capture pool's
``~MemPool`` can never fire via a stray refcount drop during a
subsequent swap's capture sweep. A cudagraph pool is kept
parked for the whole process lifetime anyway (``reap_leaked_pools``
never reaps ``_holds_cudagraphs`` pools), so the permanent pin costs
no extra VRAM beyond what is already retained — it merely makes the
"no ``~MemPool`` for a capture pool, ever" invariant independent of
whether the wrapper is still reachable from the (re-created) registry
or ``_LEAKED_POOLS`` when the next sweep snapshots them.

True once a cudagraph capture RAISED while recording into this pool.

A capture that raises mid-``torch.cuda.graph(...)`` (an OOM, a cuMem
cap denial, a stream-capture invalidation) can leave a PARTIALLY-
captured ``torch.cuda.CUDAGraph`` holding DEVICE pointers into this
pool's physical. That partial graph is often NOT reachable from the
engine's capture-pool dict (the capture context raised before the
graph was stored), so ``clear_all_cudagraphs`` never ``reset()``s it —
and its device references outlive every Python ref. ``use_count() <= 1``
is therefore NOT proof the graph is gone.

Teardown honours this latch as a hard invariant: ``free_all`` parks the
pool WITHOUT unmapping its physical, and ``reap_leaked_pools`` NEVER
destructs it (even on the ``force_cudagraph`` path). The pool stays
parked in ``_LEAKED_POOLS`` for the process lifetime; the OS reclaims
its VRAM at exit. Freeing it any earlier would pull the physical out
from under the still-live partial graph → SIGSEGV when that graph is
finalized. The stranded VRAM is the deliberate price of not crashing a
boot (or a still-serving engine on a mid-serving swap) that already
failed loud on the capture OOM.

True once every CUDAGraph referencing this pool has been ``reset()``.

Set by :meth:`Engine.clear_all_cudagraphs` AFTER it resets every owned
``torch.cuda.CUDAGraph`` (and only if all resets succeeded). It is the
POSITIVE PROOF that no live graph references this pool's physical, which
is the sole precondition under which ``reap_leaked_pools`` may
force-destruct a ``_holds_cudagraphs`` pool.

Why ``use_count() <= 1`` is not a substitute: a ``_holds_cudagraphs``
pool retains one unbalanced ``begin_allocate`` registration per capture-
sweep ``use()`` cycle, so ``reap``'s ``force_cudagraph`` path
``releasePool``s the count down to 1 to reclaim it — but that same
force-release drops a STILL-LIVE captured graph's registration if the
graphs were never reset, freeing the pool's physical out from under the
graph (SIGSEGV on any post-capture boot abort: a KV-floor refusal, a
first-forward failure). Gating the force path on this latch means a pool
parked on a teardown that never reset its graphs stays parked for the
process lifetime (the OS reclaims its VRAM at exit) instead of being
destructed while a graph still points into it.

Issue the ``releasePool`` calls the :meth:`use` machinery defers
to teardown, so the underlying ``torch.cuda.MemPool`` can be
destructed cleanly. Returns ``True`` when the pool reaches
``use_count() == 1`` (safe to destruct / its VRAM can be reclaimed
by ``empty_cache``), ``False`` when a residual ref remains (e.g. a
live ``CUDAGraph`` still pins the pool) — in which case the caller
MUST keep the pool object alive, since destructing it would trip
``~MemPool``'s ``use_count() == 1`` assert and ``terminate``.

``use_count()`` is 1 (the MemPool's own ref) plus one per
unbalanced ``begin_allocate`` registration the CUDA caching
allocator still holds — one per ``use()`` cycle that ran on this
pool. Freeing tensors / ``empty_cache`` does NOT drop these; only
``releasePool`` does. We issue one release per registration until
the count is back to 1; the ``before`` guard bails if a release
stops lowering the count (a ref we can't drop this way), so we
never spin or over-release a never-used pool.

The iteration bound is high on purpose: a hot scratch pool
(``gdn_workspace_pool`` enters ``use()`` once per GDN layer per
capture/warmup pass) accumulates hundreds of registrations, and a
cap that is too low leaves such pools partially-released →
parked in ``_LEAKED_POOLS`` → never freed (multi-hundred-MB leak
per boot). The ``before`` guard is the real terminator; the cap
is just a runaway safety net.

Route every allocation inside the ``with`` block into this pool.

Wraps the underlying ``_cuda_beginAllocateCurrentThreadToPool``
/ ``_cuda_endAllocateToPool`` primitives and updates the peak
high-water on context exit. When ``enforce_size_target`` is
True, raises :class:`PoolSizeExceededError` on context exit
when the live allocation count would push the pool past its
configured ``size_target_bytes``.

We snapshot ``torch.cuda.memory_allocated(device)`` on enter
and re-read on exit; the delta is the upper bound on bytes
the block contributed. The exit-check is the most reliable
signal we can produce without an allocator hook.

Re-entrancy + cudagraph-capture awareness
-----------------------------------------
PyTorch's CUDA allocator rejects two scenarios this method
guards against:

1. **Same-pool nested entry.** A second
   ``_cuda_beginAllocateCurrentThreadToPool`` for the same
   pool while already recording raises ``RuntimeError:
   beginAllocateToPool: already recording to mempool_id``.
   Nested wrappers
   (``profile_activation_peak`` calls ``_build_metadata`` while
   already inside ``activation_arena.use()``) hit this. Fix:
   refcount the entries; only call begin/end on the outermost
   pair.
2. **Cross-pool entry during cudagraph capture.** When a
   ``torch.cuda.graph(pool=other_pool.mempool)`` block is
   active on the current stream, the allocator is already
   routed to ``other_pool``. Calling
   ``_cuda_beginAllocateCurrentThreadToPool`` for a DIFFERENT
   pool inside that block corrupts the capture state — every
   subsequent allocation inside the capture lands in our pool
   instead of the cudagraph's pool, the capture's
   ``endAllocateToPool`` never matches a live ``begin``, and
   on engine teardown ``c10::Error: captures_underway.empty()
   INTERNAL ASSERT FAILED`` core-dumps the process. Fix:
   detect via :func:`torch.cuda.is_current_stream_capturing`
   and no-op the begin/end; allocations during capture flow
   into the cudagraph's pool, which is fine because the
   cudagraph IS one of our named pools (``graph_pool``) and
   its bytes are accounted there.

Both cases also skip the enforce-size-target check on exit:
when we no-op'd the begin we have no meaningful delta window,
and the outer ``use()`` will re-check on its own exit.

Drive begin-allocate-to-pool + open the cuMem scope for ``use()`` enter.

Mirrors :meth:`_exit_pool_routing`. Resolves whether begin/end must be
suppressed (nested same-pool entry, or a cudagraph capture already
active on the stream — see :meth:`use`), opens the cuMem
``expandable_segments_off`` scope and builds the per-alloc tag context
for cuMem-backed pools, issues the begin-allocate registration, and
bumps the re-entrancy depth. Returns the :class:`_PoolRouting` state the
caller threads to :meth:`_exit_pool_routing` and
:meth:`_check_size_target_on_exit`.

We do NOT call ``torch.cuda.use_mem_pool`` directly. Its ``__exit__``
calls ``_cuda_releasePool`` which DECREMENTS the mempool's
``use_count``. Since ``MemPool.__init__`` sets ``use_count = 1``, the
first context exit takes it to 0; the NEXT entry asserts
``use_count > 0`` and crashes (observed under
``torchrun --nproc-per-node=2`` worker ranks at boot during
``recurrent_pool.attach_mtp_snapshot_buffers``). The public API was
designed for one-shot ``with use_mem_pool`` blocks, not the re-entrant
``NamedMemPool.use()`` pattern. Fix: drive begin/end allocate-to-pool
ourselves and skip ``releasePool`` on exit. The pool stays registered
until this ``NamedMemPool`` is destroyed, at which point we issue one
``releasePool`` to balance the ``__init__`` increment.

``_cuda_beginAllocateCurrentThreadToPool`` / ``_cuda_endAllocateToPool``
are bound at module import (top of file) rather than re-imported per
call — this method runs on every ``use()`` (per decode step AND per
sample), so the per-call ``from torch.cuda.memory import ...`` was pure
Python churn on the hot path. The symbols are stable torch built-ins.

Reverse :meth:`_enter_pool_routing` on ``use()`` exit.

Drops the re-entrancy depth, issues the end-allocate-to-pool
registration (when begin was driven), and closes the cuMem
``expandable_segments_off`` scope. The cuMem per-alloc ``tag_ctx`` is
closed by the caller's ``with`` block BEFORE this runs, preserving the
original begin → tag → end ordering.

Enforce ``size_target_bytes`` on ``use()`` exit (outermost frame only).

No-op unless enforcement is on, a target is configured, this is the
outermost ``use()`` frame, and no cudagraph capture was active on enter
(a no-op'd begin leaves no meaningful delta window). Raises
:class:`PoolSizeExceededError` when the observed live bytes exceed the
target.

Skips the per-call ``_current_allocated_bytes()`` walk on the
``_enforce=False`` default — it calls ``torch.cuda.memory._snapshot()``
and segment-filters by pool, which py-spy profiling identified as a
decode-hot-path overhead (~17% MainThread CPU). Peak allocated bytes is
sampled via the public :meth:`snapshot` accessor at metrics-tick
cadence, not on every ``use()`` exit.

O(1) read of ``torch.cuda.memory_allocated`` on this device.

Wrapped so tests can monkeypatch a single attribute rather than
the global ``torch.cuda.memory_allocated`` symbol (which is
also called by other modules during a test run).

Best-effort live read of bytes allocated inside this pool.

``torch.cuda.MemPool.use_count()`` is the closest stable
introspection signal; for the byte view we walk
``torch.cuda.memory._snapshot()`` and sum segments whose
``segment_pool_id`` matches our pool's id. The snapshot call
is cheap (a list of dicts) but we skip it on the hot path —
callers query through :meth:`snapshot`.

Return a dict suitable for the OTEL gauges in :mod:`arbi_serve.server.metrics`.

Keys:

  - ``caching_allocator_allocated_bytes`` — bytes the pool's
    outstanding tensors hold (live).
  - ``caching_allocator_reserved_bytes`` — bytes the pool
    reserved from the driver (allocated + free-but-cached).
  - ``fragmentation_ratio`` — ``1 - allocated/reserved`` (0.0
    when the pool is unused).

Plus operator-hint fields (``name``, ``size_target_bytes``,
``peak_allocated_bytes``) so admin endpoints / dashboards can
spot drift.

``segments`` is an allocator walk the caller already took (see
:func:`_walk_pool_segments`). Reading a whole registry pool-by-pool
without it takes two ``torch.cuda.memory._snapshot`` walks PER POOL for
a number that is the same walk every time.

Cudagraph-capture MemPool keep-alive + leaked-pool reaping.

The module-level pool-tracking globals (``_MEMPOOL_REGISTRY``,
``_CUDAGRAPH_POOL_PINS``, ``_LEAKED_POOLS``, ``_CAPTURE_KEEPALIVE``) and the
low-level primitives (``registered_mempools``,
``_is_cudagraph_capture_active``) live in :mod:`arbi_serve.runtime.named_pool`;
this module references them through the ``_np`` module alias so every mutation
targets the single canonical instance (and monkeypatched symbols resolve
correctly). Every public name here is re-exported from
:mod:`arbi_serve.runtime.named_pool`.

Pin a strong ref to every live + parked ``torch.cuda.MemPool`` for the
duration of a cudagraph capture sweep.

Must bracket the ENTIRE multi-bucket sweep (called before the first
``BeginCapture``, paired with :func:`end_capture_keepalive` after the
last ``EndCapture``). Idempotent-ish: a second call replaces the prior
keep-alive set (the sweep is single-threaded, boot/reload-only). Returns
the number of pools pinned (logging only).

PRIMARY path: :data:`_MEMPOOL_REGISTRY` — the canonical registry every
arbi MemPool is registered into at construction (:func:`create_mempool`).
Iterating it pins EVERY arbi-created ``torch.cuda.MemPool`` deterministically,
by census, so no untracked ("wrapper-less") pool can slip the keep-alive.

``registry`` (the engine ``NamedPoolRegistry``), ``_LEAKED_POOLS``, and
:data:`_CUDAGRAPH_POOL_PINS` (every cudagraph capture pool that EVER
existed) are unioned in as belt-and-braces; their MemPools are already
in the canonical registry, but pinning from them too costs nothing and
keeps the cross-swap invariant explicit.

COMPLETENESS ASSERT (the demoted heap-scan): after pinning the
canonical registry we scan ``gc.get_objects`` for any live
``torch.cuda.MemPool`` NOT already pinned. Finding one is either (a) an
escaped arbi pool — a footgun the chokepoint lint should have caught — or
(b) a genuine torch/cuBLAS-internal pool created outside arbi. EITHER case
is LOGGED LOUD with id + use_count and then pinned (safety preserved). The
scan now PROVES the registry is complete instead of silently compensating
for it. Holding the underlying ``MemPool`` objects guarantees no
``~MemPool`` can fire via refcounting OR a tail ``gc.collect`` while the
sweep records — the ``use_count() == 1`` ``~MemPool`` abort this closes.

Release the capture keep-alive set, then deterministically reap.

Call ONLY after the sweep's last ``EndCapture`` (no capture underway).
Dropping the pinned ``MemPool`` refs here lets any pool whose last
OTHER ref already vanished destruct at THIS safe point; we then run
:func:`reap_leaked_pools` so parked-but-now-unpinned pools are freed
deterministically rather than by a later stray refcount drop. Never
raises (boot/reload/teardown-only).

A pinned ``MemPool`` still at ``use_count() > 1`` here is one a captured
CUDAGraph (or a captured-tensor alias) still registers into. Dropping its
pin would let the immediately-following ``reap_leaked_pools`` ``gc.collect``
destruct it — and ``~MemPool`` asserts ``use_count() == 1``, the exact
abort this guard prevents. So we PERMANENTLY retain every such pool (the
same lifetime the ``capture.cudagraphs`` pool already gets) and LOG it
by tag + use_count, so the dangling holder is identified loud instead
of crashing the boot. Pools
back at ``use_count() <= 1`` drop their pin normally and reap.

Pin every registered ``torch.cuda.MemPool`` for one allocate-to-pool window.

Paired with :func:`end_routing_keepalive`, called from
:meth:`NamedMemPool._enter_pool_routing` / ``_exit_pool_routing`` on the
non-suppressed path — i.e. exactly when ``beginAllocateToPool`` /
``endAllocateToPool`` are actually issued, which is exactly when the
allocator's ``captures_underway`` list is non-empty.

Why this is needed even with cudagraphs OFF: ``captures_underway`` is the
allocator's pool-ROUTING list, not a cudagraph list. Any ``~MemPool`` that
fires while it is non-empty runs ``synchronize_and_free_events``, whose
``captures_underway.empty()`` ``TORCH_INTERNAL_ASSERT`` throws a
``c10::Error`` out of a destructor → ``terminate`` → SIGABRT, uncatchable
from Python. ``run_model_forward`` holds one routing block open across the
whole forward, and at boot that forward contains a Dynamo/Inductor compile
whose tracing allocations drive many generational collections; one
unreachable MemPool in any of them aborts the process.

Only the OUTERMOST block pins (the depth counter is process-wide, matching
``captures_underway``'s own scope). Materialising the canonical
:data:`_MEMPOOL_REGISTRY` ``WeakSet`` into strong refs makes every
unreachable-but-uncollected pool reachable again, so the cyclic collector
cannot destruct it inside the window. Never raises.

Release the routing keep-alive on the outermost block's exit.

Called AFTER ``endAllocateToPool``, so when the depth reaches zero
``captures_underway`` is empty and any pool that became garbage during the
window destructs HERE, at a safe point, instead of aborting mid-window.
The pinned refs are handed to a local and dropped OUTSIDE
:data:`_ROUTING_LOCK` so ``~MemPool`` never runs under it. Never raises;
an unbalanced call (depth already zero) is a no-op.

Current allocate-to-pool routing depth (0 = no window open).

Non-zero means the allocator's ``captures_underway`` is non-empty and no
``~MemPool`` may fire. Exposed for tests and the pool observability
endpoint; not used on any hot path.

Context manager bracketing a capture sweep with the MemPool keep-alive.

Equivalent to :func:`begin_capture_keepalive` on enter +
:func:`end_capture_keepalive` on exit, but the drain runs in a
``finally`` so an exception mid-sweep still releases the pin (after the
capture has been force-ended by the sweep's own error handling).

Introspect the module-global ``_LEAKED_POOLS`` parking lot.

Every pool here is one that :meth:`NamedPoolRegistry.free_all` could
NOT drive to ``use_count() == 1`` (e.g. a CUDAGraph or forward-thread
CUDA-TLS ref outlived teardown), so its ``~MemPool`` is suppressed and
its VRAM stays resident until process exit. Across repeated hot-swap
reloads this is the prime accumulator of per-swap framebuffer overhang.

Returns one entry per parked pool with its name, the live
``use_count`` (why it could not be freed), and its currently-reserved
bytes (the resident VRAM it still strands). Cheap (a short list walk +
one allocator snapshot per pool); only called from the observability
endpoint, never the hot path.

Destruct + free any parked pool that has since reached ``use_count() <= 1``.

``force_cudagraph=True`` (teardown only — shutdown / model reload / a
per-member build's failure cleanup, AFTER ``clear_all_cudagraphs`` has
reset every captured graph THIS teardown owns and with NO capture
underway) ALSO reclaims ``_holds_cudagraphs`` pools — but ONLY those
``clear_all_cudagraphs`` actually marked :attr:`NamedMemPool.graphs_cleared`
(its positive proof the graphs are reset). For such a pool the residual
``use_count`` is the ``.use()`` begin_allocate registrations left by the
capture sweep — NOT live graphs — so ``_cuda_releasePool`` safely drives the
count back to 1 (verified: no ``uc >= 0`` / ``captures_underway`` abort once
the graphs are gone). A ``_holds_cudagraphs`` pool NOT marked
``graphs_cleared`` (a post-capture boot abort — KV-floor refusal, first-
forward failure — that ``free_all``s while COMPLETE graphs still pin the
pool) is NEVER force-destructed: ``use_count() <= 1`` is not proof the
graphs are gone, and freeing its physical SIGSEGVs the live graph. It stays
parked for the process lifetime (OS reclaims at exit).

It also drops the matching entries from the permanent
:data:`_CUDAGRAPH_POOL_PINS` safety pin so a freed capture pool's
multi-GiB cuMem pages don't stay mapped forever — otherwise every
teardown compounds into a permanent per-cycle VRAM leak.

``only_pool_ids`` (a set of ``id(mempool)`` values) scopes WHICH pins get
dropped. Pass the id set of the pools THIS teardown's own
``clear_all_cudagraphs`` actually reset — never omit it on a SHARED,
multi-member engine (stable-VA residency): :data:`_CUDAGRAPH_POOL_PINS`
is a single process-wide list covering every cudagraph pool that has
EVER existed, including pools belonging to OTHER, still-live residency
members (a parked member, or the member that was active before this
build started) whose OWN captured graphs this teardown never touched.
Unconditionally clearing the whole list (``only_pool_ids=None``) strips
THEIR permanent pin too — the pool itself isn't destructed yet (still
referenced by that member's own ``named_pools``), but the pin that
guaranteed its ``~MemPool`` could NEVER fire is gone. Whenever that
member's last Python reference is later dropped (a subsequent build's
``reset_model_state_for_build`` repointing ``eng.named_pools``, a
reload, …) its ``~MemPool`` runs
wherever the refcount happens to hit zero — and if ANY OTHER cudagraph
capture is recording on the device at that exact moment,
``synchronize_and_free_events`` asserts ``captures_underway.empty()``
and the process aborts (uncatchable ``c10::Error`` / ``terminate``).
``only_pool_ids=None`` is reserved for the genuine full-process-exit /
single-model-replaced case, where nothing else needs the pin.

A pool lands in ``_LEAKED_POOLS`` when :meth:`NamedPoolRegistry.free_all`
can't drive it to ``use_count() == 1`` at teardown — typically a
``graph_pool`` still pinned by captured CUDAGraphs. On a hot-swap RELOAD
those graphs are cleared (``clear_all_cudagraphs``) on the NEXT teardown,
so the parked pool's count finally drops to 1 — but its wrapper is only
referenced by ``_LEAKED_POOLS``, so Python GC may run its ``~MemPool``
LATER, at an arbitrary point — including DURING the next model's cudagraph
capture, where ``~MemPool`` → ``synchronize_and_free_events`` aborts with
``captures_underway.empty() INTERNAL ASSERT FAILED``.

Call this at a SAFE point (reload teardown, no capture underway, no live
swap-in graphs) so those now-unpinned pools are destructed deterministically
here, returning their VRAM to the driver, instead of by a stray GC mid-
capture. Pools still pinned (use_count > 1) stay parked. Two latches keep a
pool parked on ANY path (including ``force_cudagraph``) — a live CUDAGraph's
device references outlive its Python ``use_count``, so that count is never
proof the graphs are gone: a pool a FAILED capture poisoned
(:attr:`NamedMemPool.capture_failed`), and a ``_holds_cudagraphs`` pool not
proven graph-free (:attr:`NamedMemPool.graphs_cleared`). Both stay parked
for the process lifetime. Returns the number reaped. Idempotent; never
raises (best-effort, teardown-only).

Engine-side registry of named pools (:class:`NamedPoolRegistry`).

The ``NamedMemPool`` class and the ``_LEAKED_POOLS`` teardown parking-lot global
live in :mod:`arbi_serve.runtime.named_pool`; this module references them
through the ``_np`` module alias so ``_LEAKED_POOLS`` stays single-instance and
``make()`` constructs the canonical ``NamedMemPool``. ``NamedPoolRegistry`` is
re-exported from :mod:`arbi_serve.runtime.named_pool`.

One synthetic per-pool snapshot row, shaped like ``NamedMemPool.snapshot``.

``fragmentation_ratio`` is ``1 - allocated/reserved`` where that means
something and 0.0 for ``address_space.released_pool_va``, whose reserved bytes are
VA with no physical behind them — nothing there is fragmented, or resident.

Engine-side registry of named pools.

Holds construction order so :meth:`snapshot_all` returns a stable
list for metric callbacks (Prometheus expects label-stable
cardinality across export windows).

The registry is per-engine (each rank constructs its own) so
multi-rank deployments report ``rank``-labeled gauges via the
metrics path; the registry itself is rank-agnostic.

The cuMem allocation tags ``pools`` may carry.

A pool's allocations are tagged at ``use()`` time: raw ``<name>``
outside a tag namespace, ``"<namespace>/<name>"`` when the build ran
under one (stable-VA residency members). Target BOTH forms — a
member's pool can hold a mix (build-time allocs namespaced, runtime
growth raw). Raw-name-only targeting silently releases ~nothing of a
namespaced member.

Unmap physical pages for every cuMem-backed (sleepable) pool.

Computes ``offload_tags`` / ``discard_tags`` from each sleepable
pool's ``sleep_strategy`` (``keep``-strategy pools are left
mapped) and issues ONE :meth:`CuMemPoolAllocator.sleep` so the
physical VRAM of the targeted pools is released behind stable
VAs. ``exclude`` names pools that must stay mapped regardless of
strategy — the residency park passes ``{"model.weights"}`` so a
parked member's weights stay readable (donor-share invariant).
``namespace`` is the member's tag namespace; its pools' allocations
carry ``"<namespace>/<name>"`` tags, so it must be passed or the
sleep targets nothing of a residency member (see
:meth:`_effective_tags`). Returns the total bytes released. No-op
(returns 0) when no sleepable pools are registered or the cuMem
driver is absent.

``discard`` names pools whose CONTENT this caller declares dead for
this sleep: they are unmapped without the D->pinned-host copy their
``sleep_strategy`` would otherwise take, and their pages come back
zeroed on the next wake. A pool's strategy states what a sleep must do
to keep the pool usable in general; ``discard`` is the caller stating
that on THIS sleep there is nothing left to keep — the drained
residency park, whose woken member re-prefills from empty (see
``StableVaResidencyController.discard_sequence_state_on_park``). The
offload it skips is unswappable pinned host RAM, so this is the
difference between a park that costs the host nothing and one that
can exhaust it.

Re-map physical pages at the SAME VAs for every sleeping pool.

Mirror of :meth:`sleep_all`. With ``namespace`` the wake is SCOPED
to this registry's raw + namespaced tags so it can never remap
another parked member's slept allocations; without it the legacy
global wake runs (solo engine — every unmapped alloc is its own).
Returns total bytes mapped back. No-op (returns 0) when no
sleepable pools / no cuMem driver.

Remove + return the registered pool for ``name`` (None if absent).

Used by the shutdown drain's weight-share ordering: a DONOR member's
``weights_pool`` whose storage is still aliased by the ACTIVE sharer
is popped out of the donor's registry before its ``free_all`` so the
pool is torn down ONCE, after the last sharer's model is dropped —
never while a live module still reads its storage.

``MemPool.id`` -> former pool name, for every pool destroyed by
:meth:`release_empty_pool`.

The segments torch still records under these ids hold NO physical: the
cuMem tag behind them was unmapped and released before the pool object
was dropped. Reporting surfaces label them ``address_space.released_pool_va``
so the bytes are not read as an unattributed leak.

Stable-ordered list of every pool's snapshot. Used by metrics.

``segments`` is an allocator walk the caller already took; every pool
is then read off that one walk instead of taking two of its own (see
:meth:`~arbi_serve.runtime.named_pool.NamedMemPool.snapshot`). The
walk is indexed by pool id ONCE here, so reading N pools is linear in
the walk rather than N passes over it.

:meth:`snapshot_all` + the THREE synthetic rows that carry every
caching-allocator byte no named pool accounts for.

The residual (``memory_allocated/​reserved(device)`` minus the named-pool
sums) is not one thing, and emitting it under one name is what let a
dashboard read live default-pool tensors as untagged bytes. It is split
here by the same rule the boot ledger uses
(:func:`arbi_serve.engine.memory_budget.pool_residency.residency_by_pool`),
from one snapshot walk:

  * ``unpooled.torch_default_pool`` — segments outside every private
    ``MemPool``. Resident, live, unbudgeted, and NOT a leak: these are
    real tensors that escaped ``NamedMemPool.use()``. Relocating them
    frees nothing; only ``reserved − allocated`` is recoverable.
  * ``address_space.released_pool_va`` — segments of a pool
    :meth:`release_empty_pool` destroyed. ``memory_reserved`` still
    counts their VA, the driver already took the physical back, and
    ``allocated`` is 0 by construction (the pool had to be empty to be
    released). Nothing here is reclaimable because nothing here is
    resident.
  * ``unpooled.unregistered_pool`` — what neither of those explains: a private
    ``MemPool`` that is neither registered nor known-released. This is
    the only one of the three that is a leak signal: the bytes have an
    owner, and nobody declared it.

    It is NOT the boot ledger's ``driver.residual`` row. That is a
    different quantity on a different surface — the WHOLE-card identity's
    remainder, read from ``mem_get_info`` rather than from the torch
    segment walk — and it legitimately absorbs driver-resident cubin
    growth, so a non-zero value there is not by itself a leak. Reading
    one's docstring at the other is what sends a reader hunting a
    ``MemPool`` that does not exist.

The three sum to the residual exactly: the first two are MEASURED off
the walk and the third takes the difference, so ``Σ series`` still
equals ``torch.cuda.memory_reserved`` / ``memory_allocated``.
``unpooled.unregistered_pool`` is not clamped — a negative value means one of
our own counters disagrees with torch, which is information.

``device``: defaults to the first registered pool's device.
``None`` returns the named-pool-only snapshot when CUDA is
unavailable so test fixtures don't crash.

``{driver label: (allocated, reserved)}`` from ONE snapshot walk.

Best-effort: a walk that fails yields nothing, and the whole residual
then lands on ``unpooled.unregistered_pool`` — the row whose documented job is
to absorb what nothing else measured.

Reserved-but-unallocated bytes held by the TORCH caching allocator,
which are RESIDENT physical yet INVISIBLE to ``CuMemPoolAllocator.
mapped_bytes``.

cuMem-backed pools (``cumem_backed``) map their physical through the
pluggable allocator, so their bytes ARE counted by ``mapped_bytes`` /
the freeze cap. The torch-caching-backed pools (the non-sleepable ones —
``activation_arena`` and any scratch pool the cuMem driver was
unavailable for) and the ``unpooled.torch_default_pool`` residual instead
live in PyTorch's own caching allocator, which holds ``reserved −
allocated`` free segments resident behind the stable VA (an
``empty_cache`` between the capture grow and the freeze does not return
per-MemPool segments). That overhang is the bulk of the post-capture
``driver_free vs accounted_free`` gap: genuine cuMem-invisible resident
scratch, NOT a leaked pool. Surfacing it as a MEASURED quantity lets the
freeze reconcile recognise it instead of false-flagging it.

Sums ``max(0, reserved − allocated)`` over every non-cuMem-backed pool
plus the default-pool overhang. Returns 0
when CUDA is unavailable. No double-count: cuMem-backed pools are
excluded because their physical is already in ``mapped_bytes``.

Hand one EMPTY pool's reserved physical back to the driver, then
re-create it empty under the same name. Returns the bytes released.

THE ONLY mechanism that returns a named pool's reserved-but-free bytes.
``torch.cuda.empty_cache()`` does not visit a private
``torch.cuda.MemPool``'s block pools (pytorch#145168 — re-measured on
``memory_reserved`` unchanged across ``empty_cache``), so both reclaim
passes in the capture handoff step over these bytes and the freeze
ledger reports them, correctly, as resident held-free. Destroying the
pool is what frees them — this is :meth:`free_all`'s per-pool
release → drop → ``empty_cache`` sequence, applied to a single pool and
followed by a fresh registration so callers keep a working pool.

REFUSES a pool with live allocations. The bytes behind a live tensor are
not reclaimable at any price, and tearing the pool down under one is a
use-after-free, not a reclaim. A non-empty pool here is also the exact
signal :mod:`~arbi_serve.runtime.pool_taxonomy` documents for the loader
scratch pool ("stays 0 after boot; non-zero ⇒ a leaked load path"), so
it is reported at WARNING rather than silently skipped.

BOOT-ONLY by contract. Never call this after a cudagraph capture has
touched the pool: captured graphs bake ``base_va + offset``, and this
releases the physical behind that VA. Callers must run it before the
capture sweep (the loader seam does).

Release (``cuMemUnmap`` + ``cuMemRelease``) one cuMem-backed pool's
physical pages, keeping only its VA reservation.

A TEARDOWN release, not a park: the caller is destroying the pool, so
the allocator must mark these allocs released
(:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.release_tags`)
rather than asleep. Their records outlive the teardown — torch keeps
segment records for a destroyed

Targets BOTH tag forms — raw ``<name>`` and the namespace-qualified
``"<namespace>/<name>"`` a stable-VA residency member's allocations
carry — so a namespaced member's physical is actually released (a
raw-name-only discard silently frees nothing of it, the same miss
:meth:`_effective_tags` guards for sleep/wake). Callers MUST only
invoke this for a pool proven graph-free: unmapping physical a live
CUDAGraph still references SIGSEGVs on that graph's later destruct.
No-op when the cuMem driver is absent.

Release every pool's underlying ``torch.cuda.MemPool`` and
return its VRAM to the driver.

Engine ``shutdown()`` calls this after dropping every tensor and
every CUDAGraph that referenced the named pools, so by the time we
get here each pool's ``use_count()`` should be 1 (its own ref) plus
the CUDA caching allocator's deferred ``begin_allocate``
registrations. :meth:`NamedMemPool._release_mempool_registrations`
balances those so the MemPool reaches ``use_count() == 1`` and can
be destructed cleanly — at which point ``empty_cache`` hands the
segments back to the driver. This is load-bearing for arbi-serve's
fast sleep/wake + restart story: leaking the pools instead would
accumulate VRAM per boot and OOM long-lived processes that
boot/teardown engines repeatedly.

A pool that does NOT reach ``use_count() == 1`` (a residual ref we
can't drop — e.g. a CUDAGraph that outlived shutdown) is parked in
the module-level ``_LEAKED_POOLS`` so its destructor never runs and
never aborts; the OS reclaims that VRAM at process exit. This is
the safety net, not the common path.

INVARIANT: never ``cuMemUnmap`` physical that a still-live CUDAGraph
references. A capture OOM leaves a PARTIALLY-captured CUDAGraph
pinning ``capture.cudagraphs`` (``use_count() > 1``). The cuMem
discard-sleep (``cuMemUnmap`` + ``cuMemRelease``) is therefore issued
per pool AFTER that pool proves graph-free — a still-pinned pool is
parked with its physical INTACT so the graph's baked ``base_va +
offset`` pointers stay valid until it is destructed (the OS reclaims
the VA at process exit). Unmapping first and parking second frees the
physical out from under the live graph, whose later destruct then
faults on freed device memory (SIGSEGV).

A pool a FAILED capture poisoned (:attr:`NamedMemPool.capture_failed`)
is parked unconditionally — even if its Python ``use_count`` has since
dropped to 1 — because the partial graph's DEVICE references outlive
its Python refs; ``reap_leaked_pools`` likewise never destructs it.

Paged ``tkv-bypass`` runtime for the Nemotron VoiceChat TTS decoder.

The TTS decoder owns a logical page table because its Gemma stack has a
different layer geometry and lifetime from the conversation model. Its slabs
use the engine's attention-KV memory pool, so boot sizing and allocator
accounting cover both decoders. Every connection receives one conditional row
and, when classifier-free guidance is enabled, one unconditional row.

Real-scheduler-driven duplex STT decode-tick stepper for
NemotronLabs-VoiceChat-11B.

STT decode path
ALONGSIDE ``EngineSttStep``
(``arbi_serve/runtime/nemotron_voicechat_stt_step.py``), touching neither
that class nor its production call site
(``arbi_serve/realtime/nemotron_voicechat_turn.py``'s
``stream_nemotron_voicechat_turn``). Per §7.7.2's "the three-region design
collapses to one region in duplex mode": a persistent duplex connection has
no turn boundary, so every tick is region 2's audio-frame-fusion step,
forever — this module only ever ports that region's logic (mirroring
``EngineSttStep._advance_frame_step``), never region 3's (ordinary
post-frame decode has no duplex analog).

=== A tick has TWO halves, because two different things can drive it ===



  * :func:`prepare_duplex_tick` — fuse this tick's audio frame with the
    session's own previous text/function tokens and park the result on
    ``request.pending_embed_override``. Must run BEFORE the step is
    scheduled.
  * :func:`finish_duplex_tick` — read both logits heads out of the
    completed step, run the frame-lockstep feedback update and the Site-3
    BOS/EOS ``agent_idle`` switch, and advance the frame counter.

:func:`step_duplex_tick` still exists and is unchanged in behavior: it is
the PULL caller, sandwiching its OWN ``scheduler.schedule()`` +
``run_step_async`` between the two halves. The PUSH caller
(:mod:`arbi_serve.realtime.duplex_lane`) cannot do that — in a real
deployment ``run_forever``'s own always-on loop already issues the step
(design doc §7.8's blocking finding) — so it calls the halves from its
pre-schedule and post-step hooks around the loop's own step instead.

=== The dual-head gap this module closes ===

The generic engine dispatch (``forward_exec.collapse_dual_head_logits``)
drops every logits head but "text" before ``sample()``/``run_step.step()``
ever sees the rest — correct for an ordinary single-head caller, but a
duplex tick needs the function head too: greedily sampled and fed back as
this session's ``prev_function_token_id`` (load-bearing for the model's own
trained recurrence, not cosmetic — §7.7.2).

The heads are retained PER REQUEST (design doc §7.42): a duplex request
sets ``Request.retain_logits_heads``, and ``forward_exec
.stash_retained_logits_heads`` writes that request's OWN row of every head
to ``Request.logits_heads`` before the collapse, while still returning the
SAME collapsed value it always would — so ``run_step``'s sampling/commit
logic downstream is bit-identically unaffected and only the opted-in
request observes the extra head.

This is the THIRD shape of this mechanism, and the first one whose state
lives where the state belongs. v1 rebound ``forward_exec.run_model_forward``
process-globally: safe only while the caller owned the whole engine (the
44a tests), NOT safe once the forward is issued by ``run_forever`` on its
own thread, since the rebound global would be visible to every other
concurrent caller. v2 replaced it with ``eng._dual_head_sink``, an
engine-global dict armed and disarmed around exactly one step — which fixed
the concurrency hazard but still held ONE slot, keyed by nothing, for what
is per-row state, and needed arm/disarm/rollback bookkeeping on every exit
path to keep a step that never ran from leaving a stale capture behind.
NVIDIA's own reference reached the per-request answer directly
(``custom_outputs``, design doc §7.38.5); this is that shape, fitted to
this engine's own ``Request``/``forward_exec`` conventions.

=== Function calling: detection here, execution nowhere near here ===

DETECTION is one extra
single-token detokenize plus a substring scan — cheap enough to sit on the
engine thread, which is where it has to be, since the lane's hooks run
inside ``run_forever``'s loop body. Detection flips the session's
``fc_state.in_progress`` gate, which forces this tick's TEXT token to the
checkpoint's PAD id before the frame-lockstep feedback, the Site-3 switch
and the TTS subword read it — the reference's ``_fc_in_progress`` mute.

Tool EXECUTION never touches this module: the detected calls ride out on
:attr:`DuplexTickResult.function_calls`, the WS/session side does the
``response.function_call_arguments.done`` / ``conversation.item.create``
round trip (up to 10 s), and the result comes back as
``DuplexConnection.push_context_tokens`` — the injection row,
made correct on the lane by §7.9.15.

=== Eager-only ===

Every override-carrying row is already forced onto the eager forward path
by the generic engine's own dispatch (``model_runner.py``'s
``_forward_or_replay``) — this module builds no new CUDA-graph capture
machinery (§7.7.2's "accept eager-always as the Phase-2 MVP performance
floor").

One duplex tick's outcome — everything a caller (the driver loop,
the duplex lane's post-step hook, or this module's own live-GPU parity
test) needs to inspect after the tick completes.

``text_logits``/``function_logits`` are this tick's raw, per-vocab
logits (``(vocab_size,)``, float32, already sliced to this request's
own row) — the SAME shape/dtype convention
``EngineSttStep._pending_logits`` uses, so a parity test can compare
them directly.

``text_token_id`` is the value every downstream consumer reads: the
frame-lockstep feedback, the Site-3 BOS/EOS switch and the TTS
subword alike. While a tool call is pending it is the checkpoint's
text PAD id rather than what the sampler produced — the FC pause gate
(design doc §7.9.16). ``sampled_text_token_id`` is the ungated value,
carried for observability/tests only; nothing acts on it.

``fn()``, or ``None`` when the hook is absent — the same
``getattr``-tolerant contract every session read in this module uses,
for the ones that are methods rather than attributes. Minimal session
stand-ins in CPU tests predate these hooks.

TAKE this request's own (text, function) logits off
``req.logits_heads`` — read and reset to ``None`` in one step.

Take semantics, not a plain read, is the whole reason a stale capture
is unrepresentable here: the only writer
(``forward_exec.stash_retained_logits_heads``) runs once per forward
for this row, and the only reader clears the slot, so a tick that was
prepared but never stepped finds ``None`` rather than the PREVIOUS
tick's heads. ``prepare_duplex_tick`` clears it again at the top of
every tick as a second, independent guarantee of the same thing.

Fails loud (never silently wrong) when the slot does not hold the
expected per-head tuple.

First half of a duplex tick: fuse this tick's audio frame with the
session's own previous text/function tokens
(:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.fuse_stt_step_embeds`)
and park the result on ``session.request.pending_embed_override``.

``frame_embed``: ``(1, hidden_size)`` — ONE frame's already-encoded
audio embedding (the same thing one row of
``mm_bindings["audio"].encode(...)``'s output is). CPU tensor.

Note this does NOT advance ``session.frame_seq`` — :func:`finish_duplex_tick`
does, once the step has actually run. That ordering is what lets the
duplex lane roll a prepared-but-not-scheduled tick back cleanly
(clear the override, no frame consumed, retry next tick) instead of
silently dropping a frame of the user's audio.

**Injection ticks are ordinary frame ticks with a wider row.** When the
request carries queued ``pending_context_token_ids`` (the tool-response
resume), the scheduler serves it a ``1 + N`` token row whose
ids are ``[last_sampled_token, *context]``
(``_batch_build_materialize``'s ``RowKind.INJECT`` branch). Slot 0 is
the position the request still owes for its own pending sample — and in
THIS model that position is a frame-lockstep position like any other,
so it gets the SAME fusion every other tick's row gets, not a bare token
embedding. Building slot 0 without the fusion would drop that tick's
audio frame AND feed a region-3-shaped bare token embedding into a
sequence this module documents as region 2 forever — see the module
docstring. Design doc §7.9.15.

**The context tokens that follow it are fused too** (design doc
§7.64.3). They used to get plain ``model.embed_input_ids`` embeddings,
justified by parity with ``EngineSttStep.inject_tokens`` -> ``_extend``.
That is parity with arbi-serve's OWN legacy path, not with the
checkpoint: §7.31.1 measured what a bare token embedding costs this
model — the omitted ``cos(embed(tok), fused)``
averages 0.30 over random vocab rows, i.e. very nearly orthogonal to
anything the model was trained to read. §7.31.1 fixed exactly this for
the seed prompt (``render_duplex_seed``) and the fix was never carried
here, so every ``<TOOL_RESPONSE>`` token — the last ~30-60 positions
before the model must speak the answer — arrived as noise. These rows
now use the same ``fuse_stt_step_embeds`` shape the seed does.

The override is an inference tensor, matching
:func:`arbi_serve.runtime.forward_exec.run_model_forward` and its boot
compile warmup. Tensor dispatch keys are part of Dynamo's guards, so this
boundary must not produce an ordinary no-grad tensor for a compiled
inference forward.

Feed this tick's function-channel token through the connection's
incremental ``<TOOLCALL>[...]</TOOLCALL>`` detector; return the calls
whose block closed on THIS token.

Runs on the engine thread, inside ``run_forever``'s loop body, which
is why it is deliberately only this: one single-token detokenize plus
a substring scan over a short buffer. Nothing here blocks, allocates
on the GPU, or reaches the WS side — the 10 s
``tool_bridge.wait_for_result`` that a tool call eventually implies
stays entirely on the WS/session side (design doc §7.9.15's own
statement of why detection must move here and execution must not).

A detection while a call is ALREADY pending is dropped with a debug
log rather than started: one ``<TOOLCALL>`` block can already name
several calls, so a second block arriving before the first resolves is
the model talking over its own pending call, not a second intent to
satisfy — and starting it would leave two overlapping pause/resume
cycles racing on one gate. The detector is still fed, so its buffer
stays consistent for the block after this one.

This tick's forced ``ack_messages`` token, or ``None``.

``getattr``, not a bare call: minimal session stand-ins in CPU tests
predate :meth:`NemotronVoiceChatDuplexSession.next_ack_token` and must
keep getting the plain PAD mute (design doc §7.9.18).

Second half of a duplex tick, once the step has run.

Steps, in order (mirrors ``EngineSttStep._advance_frame_step`` +
``run_turn``'s Site-3 body, ported onto the session object per
§7.7.2/§7.7.4):

  1. TAKE both logits heads off this tick's own request
     (``Request.logits_heads``, already row-sliced by the forward —
     design doc §7.42). No batch-wide tensor and no row index cross
     this boundary, so there is nothing for a caller to re-derive
     wrongly against a slate the scheduler may have reordered.
  2. Take the text token the engine's OWN sampler already committed
     (``req.output_token_ids[-1]``) — so this tick's feedback is the
     token actually generated, honoring this request's real sampling
     params, not a re-derived argmax. Falls back to argmax only on
     the very first tick, before any token has been committed.
  3. Function-call detection + the FC pause gate (design doc §7.9.16)
     — see below.
  3.5. The duplicate-BOS guard, then the EOU-gated BOS admission check
     (design doc §7.30) — see below. Both run only on whatever
     ``text_token_id`` steps 3's gates left behind, so a tick the FC
     or speak gate already forced to PAD/an ack token never reaches
     either.
  4. The Site-3 BOS/EOS ``agent_idle`` switch — mirrors
     ``nemotron_voicechat_turn.run_turn``'s per-tick body
     exactly — extended with the TTS
     ratio-cap watchdog (design doc §7.23): a turn already
     in progress that is producing audio far faster than new text
     content gets its EOS forced here, before anything downstream
     sees this tick's token.
  5. ``session.observe(...)`` — the frame-lockstep feedback update,
     reading whatever step 4 left in ``text_token_id`` (a forced EOS
     included), so the next tick's fusion conditions on it exactly
     like a natural one. This is ALSO where the decided token is
     reconciled back into ``req.output_token_ids`` — see
     :meth:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
     .NemotronVoiceChatDuplexSession.observe`. Every one of the seven
     rewrite gates below funnels through this single call, which is
     exactly why the reconciliation lives there and not here: a gate
     added later cannot forget to do it.
  6. ``session.consume_frame()`` — advance the frame counter, now
     that the frame has genuinely been consumed by a real forward.

``decode_function_token``: ``(token_id) -> str``, this checkpoint's
detokenizer for ONE function-channel token (the adapter binds
``tokenizer.decode``; mirrors ``EngineSttStep.decode_function_token``,
which is what the turn-based path feeds its own detector from).
``None`` disables function-call detection entirely, which is what
every non-FC caller (the pull driver, the parity tests) wants — the
gate then never engages and this function behaves exactly as before.

**The FC pause gate.** Steps 4-6 all read ONE value,
``text_token_id``, and it is also what the caller hands to TTS as this
tick's subword. So forcing it to ``session.text_pad_id`` while a call
is pending is the whole gate: the agent's text goes blank, its
synthesized audio goes to the same "listening" frame a silent tick
produces, and the Site-3 switch cannot fire a spurious turn open/close
off a token the client will never hear — while the 80 ms grid, the
audio-in fusion and the function channel all keep running untouched.
That is the reference's ``_fc_in_progress`` semantics (design doc
§3.2.1: keep the grid advancing, mute the text channel), and it is
deliberately NOT "stop stepping", which has no duplex meaning.

**The speak gate** (design doc §7.24) forces the SAME value through
the SAME three consumers, on an entirely independent trigger:
``session.speak_gate_closed``, set by a caller that owns turn
boundaries (a repointed bounded session) rather than by tool-call
detection. It cannot reuse the FC gate's own flag verbatim because the
two must be able to be active at once without either clobbering the
other's own resume — a tool call can be pending while the caller is
ALSO still withholding permission to speak, or vice versa. Checked
only when the FC gate is NOT active: an in-flight tool call's own
ack-phrase-or-PAD forcing already produces the identical muted
text-channel effect the speak gate would, so there is nothing left
for it to add on a tick the FC gate already owns, and giving the FC
gate priority keeps its ack-phrase forcing intact rather than
silently overridden.

**The ack phrase rides the same seam.** When the connection has an
``ack_messages`` phrase queued (``session.queue_ack_tokens``, fed from
the WS side's tool definitions), the gate forces that phrase's NEXT
token instead of PAD — one id per 80 ms frame, until the phrase runs
out and PAD resumes. Because all three consumers read the one value,
the agent genuinely SPEAKS the phrase (the TTS side is a bare id
passthrough for this checkpoint) and it lands in the transcript, with
no second TTS thread and no change to this function's contract. The
forced ids also enter :meth:`observe`, which is the correct half of
the choice: the user heard those words, so the next frame's fusion
must condition on them. Design doc §7.9.18.

**The babble watchdog is independent of sampling.** Design doc §7.23:
a turn's own token stream is the ONLY thing that has ever closed it in
duplex mode, and a turn whose sampled tokens never happen to include
EOS runs forever on that alone. ``NemotronVoiceChatDuplexSession
.note_babble_frame`` tracks audio frames spoken against genuinely new
text tokens for the CURRENT turn and reports when the ratio crosses
:attr:`~NemotronVoiceChatDuplexSession.tts_ratio_cap`; when it does,
this tick's ``text_token_id`` is overwritten to ``text_eos_id`` before
anything downstream (the recurrence, the transcript, the TTS subword)
ever sees the token the model actually sampled. A client cannot tell
a watchdog-forced EOS from a model-emitted one — both take the exact
same path from here on.

**The duplicate-BOS guard.** Design doc §7.26. A second BOS while the
agent is already speaking (``not session.agent_idle``) is rewritten to
``text_eos_id`` before the Site-3 switch below ever sees it, so this
tick closes the existing turn instead of silently re-arming it: the
unguarded alternative calls ``mark_agent_active()`` on an already-open
turn, which resets the babble watchdog's counters without ever
transitioning ``agent_idle`` back to ``True`` — the turn never closes,
the WS bridge's synthetic ``response.done`` (keyed off exactly that
transition) never fires, and a re-armable connection can never stage
its next turn (``DuplexConnection.begin_turn`` requires the previous
one to have ended). Checked first, so a rewritten value flows through
the ordinary BOS/EOS branches below like any other tick's token.

**The EOU-gated BOS admission check.** Design doc §7.30. By
the time this runs, the duplicate-BOS guard above has already
rewritten an "already speaking" BOS to EOS — so a ``text_token_id``
that is STILL ``text_bos_id`` here means ``session.agent_idle`` was
genuinely ``True``: a real candidate to open a NEW turn, not a
duplicate of the current one. ``NemotronVoiceChatDuplexSession
.eou_admits_bos()`` decides whether to honor it: always yes for a
session's very first turn (a proactive greeting is legitimate — see
that method's own docstring for why the reference never gates this
either), but every turn after the first needs real confirmed user
speech since the agent's last turn opened — the confirmation coming
from the checkpoint's own RNNT endpoint detector where available, the
acoustic VAD otherwise (design doc §7.32) — the live evidence
(design doc §7.27) is that without this, the model reopens with a
self-invented topic within seconds of a turn it just closed, with the
user genuinely silent throughout. That is the WHOLE gate: it does not
also make the model wait out the FORCE-open's silence timer below
(design doc §7.51 — that clause could never refuse a turn, only delay
it to the tick the force-open would have opened it on anyway, which
left this checkpoint's own turn-taking head with no say at all). A
refused candidate is rewritten to
``session.text_pad_id`` — the same "nothing happened this tick" value
the FC/speak gates above already use, not ``text_eos_id`` (there is no
open turn here to close) and not left as ``text_bos_id`` (which would
open the very turn this check just refused). The client hears ordinary
silence for that tick, identical to the agent genuinely listening;
``DuplexTickResult.eou_gate_closed`` records whether this fired, for
observability only — nothing downstream branches on it.

**The EOU FORCE-open** (design doc §7.37, extended to the bounded path
by §7.44) is the same mechanism's additive half, applied to the
opposite case: a tick whose text channel produced something OTHER than
BOS, on a connection whose own turn-boundary signal says the agent owes
an answer (``NemotronVoiceChatDuplexSession.eou_forces_bos()`` — the
VAD/EOU confirmed-speech-then-silence rule in genuine duplex, the
client's committed-audio-consumed rule on a client-owned turn), opens
the turn anyway. The veto half alone can only ever remove a turn-open
the model proposed; since the duplex text channel is sampled greedily
and this checkpoint's logits keep BOS a persistent runner-up to A forced open is written as
``text_bos_id`` and flows through the SAME Site-3 switch and recurrence
a model-sampled BOS does, so nothing downstream distinguishes them;
``DuplexTickResult.eou_forced_open`` records it for observability.

Drive ONE real scheduler tick for ``session``'s persistent duplex
``Request`` — the design doc §7.7.6/task-44a PULL stepper.

This is the caller that owns the whole step: it calls
:func:`prepare_duplex_tick`, then ``eng.scheduler.schedule(_now=now)``
+ ``run_step_async`` itself (the SAME step-execution path
``tests/test_nemotron_voicechat_embed_override_live_gpu.py`` already
proved correct for an override-carrying duplex request, reused
verbatim), then :func:`finish_duplex_tick`.

**It must not be used against an engine whose own ``run_forever`` loop
is running** — that loop calls ``scheduler.schedule()`` unconditionally
on every iteration and will independently pick up and drive the same
duplex request, double-stepping its KV/Mamba state out of sync with
this function's own per-tick protocol (design doc §7.8). Use
:mod:`arbi_serve.realtime.duplex_lane` for a real, running engine;
this function is for standalone/offline drivers and tests that own
the engine exclusively.

Precondition (violated -> ``RuntimeError``, fails loud, never silently
wrong): ``session.request`` must already be admitted
(``new_duplex_request``/``admit_duplex_request``) and must actually be
picked up by ``eng.scheduler.schedule(_now=now)`` this tick — either
because it is already ``DECODING`` in ``eng.scheduler.running``, or
because it is still ``WAITING`` and this is its promotion tick
(``Scheduler._build_duplex_slate``'s "Source 1" fast path).

``audio_frame_embeds``: ``(num_frames, hidden_size)`` — the SAME shape
``EngineSttStep.__init__``'s own ``mm_bindings["audio"].encode(...)``
output has; the frame at ``session.frame_seq`` is consumed per call.

Real ``SttStepFn`` implementation for
:func:`arbi_serve.runtime.nemotron_voicechat_turn.run_turn` — the injected
STT decode-step callback that module's own docstring scopes OUT of itself
(see that docstring's "Why the STT step is an INJECTED callback" section).

:class:`EngineSttStep` drives
:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.forward`
against a hand-built ``ScheduledBatch``/``MultiStatePool``/``attn_ops`` —
the same real, CUDA/Triton-only backend kinds
(``paged_kv:tkv-bypass``/``mamba:mamba``) as
``tests/test_nemotron_voicechat_parity_live_gpu.py``'s
``_build_batch_and_pool``, extended one step at a time instead of that
test's single fixed-length teacher-forced pass. This module is therefore
CUDA-only and NOT exercised by this repo's CPU-only test suite (see
``run_turn``'s own module docstring for why no CPU backend exists for
either state kind) — it is verified only by a live-GPU pass; the CPU-only
realtime tests instead inject a fake ``SttStepFn`` (mirroring
``tests/test_nemotron_voicechat_turn.py``'s own convention). The one piece
of this module's own logic that IS CPU-testable without CUDA —
:class:`_FrameLockstepState`, the frame-vs-plain-decode bookkeeping below
— has its own direct unit tests in ``tests/test_nemotron_voicechat_stt_step.py``.

=== Frame-lockstep audio-in stepping ===

The reference (NVIDIA's NeMo checkpoint,
``nemo.collections.speechlm2.models.duplex_stt_model.DuplexSTTModel
._init_inference``/``_step_zero``/``_step_inference``) does NOT encode a
whole audio clip and splice it in as one block of placeholder-token
positions the way an ordinary multimodal LLM does. It runs exactly
``T = system_prompt_len + num_audio_frames`` autoregressive backbone
steps — one step per audio frame (after the system prompt) — where every
step's input embedding is the FUSION
(:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel
.fuse_stt_step_embeds`) of that frame's continuous embedding with the
embedding of the PREVIOUS step's own emitted text-channel token (mostly
``text_pad_id`` — the "blank" the reference's own
``DuplexSTTModel.text_pad_id`` docstring names for "frames when the model
is not speaking") and the previous step's own emitted function-channel
token. The model was trained to condition on its OWN mostly-blank
per-frame history interleaved with audio; feeding it one contiguous block
of placeholder tokens followed by nothing (the shape an ordinary
multimodal splice produces) is an out-of-distribution input the model was
never trained to decode from.

This class therefore drives THREE regions per turn, all through the SAME
:meth:`__call__`/:meth:`_extend` machinery:

  1. **System prompt** (``__init__``): one ordinary multi-token prefill,
     plain ``embed_tokens(input_ids)`` — ``prompt_token_ids`` carries no
     audio placeholder ids. This is a deliberate, explicitly-scoped
     simplification of the reference (which fuses PAD-channel constants
     into the prompt region too, via the same per-position formula) —
     the prompt itself carries no audio-frame term either way, and the
     fusion weights are 1.0 on the dominant (real-token) term, so this
     does not reproduce the reference BIT-EXACTLY over the prompt
     region — only over the audio-frame region, which is the one that
     determines whether the model's output is real speech/text or
     near-immediate EOS.
  2. **Audio-frame region** (:class:`_FrameLockstepState`-driven, inside
     :meth:`__call__`): exactly ``num_audio_frames`` steps, each one
     fusing (this frame's continuous embedding, the previous step's own
     sampled text-channel token id [seeded ``text_pad_id``], the previous
     step's own GREEDILY sampled function-channel token id [also seeded
     ``text_pad_id``] — greedy per the reference's own
     ``_step_inference``/``_step_zero``, "function channel always uses
     greedy, not affected by text sampling params"). Both heads are
     sampled every step; most steps land back on ``text_pad_id`` (the
     model is still "listening") and that blank is what feeds the NEXT
     step — this is correct/expected, not a bug. A step predicting the
     turn-stop token (:func:`resolve_turn_eos_ids`) WHILE frames remain
     does not end the turn either
     (:meth:`_FrameLockstepState.should_stop_turn`) — it means "pause
     speaking, keep listening," not "conversation over" (matching the
     production Triton backend's own per-step driver, which never stops
     decoding on this signal — see that method's own docstring for the
     citation); the raw sampled id still feeds forward as the next
     step's text-channel term either way.
  3. **Post-frame region** (once :attr:`_FrameLockstepState.frames_remaining`
     is False): ordinary autoregressive decode, plain token embedding,
     one head fed forward. The reference has no post-frame case at all
     (its own ``offline_inference`` is a fixed-``T`` batch scorer); a
     real-time turn must keep generating past the input utterance's own
     frame count, so this region is this integration's own
     extrapolation, not a reference port.

Two extra methods beyond the bare ``SttStepFn`` callable contract:

  * :meth:`inject_tokens` — feeds extra context into the running
    sequence WITHOUT sampling, used by
    :mod:`arbi_serve.realtime.nemotron_voicechat_turn`'s tool-call
    pause/resume flow to splice a ``<TOOL_RESPONSE>[...]</TOOL_RESPONSE>``
    chunk in between two ordinary decode steps. Always a plain-token
    extend (region 3's shape) — tool-call syntax only ever appears in the
    model's own emitted text, which (region 2's blank-heavy cadence)
    realistically only happens once the audio-frame region is at or near
    exhausted.
  * :meth:`decode_function_token` — detokenizes the CURRENT step's
    function-channel token, taken GREEDILY (argmax), which is the same
    id :meth:`__call__` feeds forward as region 2's fusion feedback.
    Greedy is not a choice this method gets to make: the reference's own
    ``_step_inference``/``_step_zero`` document "function channel always
    uses greedy, not affected by text sampling params", and the id that
    enters the recurrence is therefore always the argmax one. Re-sampling
    these logits under the TEXT channel's knobs (which is what this
    method used to do) made the detector read a token the model never
    emitted whenever a client set ``temperature`` > 0 via
    ``session.update`` — it could manufacture a ``<TOOLCALL>`` the model
    did not produce, or miss one it did. The duplex path never had this
    defect (``nemotron_voicechat_duplex_adapter._decode_function_token``
    decodes the already-committed greedy id and says why); this is the
    turn-based path brought into line with it. Design doc §7.38.

Look up ``token_str`` in the served tokenizer's own vocab, raising
loud (not silently falling back) if it's absent — shared by
:func:`_resolve_text_pad_id` and :func:`resolve_turn_eos_ids`, which
need the SAME lookup for different token strings.

This checkpoint's real text-channel PAD id
(:data:`arbi_serve.models.nemotron_voicechat.TEXT_PAD_TOKEN_STR`), the
"blank" fed for the text/function channel at every audio-frame step
before either channel has emitted a real token this turn (see module
docstring's "Frame-lockstep audio-in stepping" section, region 2).

This checkpoint's real turn-stop id set for
:class:`EngineSttStep`'s ``eos_token_ids`` — see
:data:`arbi_serve.models.nemotron_voicechat.TEXT_EOS_TOKEN_STR`'s own
docstring for why this is a single, EXPLICITLY resolved id
(``</s>``) rather than the tokenizer's generic ``stop_token_ids``
(which, for this specific checkpoint's bundled tokenizer artifact,
wrongly includes :data:`~arbi_serve.models.nemotron_voicechat
.TEXT_PAD_TOKEN_STR`'s own id — using it as a turn-stop condition
would end every turn on the first blank/listening frame). A caller
building ``prompt_token_ids`` from the SAME tokenizer's
``apply_chat_template``/``encode`` is unaffected — this only changes
what :class:`EngineSttStep` treats as "the model is done for this
turn".

This checkpoint's real turn-BOS id — the Site-3 BOS side of
:func:`resolve_turn_eos_ids` (design doc
``docs/nemotron-voicechat-duplex-design.md`` §3.2.1): the agent
TEXT-channel token id that flips
:meth:`~arbi_serve.realtime.nemotron_voicechat_duplex_session
.NemotronVoiceChatDuplexSession.mark_agent_active` when sampled.
Mirrors :func:`resolve_turn_eos_ids`'s exact lookup pattern (a single
explicitly resolved id via :func:`_resolve_special_token_id`, not the
tokenizer's generic ``bos_token_id``) — see
:data:`arbi_serve.models.nemotron_voicechat.TEXT_BOS_TOKEN_STR`'s own
docstring for why this checkpoint needs an explicit constant rather
than a generic tokenizer property.

Pure Python bookkeeping for one turn's audio-frame-fusion decode
region (module docstring, region 2) — no CUDA, no tensors beyond the
frame-embedding tensor it indexes into; this is the one piece of
:class:`EngineSttStep`'s own logic ``tests/test_nemotron_voicechat_stt_step.py``
exercises directly on CPU.

``prev_text_token_id``/``prev_function_token_id`` start at
``text_pad_id`` (the "blank" — see :func:`_resolve_text_pad_id`) and
are updated via :meth:`observe` after each step's own sampling, so
the NEXT :meth:`consume_frame` call's fusion input reflects what the
model itself just emitted (mostly still blank).

Stateful ``SttStepFn`` — a callable class instance satisfies the
plain ``Callable[[int], SttStepResult]`` contract
(:data:`~arbi_serve.runtime.nemotron_voicechat_turn.SttStepFn`) without
``run_turn`` needing to know anything about the extra methods below.

Owns a PRIVATE, turn-scoped ``MultiStatePool`` (never the shared
engine pool — this loop is explicitly not wired into the engine
scheduler, see ``run_turn``'s module docstring), so on barge-in
abandonment there is no explicit page-free/deregistration protocol
to run: the pool and its backing CUDA tensors are reclaimed by
ordinary Python refcounting once the owning
:func:`arbi_serve.realtime.nemotron_voicechat_turn.stream_nemotron_voicechat_turn`
generator (which holds the only reference) is closed/collected —
see that module's docstring for the barge-in cleanup finding this
class's design was chosen to keep true.

Should THIS step's ``SttStepResult.is_eos`` actually end the
turn?

A text-channel turn-stop token predicted WHILE frames remain
does NOT end the turn — it means "pause speaking, keep
listening," not "conversation over" (module docstring's "Frame-
lockstep audio-in stepping", region 2 — matches the production
Triton backend's own per-step driver, which never stops
decoding on this signal, only silences that frame's TTS output).
``sampled_turn_stop_token`` is the caller's own raw
``token_id in eos_ids`` check (:func:`resolve_turn_eos_ids`) —
kept as a plain bool argument (not re-derived here) so this
class stays token-id-agnostic.

Build the PERSISTENT ``attn_ops``/metadata-builder objects the
decode-shaped path (:meth:`_decode_forward`) reuses across every
call for the turn's life, and initialize this turn's CUDA-graph
capture bookkeeping (see
``arbi_serve.runtime.capture.nemotron_voicechat_decode``'s module
docstring for why the turn loop's steady-state decode step needs
its OWN small capture path rather than the shared engine one).

Persistent (not rebuilt per call, unlike the prefill/injection
path's own objects in :meth:`_extend`) because a captured graph's
kernels reference these objects' internal scratch/dispatch state
by address — the warmup forwards that precede a capture and the
capture itself must run through the SAME objects, and every
later replay must too.

One decode-shaped (exactly one new token) forward — region 2
(``inputs_embeds`` already the fused audio-frame embedding) or
region 3 (ordinary decode; looked up here via ``embed_tokens`` so
BOTH regions present the captured graph the identical shape: an
``inputs_embeds``-bypass forward, never the ``embed_tokens``
lookup itself — see
``arbi_serve.runtime.capture.nemotron_voicechat_decode``'s module
docstring). Updates ``self._pending_logits``/``self._all_token_ids``/
``self._next_pos`` exactly as :meth:`_extend` does — callers don't
need to know whether this particular call went eager, compiled, or
captured-replayed.

One forward pass appending ``token_ids`` (or, when
``inputs_embeds`` is given, that precomputed embedding — see
below) to the running sequence; caches the LAST position's
``(text_logits, function_logits)`` for the next
:meth:`__call__`/:meth:`inject_tokens` to consume. Shared by the
initial prompt prefill, an ordinary one-token decode step
(region 3), an audio-frame-fused decode step (region 2), and a
multi-token tool-response injection alike — see class/module
docstring.

``inputs_embeds``: ``(len(token_ids), hidden_size)`` — when
given, bypasses this model's own ``embed_tokens`` for this
forward (see :meth:`NemotronVoiceChatBackboneModel.forward`'s
own docstring); used for region-2 steps, whose real input is a
FUSED embedding (:meth:`~arbi_serve.models.nemotron_voicechat
.NemotronVoiceChatModel.fuse_stt_step_embeds`), not any single
token's embedding row. ``token_ids`` is still required in this
case — its VALUES are unread by the forward, but they still
become this position's bookkeeping entry in ``self._all_token_ids``
(see :meth:`_advance_frame_step` for what id region-2 records
there).

``allow_compile``: the initial prompt prefill and a tool-response
injection are each a ONE-OFF, variously-sized shape (never
repeated within a turn), so compiling them buys nothing and is
forced eager unconditionally (default False). The steady-state
one-token decode step (region 3, ``allow_compile=True``) and the
audio-frame-fused decode step (region 2, also ``allow_compile=
True`` — a second, likewise-repeated fixed shape: same
single-position ``block_table``/metadata construction below,
differing only in ``inputs_embeds is not None``) are each the
turn's hottest, most-repeated shape for their own region and DO
stay compiled — see ``block_table``'s fixed-width construction
below, the change that makes that possible: with
the previous variable-width ``[self._page_ids]`` table, every
single decode step was ALSO a distinct shape (Dynamo's own guard
system reports it as a data-dependent-expression guard on the
table width), and Dynamo's ``fullgraph=True`` recompile limit
(default 8) blew a real turn's steady-state loop up around the
8th token — GPU-verified (``torch._dynamo.exc.FailOnRecompileLimitHit``,
``arbi-serve`` boot log, live 11B checkpoint). Padding the table
out to the instance's own fixed ``self._max_pages`` ceiling (never
exceeded — ``_ensure_pages`` raises loud on overflow instead of
growing past it) makes every one-token decode step's ``block_table``
BIT-IDENTICAL IN SHAPE, so Dynamo compiles it exactly once and
replays the same graph for the rest of the turn. The unused
tail entries point at page 0, the pool's own permanent
null-sentinel page (see ``MultiStatePool``/``PagedKVStatePool``'s
module docstrings) — never read, since every kernel bounds its
block-table walk by ``seq_lens``, which never reaches the padded
region.

``attn_ops``/``metadata_builders``: when given (the decode-shaped
path — see :meth:`_decode_forward`), reuse these PERSISTENT
objects. Required for CUDA-graph capture correctness: the captured
graph's kernels reference these objects' own persistent scratch/
dispatch state, so the warmup forwards that precede a capture
(this method, called eagerly a couple of times first) and the
capture itself must run through the IDENTICAL objects, not
equivalent-but-distinct ones. When omitted (the prefill/injection
default below), ``metadata_builders`` is rebuilt fresh per call
(each one-off shape needs its own correctly-sized scratch) while
``attn_ops`` falls back to ``self._decode_attn_ops`` — see the
fallback's own comment for why it must NOT build fresh op objects.

Feed ONE audio frame's fused embedding as the next decode
step's input, priming ``self._pending_logits`` for that position
— module docstring region 2, the reference's per-frame
``AddFusion`` recurrence (``DuplexSTTModel._step_inference``/
``_step_zero``). Advances ``self._frames``; callers check
:attr:`_FrameLockstepState.frames_remaining` first (``__init__``
for frame 0, :meth:`__call__` for every later frame).

Detokenize this step's GREEDY function-channel token from
``function_logits`` (this step's raw logits, off
``SttStepResult.function_logits``) — see module docstring.

Argmax, deliberately: this is the same id :meth:`__call__` feeds
forward into region 2's fusion, so the detector reads exactly the
token that entered the model's recurrence.

NemotronLabs-VoiceChat-11B per-turn STT+TTS lockstep orchestration loop.

Drives :class:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel`
to produce one real turn: STT text/function decode, THEN one TTS frame per
STT step (lockstep), THEN codec decode to a waveform chunk — porting the
shape of the reference NIM container's
``nemo.collections.speechlm2.models.nemotron_voicechat.NemotronVoiceChat
.offline_inference`` (``t in range(1, T)`` loop body, lines ~687-727 of the
extracted reference) and
``nemo.collections.speechlm2.models.duplex_ear_tts.DuplexEARTTS
.infer_codes_one_step``/``decode_one_audio_step``.

=== Why the STT step is an INJECTED callback, not built here ===

The composite model's own docstring says a turn-orchestration loop drives
:meth:`NemotronVoiceChatModel.forward` "through the ordinary engine
scheduler" — but that backbone is NOT wired into arbi-serve's ``Engine``/
``Scheduler`` (its two-head ``(text_logits, function_logits)`` return
doesn't fit ``LayerStackModelMixin._lm_head``'s one-head contract; see
``NemotronVoiceChatBackboneModel.forward``'s own docstring), and the
approved project plan explicitly SKIPS promoting the STT half into the full
scheduler for multi-tenant concurrency (``docs/omni_serving.md`` / the
project plan's "M4" item). The only place a raw ``forward()`` call over
hand-built ``ScheduledBatch``/``MultiStatePool``/``attn_ops`` currently
exists for this model is
``tests/test_nemotron_voicechat_parity_live_gpu.py``'s
``_build_batch_and_pool`` — and every backend kind it uses
(``paged_kv:tkv-bypass``, ``mamba:mamba``) is CUDA/Triton-only (verified:
``arbi_serve/backends/__init__.py``'s lazy backend registry has no CPU
variant of either), so that plumbing cannot run on CPU at all, let alone in
a CPU-only unit test.

So this loop treats "run one STT decode step and produce this step's
token/text/function-logits/EOS-flag" as an injected dependency
(:data:`SttStepFn`) rather than reimplementing incremental
prefill-then-decode batch-extension (growing ``block_table``/
``slot_mapping``/positions/KV-pool state one token at a time) — that is
real ``arbi_serve/engine/active.py``-shaped machinery, out of scope to
duplicate here, and a live caller wires it using exactly the
``_build_batch_and_pool``-style construction the parity test already
established (extended, step to step, the way ``arbi_serve/engine/active.py``
extends a live request). What THIS module owns — and what actually needs
designing, not reusing — is the LOCKSTEP SEQUENCING: one TTS frame per STT
step, the turn-lifecycle state threading (``past_key_values``/``prev_code``),
the composite model's documented TODOs (see below), and the stopping
condition. That part is fully exercised on CPU with synthetic random/meta
weights (``tests/test_nemotron_voicechat_turn.py``).

=== Lockstep coupling: what the reference source actually does ===

The composite model's docstring (``arbi_serve/models/nemotron_voicechat.py``)
frames the STT→TTS text coupling as "sample text from the STT backbone's
text_logits, decode it to a string, and RE-TOKENIZE that string into the TTS
side's own subword vocabulary". The reference's actual per-step loop body
(``NemotronVoiceChat.offline_inference``, lines ~692-711) does NOT do that:

    current_subword_id = inference_state["gen_text"][:, t].unsqueeze(-1)
    ...
    code, past_key_values = self.tts_model.infer_codes_one_step(
        current_subword_id=current_subword_id, ...)

The STT backbone's OWN generated token id is passed straight into the TTS
side as ``subword_ids`` — no detokenize/re-tokenize round trip anywhere in
the loop. This is consistent with (not a coincidence): the checkpoint's own
shape truth has ``NemotronEarTTSConfig.subword_vocab_size == 131072``,
IDENTICAL to the STT backbone's ``vocab_size`` in ``_BACKBONE_SHAPE`` — the
strong implication (not independently confirmed against a real tokenizer
artifact, same caveat as ``AUDIO_PATCH_ID``) being that STT and TTS share
one subword vocabulary for this checkpoint, so "re-tokenization" is a no-op
identity pass-through. This module follows the REFERENCE'S ACTUAL CODE (id
pass-through) rather than the docstring's more cautious prose, and flags the
discrepancy here for that docstring to be corrected in a later pass.

One more reference detail that falls out of this: ``infer_codes_one_step``
also threads a ``prev_subword_id`` (the STT token from step t-1), but ONLY
to build ``context_hidden_state = self.embed_tokens(prev_subword_id)`` when
``context_hidden_size is not None`` — and this checkpoint's
``context_hidden_size`` IS ``None`` (see the composite docstring's own
"Why generate_tts_frame does NOT take a raw STT hidden-state tensor"
section), so ``prev_subword_id`` is dead code for VoiceChat-11B. Consistent
with that: :meth:`NemotronVoiceChatModel.generate_tts_frame` was built with
no ``prev_subword_id``/``context_hidden_state`` parameter at all — this loop
needs only the CURRENT step's token id, never the previous one, for the TTS
coupling.

=== The four composite-model TODOs, as resolved here ===

  1. ``set_char_map`` — called once, at the top of :func:`run_turn`, using
     either a caller-supplied ``subword_id_to_char_ids`` map or one built
     from a caller-supplied ``vocab`` dict via
     :func:`build_subword_char_map` (a direct, small port of the
     reference's ``ear_tts_model.build_vocabs`` — see that function's
     docstring for the one thing it deliberately drops: the vocab-dir file
     cache, irrelevant here since arbi-serve's own tokenizer is already
     loaded once at process boot, not per-turn).
  2. First-frame ``prev_code`` — seeded from ``model.tts_codec_silence_tokens``
     (the composite's own documented fallback) unless the caller passes
     ``initial_prev_code`` (e.g. a ``prepare_nemotron_voicechat_voice.py``
     custom-voice bundle's ``prompt_codes[:, -1:]``, for closer reference
     fidelity — see the composite docstring's ``guidance_enabled`` constancy — enforced BY CONSTRUCTION: this loop
     takes one ``guidance_scale`` parameter, derives
     ``guidance_enabled = bool(guidance_scale)`` exactly once, and reuses
     that same derived flag (and the same ``guidance_scale`` value) for
     :meth:`~NemotronVoiceChatModel.prime_tts_turn` and every
     :meth:`~NemotronVoiceChatModel.generate_tts_frame` call for the whole
     turn — there is no second knob that could desync mid-turn.
  4. Control-code sanitization — ``sanitize_control_codes`` is forwarded to
     :meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.generate_tts_frame`,
     which applies
     :func:`~arbi_serve.audio.nemotron_audio_codec.replace_control_speech_codes`
     (a direct port of the reference's own, ``duplex_ear_tts.py``, the
     ``silence_tokens is not None`` branch — the only branch reachable
     here, since ``NemotronVoiceChatModel`` always carries
     ``tts_codec_silence_tokens``) to the codes entering the vocoder, NOT
     to the code fed back as next-step ``prev_code`` — matching the
     reference, which only sanitizes at the decode call site
     (``decode_one_audio_step``) and keeps the raw generated code in the
     autoregressive history. This used to live here as a post-hoc
     "``.any()`` probe, then re-decode if it fired" retrofit; that shape
     costs a host sync every tick and cannot coexist with the streaming
     vocoder cache (design doc §7.20) — a second decode of the same frame
     would advance the cache twice. NOTE:
     ``torch.isin`` treats ``tts_control_codes`` as an unordered set, so
     the composite docstring's unconfirmed guess at its 3-entry ORDERING
     doesn't matter for this use — but for the record, the reference
     source (``duplex_ear_tts.py`` line ~118) shows the real order is
     ``[speech_bos_id, speech_eos_id, speech_pad_id] == codebook_size +
     {2, 1, 0}``, not the composite docstring's hypothesized
     ``codebook_size + {0, 1, 2}`` (pad, bos, eos) — a correction worth
     folding back into that docstring later.

  A fifth, smaller thing the reference does that the composite docstring
  doesn't mention at all: ``infer_codes_one_step``'s
  ``inference_force_speech_silence_on_eos`` behavior — when the current
  step's STT token is the text-channel EOS, the code FED INTO that step
  (``prev_audio_tokens``) is force-replaced with the silence frame before
  the TTS backbone runs, so the model's speech naturally trails into
  silence at end-of-turn rather than conditioning on whatever it happened
  to generate last. Small and well-scoped, ported directly (see the
  ``is_eos`` branch in :func:`run_turn`).

=== Termination ===

The reference's own ``offline_inference`` always runs a FIXED ``T`` steps
(batch/offline scoring shape: trims the output by tracked lengths
afterward) — not suited to an interactive per-turn loop. This loop instead
stops as soon as :class:`SttStepResult.is_eos` is set (after emitting one
final forced-silence TTS frame — see above) or ``max_tokens`` steps have
run, mirroring ``arbi_serve/realtime/turn.py``'s ``stream_turn`` own
``client.finish_reason is not None`` early-stop convention rather than the
reference's fixed-length batch shape.

``max_tokens`` counts EVERY ``stt_step_fn`` call this loop makes — this
module has no concept of "audio frame" at all (see the module docstring's
scope note: that's entirely :class:`~arbi_serve.runtime
.nemotron_voicechat_stt_step.EngineSttStep`'s concern). When the injected
``stt_step_fn`` is an :class:`EngineSttStep` driving real audio-in
(``mm_feats`` given), a fixed, input-duration-INDEPENDENT number of its
early calls are mandatory audio-frame-fusion "listening" steps (that
class's own module docstring, region 2) — :attr:`EngineSttStep
.num_audio_frames` of them — during which ``should_stop_turn`` can never
end the turn by design. If ``max_tokens`` is not large enough to cover
BOTH that mandatory listening region ``tool_call.wav``
finding: 84.46 s of real audio + the realtime protocol's mandatory 8 s
trailing silence is 1157 audio-frame steps at this checkpoint's 80 ms
frame grid — already past the ``max_tokens=1024`` default before a single
generation step could run). :func:`audio_turn_max_tokens` is the fix: a
caller driving an :class:`EngineSttStep` with real audio MUST inflate the
``max_tokens`` passed to THIS function (not the one passed to
``EngineSttStep`` itself, whose own page-budget headroom already reserves
``num_frames`` separately from ``max_tokens`` — see that class's
``__init__`` docstring) by that step's ``num_audio_frames``, or every long
input silently gets less real generation budget than a short one, down to
zero once listening alone exceeds ``max_tokens``.

=== Iteration shape ===

A synchronous generator (:func:`run_turn`), NOT ``async def`` — this loop
does no I/O of its own (every yield point is directly downstream of a
blocking ``torch`` forward call); ``arbi_serve/realtime/turn.py``'s
``stream_turn`` is ``async def`` only because IT wraps blocking codec work
in ``asyncio.to_thread`` at the WS-session layer. A later caller (the
"different, not-yet-written" WS-session integration task) is expected to
drive this generator's ``next()`` the same way, e.g. one
``asyncio.to_thread(next, gen)`` per step — this module intentionally
doesn't do that itself (see module docstring above on scope). Per step this
yields, in order, a :class:`TextDelta` (STT text-channel output — text is
"" when the caller's tokenizer has nothing visible to show for this token,
e.g. a control token; ``function_logits`` is passed through untouched so a
later, different function-calling pass can hook tool-call detection in
here without this loop needing to know anything about tool-call syntax) and
an :class:`AudioChunk` (this step's TTS/codec waveform, ``[B,
wav_to_token_ratio]`` float at ``NemotronAudioCodec.SAMPLE_RATE``) —
deliberately separate dataclasses from ``realtime/turn.py``'s
``TranscriptDelta``/``AudioDelta`` (different payload: raw float waveform
here vs. already-PCM16-encoded bytes there, and this module must stay
importable without pulling in the realtime/session stack — same
"off-critical-path, decoupled" motivation as
``arbi_serve/audio/token2wav.py``). A final :class:`TurnDone` names the
stop reason and step count.

One STT backbone decode step's result — whatever produced it
(real ``ScheduledBatch``/``MultiStatePool``/``attn_ops`` plumbing
against a live engine-adjacent forward loop, mirroring
``tests/test_nemotron_voicechat_parity_live_gpu.py``'s
``_build_batch_and_pool``, extended step to step) is the caller's
concern — see module docstring.

Build the ``subword_id -> char_ids`` map
:meth:`~arbi_serve.models.audio.nemotron_ear_tts.CharAwareSubwordEncoder.set_char_map`
needs, from a plain subword-tokenizer ``str -> id`` vocab.

Direct port of the reference's ``ear_tts_model.build_vocabs``, minus
its ``vocab_dir`` file-cache path (irrelevant here — arbi-serve's own
tokenizer is already loaded once at process boot, not rebuilt per
turn) and its ``subword_padding_idx`` sentinel entry (irrelevant here
too: :meth:`CharAwareSubwordEncoder._prepare_inputs` already falls
back to an empty char sequence for any id absent from the map via
``.get(x, ())``, so omitting that one synthetic entry changes nothing
functionally).

1. Character vocab: every SINGLE-CHARACTER token string already
   present in ``vocab``, re-densified to ``0..N-1`` in original-token-
   id order (matches the reference's ``_build_char_vocab``).
2. ``subword_id -> char_ids``: every subword's char-id sequence,
   silently dropping any character absent from that char vocab
   (matches the reference exactly — not this port's own choice).

Takes a plain ``vocab`` dict (not a tokenizer object) so this module
stays decoupled from arbi-serve's specific tokenizer wrapper API —
callers pass whatever ``get_vocab()``-shaped dict their tokenizer
exposes.

The ``max_tokens`` value a caller MUST pass to :func:`run_turn` for
a turn driven by an audio-in :class:`~arbi_serve.runtime
.nemotron_voicechat_stt_step.EngineSttStep` (``num_audio_frames`` from
that instance's own :attr:`~EngineSttStep.num_audio_frames`) —
``base_max_tokens`` (the real generation-step budget, e.g. what a
caller/config actually means by "how long can the assistant talk")
PLUS that turn's mandatory, input-duration-driven audio-frame
"listening" steps, which :func:`run_turn`'s own step counter cannot
tell apart from real generation (see its docstring's "Termination"
section for why, and the bug this fixes). ``0`` frames (text-only
turn / no ``mm_feats``) is a no-op: returns ``base_max_tokens``
unchanged.

Run one turn: STT step, TTS frame, codec decode — lockstep, per
step — until EOS or ``max_tokens``. See module docstring for the full
design (lockstep coupling,

Args:
    model: the composite model, audio path ON (``ARBI_ENABLE_AUDIO``)
        — see :meth:`NemotronVoiceChatModel.__init__`.
    stt_step_fn: produces one STT decode step's result per call — see
        :data:`SttStepFn` and the module docstring's scope note on
        why this is injected rather than built here.
    audio_prompt_latent: this turn's speaker-conditioning latent,
        straight into :meth:`~NemotronVoiceChatModel.prime_tts_turn`
        (e.g. ``model.tts_default_voice_prompt_latent``). Its batch
        dim ``B`` sets the batch dim for the whole turn — this loop
        is written for the realtime path's single-request shape
        (``B=1``, matching ``arbi_serve/realtime/turn.py``'s
        ``stream_turn``, one ``asubmit`` per turn) but doesn't
        hardcode ``B=1`` anywhere.
    subword_id_to_char_ids / vocab: Exactly one must be given.
    max_tokens: hard cap on STT steps (and therefore TTS frames) —
        the ``max_tokens``-based half of the stopping condition (see
        module docstring's "Termination").
    num_iter/guidance_scale/top_p_or_k/noise_scale: passed straight
        through to every :meth:`~NemotronVoiceChatModel.generate_tts_frame`
        call — see that method's own docstring for what they mean.
        DERIVED from
        ``guidance_scale`` here, once, and reused for the whole turn
        — there is no separate flag to pass.
    initial_prev_code: ``prev_code`` seed
        override (e.g. a custom voice bundle's ``prompt_codes[:,
        -1:]``); defaults to ``model.tts_codec_silence_tokens``
        broadcast to ``[B, 1, num_quantizers]``.
    prime_subword_ids/prime_subword_mask: optional system-prompt text
        conditioning for :meth:`~NemotronVoiceChatModel.prime_tts_turn`'s
        warmup call (rare — see that method's own docstring).
    sanitize_control_codes: False to skip (e.g. a
        caller that already knows its checkpoint/sampling never
        leaks control codes and wants to save the per-step
        Site-3 BOS/EOS ``agent_idle`` switch. When
        ``duplex_session`` is given (both ids then required), every
        step reads ``stt_result.token_id`` directly against
        ``text_bos_id``/``text_eos_id`` and calls
        ``duplex_session.mark_agent_active()``/``mark_agent_idle()``
        accordingly — independent of ``stt_result.is_eos``/
        ``should_stop_turn`` (which intentionally swallow an
        EOS-shaped token while audio frames remain; see module
        docstring's "Frame-lockstep audio-in stepping", region 2, and
        ``_FrameLockstepState.should_stop_turn``'s own docstring).
        Defaults to ``None`` — the ordinary turn-based path (no
        persistent duplex connection) leaves this a pure no-op, so
        existing callers/tests are unaffected. See design doc
        §3.2.1's "Where this hooks into the per-tick loop" for why
        this runs here (right after this step's token id is known,
        before this same step's TTS-frame call below) rather than in
        the async wrapper around this generator.

        ``duplex_session`` also
        changes where the ``cond_on_prev_audio_tokens`` state
        (``prev_code``/``past_kv``) is read from and written to.
        With a session given, this function reads
        ``duplex_session.prev_tts_code``/``duplex_session.tts_past_kv``
        as the seed for THIS call instead of always rebuilding fresh
        state, and writes every step's result back onto the session
        (:meth:`~arbi_serve.realtime.nemotron_voicechat_duplex_session.NemotronVoiceChatDuplexSession.observe_tts_frame`)
        instead of only a loop-local — so the state survives across
        successive calls on one connection: a LATER call to this
        generator on the SAME session object continues the TTS
        backbone's KV cache and audio-code feedback exactly where the
        previous call left off (including skipping
        :meth:`~NemotronVoiceChatModel.prime_tts_turn` entirely once
        the session already carries a primed ``past_kv`` — re-priming
        would discard that history and restart the backbone's KV
        cache at the priming baseline). A fresh session (``prev_tts_code``
        still ``None``) falls through to the SAME silence-broadcast/
        ``initial_prev_code`` seed the turn-based path always uses,
        seeded lazily on this first call rather than at session-open
        time, since the seed's batch/device/dtype come from this
        call's own ``audio_prompt_latent``, which the session doesn't
        own. Defaults to ``None`` — the ordinary turn-based path is
        unaffected (every one of these branches is gated on
        ``duplex_session is not None``).

Yields: per step, a :class:`TextDelta` then an :class:`AudioChunk`;
    finally one :class:`TurnDone`.

Publish small host vectors to the device without stalling the host.

THE DEFECT THIS CLOSES. ``torch.tensor(<python list>, device=cuda)``,
``torch.from_numpy(a).to(cuda)`` and ``dev_buf.copy_(cpu_tensor)`` all source
their copy from PAGEABLE host memory, and torch serves such a copy as
``memcpy_and_sync`` — ``cudaMemcpyAsync`` plus an explicit stream
synchronize. The host then blocks until everything already queued on that
stream has drained: a stall proportional to the forward it is waiting
behind, not to the handful of ints being copied. Measured on 201/cuda:0 with
four 2048³ GEMMs queued, one step's seven publications cost **1513 µs**
pageable against **970 µs** staged through pinned memory; on an idle stream
the same seven cost 44 µs and 54 µs, so the saving is the stall and nothing
else. It is the obvious spelling, it is invisible in review, and
``torch.cuda.set_sync_debug_mode`` reports it as a sync (see
:mod:`arbi_serve.runtime.sync_census`).

WHY THE STAGING IS OWNED, AND NOT A PER-CALL ``pin_memory=True``. The one-line
version of this fix — allocate a pinned tensor per call and let torch's
caching host allocator recycle it — measures beautifully and is a trap. That
allocator can only hand a block back once the copy reading it has EXECUTED,
so while the host runs ahead of the GPU every in-flight staging buffer is a
distinct block, and it never returns one to the OS (``num_host_free`` stays
0 for the life of the process). Measured, same box: 2000 same-sized stagings
against an idle stream held 32 KiB in 2 allocations, but the identical
pattern behind a busy stream grew to **36.8 MiB in 1288 ``cudaHostAlloc``
calls**, and **68.9 MiB** at a wide prefill shape. That footprint is
(bytes per step) x (how far the host has run ahead) — a quantity this module
does not control and cannot state. Page-locked pages are neither swappable
nor reclaimable, so an unbounded one is a reservation the kernel's OOM killer
prices for you.

So the staging is a RING this module owns, and the bound is structural:

* ``_RING_DEPTH`` pinned arenas, bump-allocated within an arena and advancing
  to the next when a publication does not fit. Leaving an arena records a
  CUDA event on it; entering one waits for that event, so an arena is
  rewritten only after every copy that reads it has executed. This is the
  discipline ``PiecewiseBuffers.wait_staging_free`` / ``mark_staging_consumed``
  already imposes on the persistent path, and the same argument applies: the
  copies were enqueued at the HEAD of a step, ahead of its forward, so by the
  time the next step wants the arena back they have long executed. It stalls
  only when the host has lapped the GPU — and a host that far ahead loses
  nothing by waiting.
* The total is therefore ``_RING_DEPTH x cap``, where ``cap`` is the largest
  single publication seen, and every publication is per-step scheduling
  metadata bounded by the config's own ceilings: one int32/int64 per token on
  the flat axis (``max_batched_tokens``), one per row (``max_batch``), and one
  ``max_batch x max_pages`` int32 block table. Nothing here scales with the
  MODEL, with concurrency beyond ``max_batch``, or with how far the host runs
  ahead — the three ways the per-call form was unbounded.
* That total is PRICED before it is taken, through
  :mod:`arbi_serve.runtime.pinned_host_budget`, and reported to the ledger as
  residency because it moves. A host that cannot afford the growth does not
  get it: the call falls back to the pageable copy — slower, exactly today's
  behaviour, never an OOM kill.

A publication larger than the ring can grow to therefore still works; it just
pays the stall the ring exists to remove, and says so in the ledger by not
growing.

Return ``host`` as a fresh tensor on ``device``, copied asynchronously.

Falls back to a plain transfer on a non-CUDA device (CPU stub / fixtures),
where pinning does not exist and no sync is possible, and on a host that
cannot afford to grow the ring.

Fill an EXISTING device slice from ``host``, asynchronously.

For the persistent-buffer callers: the destination's ``data_ptr`` is the
address a captured graph references, so the bytes must land in it rather
than in a fresh allocation. Ordering is the caller's already — the copy is
enqueued on the current stream, ahead of whatever reads ``dst`` next.

The host-RAM ledger: every page-locked reservation is priced before it is taken.

The card has had a byte-exact ledger for a long time
(:mod:`arbi_serve.engine.vram_ledger`, ``docs/memory-accounting.md``): every row
declares how it was measured, the rows sum to the framebuffer, and the one row
allowed to absorb error is named. The HOST had no such ledger at all, and a
reservation nobody prices is a reservation nobody can refuse.

Pinned host memory is where that gap bites, because page-locked pages are
neither swappable nor reclaimable. An over-reservation of pageable memory
degrades — the kernel pages something out and the process gets slower. An
over-reservation of PINNED memory does not degrade: the kernel reaches for its
global OOM killer, and which process it picks is not our choice. On a shared
box the engine is the fattest target, so the engine is what dies, mid-request,
with no line in any log of ours saying what it had asked the host for.

So the rule this module enforces is the VRAM ledger's rule, moved to the host:

    **measure what the host can afford, price the reservation against it, and
    when it does not fit, shrink to what does and SAY SO.**

Three things follow, and each is a separate function here so each is testable
on its own.

``read_host_memory``
    What the host can afford, MEASURED, at the moment of the ask. Two ceilings
    exist and either can be the binding one: ``MemAvailable`` from
    ``/proc/meminfo`` (the host's), and the cgroup memory limit (the
    container's). Both are read; the smaller binds. Neither is assumed — a
    reading that could not be taken comes back as ``None`` and every consumer
    treats that as *refuse*, never as *plenty*.

``working_set_floor_bytes``
    What must still be allocatable AFTER the reservation. It is the engine's
    own non-pinned resident host working set, measured from ``VmRSS`` minus the
    bytes already booked here — not a percentage and not a fixed number of
    gigabytes. That quantity is a measured LOWER BOUND on what the engine will
    still want of the host: everything the boot built on the host is in it, and
    serving only adds to it (per-request structures, tokenizer state, output
    staging). It also scales with the inputs that actually move host demand —
    model, batch, dtype, vocab — because those are what produced the RSS.

``fit_pinned_host``
    The decision. ``available - reservation >= floor``, in whole granules of
    whatever the caller allocates in, returning how many granules it may have
    and — always, whether it shrank or not — the numbers that produced the
    answer.

**This is not a boot-time check.** Two of the consumers here take their bytes at
boot (the savepoint ring, the copy-engine collective) and two do not: the sleep
pool pins when it is asked to sleep, and the savepoint STORE raises its own byte
cap at the FIRST SNAPSHOT, mid-serving, because ``min_entries`` outranks the
configured cap. Every reading this module takes is taken at the moment of the
ask, so a ceiling raised during serving is priced against the host as it is
then, not against the host the boot saw.

**The conservative direction, named.** Both errors here are real, and they are
not symmetric. Shrinking too far costs exactly what the shrunk thing bought
(fewer savepoint slots ⇒ a resume becomes a re-prefill), and that is measured
and graceful. Pinning too much costs the process a SIGKILL. So where an input
is uncertain this module rounds toward LESS pinning, and says so at the site
(see ``cgroup_available_bytes``).

**One seam, so the next pinned allocation inherits this.** The savepoint ring
is the first large pinned reservation in the engine and will not be the last —
the copy-engine TP collective page-locks its ``/dev/shm`` slots, the sleep pool
parks a whole model's VRAM in pinned host RAM, and any host/NVMe tier for
recurrent state pins too. All of them go through :func:`reserve_pinned_host`
(taken once and held) or :func:`set_pinned_host_bytes` (a footprint that moves,
whose CEILING was priced where it was raised), which is what makes the ledger
complete and what makes a future one budgeted by default rather than by
remembering. ``tests/test_pinned_host_seam.py`` fails on a new large pinned
allocation that does not.

**The rows are per OWNER, and that is not cosmetic.** "9.7 GiB pinned" is not
something an operator can act on; "the ring took 2.31 GiB at boot and the store
grew to 7.4 GiB at the first snapshot" is.

Pure stdlib and pure arithmetic: no torch, no CUDA, no engine. Every input can
be injected, so both directions of every decision are exercisable off-GPU.

A pinned-host reservation the host cannot afford, refused before it is taken.

Raised by :func:`reserve_pinned_host` for a caller whose size is FIXED.
A caller whose size is divisible asks :func:`fit_pinned_host` instead and
shrinks; this exception is for the ones with nothing to shrink, where the
only alternative to refusing is being killed by the kernel later with no
diagnosis at all.

Headroom inside a cgroup memory limit, or ``None`` when there is no limit.

``limit - current`` alone understates it, because ``memory.current`` counts
clean page cache the kernel drops before it ever reaches the OOM killer.
Adding back the cgroup's own ``inactive_file`` is the same allowance
``MemAvailable`` makes on the host side, and without it a long-lived
container with a warm page cache would read as having no headroom at all
and refuse every reservation for a reason that is not true.

When ``memory.stat`` could not be read, ``reclaimable_file_bytes`` is
``None`` and the result falls back to ``limit - current``. That is the
conservative direction for this quantity (see the module docstring): it
under-states headroom, which costs slots, rather than over-stating it,
which costs the process.

One reading of what the host can afford, and of what this process holds.

Every field is ``None`` when the file behind it could not be read. That is
a SENTINEL, never a zero and never an optimistic default: a reservation
priced against a reading that was not taken is exactly the reservation this
module exists to prevent.

Host RAM that must stay allocatable after the reservation.

It is this process's own NON-PINNED resident working set: ``VmRSS`` minus
everything already booked through this ledger. ``VmRSS`` counts anon, file
and shmem residency, so a page-locked page is inside it whichever of those
the driver's mapping lands in, and subtracting the booked bytes leaves the
part that is the engine itself.

Why that quantity and not a percentage or a number of gigabytes: it is a
MEASURED lower bound on what this engine will still ask the host for.
Everything the boot built on the host is already inside it, and serving
only adds — per-request structures, tokenizer state, output staging — so a
host that cannot spare it again cannot serve. And it moves with the inputs
that genuinely move host demand (model, batch, dtype, vocab), because those
are what produced the number. A constant would be right on the
configuration it was calibrated on and silently wrong on every other one.

``None`` when ``VmRSS`` could not be read — a floor that could not be
measured must not read as a floor of zero.

How much of a pinned request the host can afford, and the arithmetic behind it.

Carries the inputs as well as the answer, because a shrink an operator
cannot reproduce from the log line is a shrink they will attribute to a
tuning constant somebody chose.

Granules of ``granule_bytes`` this host can afford to page-lock.

``available - reservation >= floor``, solved for the granule count and
clamped to what was asked. ``host`` and ``already_pinned_bytes`` are
injectable so both directions — the ample host that must NOT shrink and the
starved host that must — are exercisable without a real host in that state.

Grants ZERO, with the reason, when either input could not be measured. An
unpriced pinned reservation is the defect; permitting it because the
measurement failed would be the same defect wearing an excuse.

Book ``nbytes`` of page-locked host RAM to ``owner``; return the total held.

THE seam. Call it immediately before page-locking, never after: the point
is to refuse a reservation the host cannot afford, and a reservation booked
after the allocation has already been made can only describe the kill.

Raises :class:`PinnedHostBudgetError` when it does not fit. A caller with
something to shrink calls :func:`fit_pinned_host` first and asks for a size
that does; by the time it reaches here the answer is already yes, so a
raise from a shrinkable caller means the host moved between the two calls,
which is exactly when it should refuse.

``owner`` must be declared in :data:`PINNED_HOST_OWNERS`.

Level-set what ``owner`` currently holds, replacing its row.

For a consumer whose page-locked footprint MOVES — admitted and evicted
continuously — rather than being taken once and held. Such a consumer
prices its CEILING through :func:`fit_pinned_host` where the ceiling is
raised (that is the decision point, and it is not at boot), and reports its
residency here, so the ledger row is what is actually resident rather than
a high-water mark nobody released.

Deliberately does NOT re-price: this is called from inside the consumer's
own accounting, which runs on the snapshot path, and a ``/proc`` read per
admitted entry would put a syscall where a counter belongs. The refusal
lives at the ceiling; this is the meter.

Lifetime token for one reservation held by several objects.

Weak-referenceable and nothing else: every holder keeps it alive, so the
finalizer that releases the bytes runs when the LAST of them is collected.

Reserve ``nbytes`` and release them when the LAST holder is collected.

For a reservation whose bytes are handed BACK to a caller — several objects
that share one page-locked allocation and are dropped independently. A
finalizer on any single one of them releases while its siblings are still
resident, and the ledger then under-reports exactly the bytes that are
still page-locked, which over-grants the next reservation. So the token is
attached to every holder and the release rides the token.

Refuses like :func:`reserve_pinned_host` when the host cannot afford it,
and attaches nothing in that case. Returns the token, so a caller with no
natural holder can keep it in a scope of its own.

Give ``nbytes`` back; return what ``owner`` still holds.

A reservation that is freed and not released here leaves the ledger
over-reporting, which raises the measured floor of the NEXT ask and shrinks
it for bytes nobody holds.

The boot block, printed beside the VRAM ledger. ``None`` when nothing is pinned.

Deliberately its own table rather than rows inside the VRAM ledger: that
ledger's whole contract is that its rows sum to the card, and host bytes
inside it would break the identity it exists to hold.

One sentence naming the ceiling, the floor and what was granted.

Rendered into the owner's own boot line, so the shrink is stated where
the reader is already looking rather than in a second log line they
have to correlate.

Pool registry — thin Protocol re-export over :mod:`named_pool`.

The :class:`NamedPoolRegistry` / :class:`NamedMemPool` classes in
:mod:`arbi_serve.runtime.named_pool` provide the registration / lookup /
snapshot surface. This module exposes the :class:`Pool` Protocol the
rest of the code types against, without changing those classes. The
Protocol formalizes the contract surfaced by :class:`NamedMemPool`
(``name`` / ``bytes_inuse`` / ``bytes_capacity`` / ``snapshot``) and
documents the ``predict_bytes(...)`` interface that
:mod:`arbi_serve.engine.memory_budget` exports as module-level
predictors.

:class:`Pool` Protocol — what every named allocator exposes.

Matches what :class:`NamedMemPool` provides; this
module documents the contract so non-allocator code can typecheck
against the abstract surface without importing
:class:`NamedMemPool` itself.

Implementations register with the engine's :class:`PoolRegistry` via
:meth:`PoolRegistry.register` (existing) and route allocations
through ``with pool.use():`` (existing on :class:`NamedMemPool`).

The ``predict_bytes(...)`` family is documented here as a contract
(the boot phase pre-flights every pool's predicted footprint into
:func:`arbi_serve.engine.build.compute_kv_budget` so
:exc:`BudgetExceeded` raises fail-loud BEFORE the alloc), but the
actual predictors live in :mod:`arbi_serve.engine.memory_budget` as
module-level functions. The contract here pins the shape.

Post-step leak assertion: the ``unattributed`` residual must not

Reads :meth:`PoolRegistry.snapshot_with_residual`, whose trailing
synthetic entry (``name == "unattributed"``) carries every byte of
caching-allocator allocation NOT routed through a NamedMemPool — a
leak signal. Raises :class:`AssertionError` when it drifts past the
tolerance. No-ops when CUDA is unavailable or no residual entry is
present (e.g. CPU test fixtures), where ``snapshot_with_residual``
returns named pools only.

One-segment slab fold for small, co-live, engine-lifetime buffers.

A :class:`~arbi_serve.runtime.named_pool.NamedMemPool` gets its physical from
torch's caching allocator, which never hands out a block smaller than the
SEGMENT it carves it from. The segment size is a step function of the request
(:data:`_K_SMALL_SIZE` / :data:`_K_SMALL_BUFFER` / :data:`_K_MIN_LARGE_ALLOC` /
:data:`_K_LARGE_BUFFER` / :data:`_K_ROUND_LARGE`), and the step is steepest in
``torch.cuda.MemPool`` cannot give that overhang back either: ``empty_cache``
does not visit a private pool's block pools (pytorch#145168), so a pool holding
one few-MiB buffer holds a whole segment for the process lifetime.

Several such pools each pay that overhang separately. Folding their buffers into
ONE slab allocation pays it once: a single request at or above
:data:`_K_MIN_LARGE_ALLOC` is rounded only to :data:`_K_ROUND_LARGE`, not up to
a whole large buffer.

:class:`PoolSlab` is that slab. It allocates one byte tensor inside a host pool
and hands out ``narrow()`` views of it. Each view is charged to a SUB-TAG
(:meth:`~arbi_serve.runtime.cumem_allocator.CuMemPoolAllocator.charge_subtag`),
so every consumer that used to own a pool keeps its own VRAM-ledger row and its
own cap key — a fold that saved bytes by collapsing the attribution would be
trading the thing this accounting exists for.

Only for buffers that are ALL of: allocated once at boot, live for the engine's
lifetime, and co-live with each other. A transient belongs in a per-step arena,
where its bytes are reused; folding it here would pin them forever.

A carve would take more than its sub-tag or the slab was sized for.

Raised rather than silently growing: the slab's bytes were held out of the
KV budget as a specific number, so a carve past it spends VRAM the KV pool
was already handed.

Physical bytes the caching allocator maps to serve a request of ``nbytes``.

The segment size, not the request: the difference is the overhang the fold
exists to pay once instead of once per pool.

Bytes a fold of ``sizes`` returns, versus one pool per entry.

A LOWER bound: each entry is priced as ONE segment, while a pool holding
several buffers can take several. Used as the gate on whether a fold is
worth doing, so under-claiming is the safe direction.

One slab allocation in a host pool, carved into per-sub-tag views.

``budgets`` maps sub-tag → the bytes that sub-tag may carve; the slab is
sized to their sum (each entry rounded to the carve alignment) and every
carve is checked against its entry. The budgets are the SAME closed forms
the memory plan prices the folded pools at, so the slab spends exactly what
was held out for them.

Views alias the slab, so the slab holds them all resident; dropping the
:class:`PoolSlab` does not free a view still referenced by a live tensor.

One slab belongs to ONE model: the budgets are closed forms of that model's
vocab / width / dtype. An engine that builds a second model must retire the
first slab rather than carve from it.

Return a zeroed ``(shape, dtype)`` view of the slab, charged to ``subtag``.

Bump-allocated and never freed WITHIN one slab: its consumers live for
the model's lifetime, so a free list would never be exercised. The
cursor is therefore only meaningful against the model the budgets were
computed for — a second model in the same process gets a NEW slab, and
the previous one is retired at the build seam
(:func:`~arbi_serve.engine.persistent_fold.retire_persistent_slab`).
Reaching a spent budget here means that retirement did not happen.

Raises :class:`SlabBudgetExceeded` when the carve would take more than
``subtag``'s budget or run past the slab.

:meth:`carve`, but idempotent for a repeated ``(subtag, shape, dtype)``.

Returns the SAME view (re-zeroed) rather than a second one. For a
buffer that is engine-lifetime and shape-identical, asking twice means
the caller was re-installed, not that it wants two buffers — and
carving twice out of a bump allocator that never frees would spend the
sub-tag's budget again and raise :class:`SlabBudgetExceeded`.

This is what a live config-override variant build does: it shares the
donor member's pools, so it re-installs against a slab the baseline
already carved to its budget.

Move live ``tensors`` into the slab in place, returning bytes moved.

Copies each tensor into a fresh carve and rebinds its storage
(``Tensor.set_``), so every tensor keeps its Python identity and no
caller re-points anything. Two phases — carve and copy ALL first, rebind
only then — so a failure mid-way leaves nothing half-rebound.

BOOT-ONLY, and only before any cudagraph capture: a captured graph bakes
the address it read, and rebinding under one would leave the graph
replaying out of freed storage.

Drop the slab and every sub-tag charge it booked.

The physical returns only once the last view is also dropped and the
host pool is destroyed — a private MemPool does not release segments to
``empty_cache``.

Single source of truth for VRAM pool names, categories, descriptions.

Every VRAM bucket arbi-serve tracks shows up on four surfaces that have
to AGREE so an operator reading any one of them knows what each pool is:

  * the :data:`arbi_serve.engine.engine.NAMED_POOL_NAMES` registry (the
    live ``torch.cuda.MemPool`` buckets),
  * the metrics ``pool=`` label values
    (``caching_allocator_reserved_bytes{pool=...}`` etc. in
    :mod:`arbi_serve.server.metrics`), which are emitted straight from the
    registry name,
  * the phase-2 freeze SERVING POOL TABLE rows + the cuMem-invisible
    reconcile terms (:mod:`arbi_serve.engine.phase2_freeze`),
  * the admin UI's GPU-memory panel, which renders group + row captions
    from :func:`taxonomy_payload` rather than a copy of this table.

Every pool name is a self-documenting ``category.role`` string, defined
ONCE here, and all four surfaces draw from this registry.

TWO axes, and the prefix carries both. ``category`` says what the bytes
are FOR (a lifetime and a scaling law); :data:`CATEGORY_TENURE` says what
KIND of number the row is — physical that EXISTS, physical that is merely
PROMISED, physical that is not ours, or address space that is not memory
at all. A surface that renders only the first axis puts a weight slab and
a reservation nothing has taken in one list, where both read as memory in
use. Both axes travel with every name, and ``category`` alone determines
``tenure``, so the two cannot disagree.

Naming scheme — ``category.role``:

  * ``model.*``   — engine-lifetime model parameters (weights, adapters,
    the loader's cold-boot scratch).
  * ``state.*``   — per-sequence runtime memory: the "KV" of each layer
    type (attention KV slab, GDN recurrent state, short-conv state, MLA
    shared KV). Reserved-VA vs mapped-physical KV are named APART
    (``state.attn_kv`` is the reserved-VA pool; ``state.attn_kv.mapped``
    is the physically-mapped growable slab).
  * ``spec.*``    — speculative-decode overhead (drafter slab, the GDN
    recurrent-state rollback snapshot for partial-accept).
  * ``capture.*`` — CUDA-graph capture memory (the captured-graph mempool
    and the persistent kernel I/O buffers captured graphs reference).
  * ``scratch.*`` — transient per-step compute arenas (activation arena,
    attention codec scratch, GDN workspace, RoPE cache, cuBLAS GEMM
    workspace).
  * ``comms.*``   — collective-communication buffers (NCCL staging).
  * ``unpooled.*`` — OUR OWN live tensors that no named pool claims:
    real, measured, in use, unbudgeted only because no pool claims them.
    Two ways that happens — an allocation that escaped placement into
    torch's default pool, and a private ``MemPool`` nobody registered.
  * ``driver.*``  — memory OUTSIDE every allocator we own: the CUDA
    primary context, TP sibling contexts + NCCL comm workspace, the
    driver's kernel-stack and cubin residency, another tenant's VRAM,
    and the whole-card remainder.
  * ``address_space.*`` — NOT MEMORY. Virtual address space this process
    still holds while the driver already has the physical back. Takes
    part in no sum on the card, and is named apart from ``driver.*`` so a
    reader of one number never mistakes it for driver residency.
  * ``transient.*`` — device-free physical CLAIMED by one in-flight
    serving step: the bytes the KV grow held back so a step can
    allocate them. Itemised by the step phase that takes them.
  * ``unclaimed.*`` — device-free physical NOTHING will allocate: a
    ceiling the operator asked for, a gate the residency controller
    needs, allocator round-up slack, and — the row that matters — the
    bytes no row on the card declares at all.

The taxonomy carries, per name, what the pool HOLDS, WHEN it is
allocated, and what it SCALES with, so the name plus its one-line
description is self-explanatory.

:func:`metric_pool_label` /
:data:`POOL_TAXONOMY` are the single import point; renaming a pool means
editing ONE entry here (plus the registry tuple that references the same
string). ``tests/test_pool_taxonomy_sync.py`` asserts the registry, the
freeze table, the metric labels and the admin UI all draw from this name
set — no orphan or synonym names anywhere, and no category the panel
cannot render.

Taxonomy entry for one VRAM bucket.

``category`` is the ``category`` half of the ``category.role`` name;
``description`` is a one-line "holds X, allocated WHEN, scales with Y";
``provenance`` is one of :data:`PROVENANCE` — how the row's number is
obtained, which is the difference between a measurement and a guess and so
travels with the name rather than living in a comment.

Return the :data:`TENURE` of row ``name`` — what KIND of number it is.

Derived from the row's category rather than stored per row, so a name
cannot declare a category and contradict it. Raises on an unknown name for
the same reason :func:`describe` does.

Return the taxonomy as JSON for a UI: categories, rows, provenance.

Served by ``GET /v1/admin/memory/live`` so the admin memory panel renders
every group, row caption, provenance badge and leak flag from THIS module
instead of a hand-maintained copy that silently leaves a new category blank.

``provenance`` on a pool entry is that row's DEFAULT; a payload row may
override it where the runtime knows better (a reserve that fell back from a
measurement to its analytic bound is not measured any more, and must not
keep saying it is).

``tenure`` is NOT overridable: it says whether the row's memory exists at
all, it is a pure function of the category, and a surface that let a row
claim a tenure its category does not have would be back to two schemes.

Return the metric ``pool=`` label value for a registry pool ``name``.

Identity — the registry name IS the label value IS the taxonomy
key. Exists so the three surfaces route through one function rather than
re-deriving the string, keeping a single seam for any label/registry
divergence.

Boot-time activation / vision peak profiling for
:class:`EagerModelRunner`.

These collaborators run ONCE at boot (from
:func:`arbi_serve.engine.build.profile_and_size_kv_pool`) to derive the
auto-sized KV pool depth — they are firmly off the per-step hot path.
They are module-level free functions taking the runner explicitly;
:class:`EagerModelRunner` keeps thin delegating method wrappers
(``runner.profile_activation_peak`` / ``runner._profile_vision_peak`` /
``runner._worst_case_vision_inputs`` /
``runner._clamped_synthetic_state_indices``) for the boot path + tests.

Turn a probe OOM into the same loud refusal the multi-shape probe raises.

THE SEAM. Two boot probes run the same synthetic prefill: this one (which
sizes the KV budget) and :func:`~arbi_serve.runtime.activation_profile.
_profile_one` (which sizes the serving floor). The second wraps its forward
and raises :class:`~arbi_serve.runtime.activation_profile.
ActivationProbeFailed` — a refusal naming the shape and the knobs that move
it. This one did not, so a config too wide to profile died with a raw
``torch.OutOfMemoryError`` and a Triton-deep traceback, from the probe that
runs FIRST and therefore decides the outcome for every such config.

Both halves of the boot contract want the same thing here: an inadmissible
configuration must refuse, and say what to change. The refusal was already
written; it was reachable from only one of the two probes.

The engine is not poisoned by this — the probe runs long before the Phase-2
layout freeze and the process exits — so the conversion is about what the
operator is told, not about recovery.

Per-sequence q-token lengths for the boot activation-peak prefill probe.

A real scheduler never emits a single prefill row wider than
:func:`~arbi_serve.runtime.activation_profile.prefill_row_cap` —
``min(max_context, chunk_prefill)``: admission rejects any prompt ``>=
max_context``, and every scheduler path sizes a row at
``min(remaining_prompt, chunk_prefill, token_budget)`` (recurrent rows
included). The ``max_batched_tokens`` token budget is a per-STEP aggregate
spread across multiple admitted prefills, NOT one contiguous sequence. So
the synthetic activation-peak prefill must be spread the same way:
``ceil(N / cap)`` rows (:func:`prefill_probe_num_seqs`), each ``<= cap``,
summing to exactly ``N = max_batched_tokens`` so the full-budget
activation peak is still measured while every per-row attention shape is
one the scheduler can actually build.

Returns the list of per-row lengths (length == the probe ``num_seqs``).
Every entry is ``>= 1`` and ``<= cap`` (when the cap is bounded), and
``sum(result) == max_batched_tokens``. Degenerate inputs
(``max_batched_tokens <= 0``) return ``[]``; an unbounded cap falls back
to a single row.

Run one synthetic forward to measure the activation peak.

Builds a transient :class:`MultiStatePool` + per-layer attn ops
+ per-kind metadata builders sized just large enough to absorb
a single ``max_batched_tokens``-wide prefill, runs that
forward, captures :func:`torch.cuda.max_memory_allocated`, and
returns ``(peak_activation_bytes, scratch_pool_bytes)``.

``peak_activation_bytes`` is ``max_memory_allocated -
weights_bytes - scratch_pool_bytes`` — the part of the peak
that is neither weights nor the (transient) KV pool. That is
what :func:`compute_kv_budget` subtracts from the usable
budget. Only the scratch pool's CACHING-ALLOCATOR bytes are
subtracted, because only those are inside
``max_memory_allocated``; ``scratch_pool_bytes`` is returned as
the pool's full physical footprint, cuMem arenas included.

The scratch state is fully torn down before returning;
``empty_cache`` is called and ``memory_allocated`` is asserted
to fall back to the pre-profile baseline (modulo allocator
granularity). No engine-visible state survives the call.

Side effect: the engine's ``pool`` / ``attn_ops`` /
``metadata_builders`` are mutated in-place during the profile.
We restore them to ``None`` / ``[]`` / ``{}`` on exit; the
caller (:func:`arbi_serve.engine.build.profile_and_size_kv_pool`)
is followed by :func:`arbi_serve.engine.build.build_active`
which constructs the production state fresh.

Run a worst-case dummy vision-tower forward to fold the ViT
activation transient into the profiled peak.

No-op unless vision is enabled AND the model carries a
``visual`` tower — keeping the text-only peak (and the whole
memory budget) byte-identical when vision is off. The dummy
encode runs inside the caller's ``reset_peak_memory_stats …
max_memory_allocated`` window, so the peak naturally becomes
``max(text_peak, vision_peak)``. The tower is TP-replicated, so
every rank measures the same vision peak and the profiler's
cross-rank ``all_reduce(MAX)`` reconciles them.

Run a worst-case dummy audio-tower encode inside the profiled peak.

No-op unless audio is enabled AND the model carries an
``audio_tower`` — text-only (and vision-only) peaks stay
byte-identical. The worst case per chunk is a full 25 s mel
(2502 frames → 1251 post-conv positions, within the tower's 1500
pos-emb capacity); ``ARBI_MM_MAX_AUDIO_CHUNKS`` scales the batched
chunk count a single request may carry.

Run a worst-case dummy perception encode inside the profiled peak.

The perception-tower sibling of :func:`profile_audio_peak`, duck-typed on
``profile_worst_case_perception_encode`` rather than on a tower attribute,
so any model whose speech-in encoder runs eagerly outside every pool and
every captured graph is priced the same way. The transient folds into
``media_encode_peak_bytes``, which the post-capture grow floor holds back as
``media_reserve``.

No-op on a model without the hook — text / vision / Step-Audio-2 peaks stay
byte-identical.

Run a worst-case dummy TTS-frame generation inside the profiled peak.

No-op unless the served model exposes ``profile_worst_case_tts_frame`` —
a duck-typed hook, never a concrete model check, so this stays generic
across any model whose speech-output path runs eager per tick outside the
paged-KV pool and every captured graph (today: NemotronVoiceChat's
frame-lockstep TTS backbone + MoG head + audio-codec decode chain — see
:meth:`~arbi_serve.models.nemotron_voicechat.NemotronVoiceChatModel.profile_worst_case_tts_frame`).
Models with no such hook (including Step-Audio-2, whose token2wav vocoder
is a genuinely separate process/pool with its own
``ARBI_AUDIO_T2W_RESERVE_GB`` reserve — see
:mod:`arbi_serve.audio.t2w_pool`) leave the peak byte-identical.

Synthesize the worst-case dummy ``(pixel_values, image_grid_thw)``
matching the Qwen-VL processor output layout.

The per-image pixel budget is the checkpoint's processor
``max_pixels`` (the true runtime ceiling) unless
``override_max_pixels`` > 0. Patches-per-image is derived from
that budget, laid out as a square ``h × w`` grid with ``t=1``,
``h`` and ``w`` rounded down to a multiple of
``spatial_merge_size`` so the merger and 2-D RoPE coord math
don't trip. Pixel values are zeros — the profiler measures
memory, not numbers.

Build a synthetic per-row recurrent slab-row mapping clamped
into ``[0, recurrent_slab_capacity)``.

Used by the profile-time forward to pin
:attr:`ScheduledBatch.state_indices` so the recurrent backends
(Mamba / GDN / ShortConv) can never index OOB on a scratch
slab. The default-resolve path (:class:`StateIndicesSubBuilder`)
falls through to ``arange(B)``; on a profile shape with
``B > recurrent_slab_capacity`` the GDN block raises
:class:`IndexError` deep inside :meth:`GDNBlock.forward`. That
exception could be swallowed by
:func:`profile_activation_peaks._profile_one`, leaving the
per-shape budget math silently under-counting the GDN
allocation footprint (the named-pool delta still records
what landed before the throw, but the post-throw allocations
the block would have made do not get attributed).

Returns ``arange(B) % recurrent_slab_capacity`` when ``slab``
exposes any recurrent pool view; falls back to ``arange(B)``
otherwise (paged-only models — no slab to clamp against).

LIFO teardown registry for process / GPU resources.

A site that acquires a GPU resource — or sets a *module-global* that
pins one — registers a ``release()`` callback at acquire time via
:class:`ReleaseRegistry`; shutdown drains the registry in reverse
(LIFO) order. The per-resource cleanup knowledge lives at the acquire
site, not in ``shutdown``, so a new global can't be forgotten —
registering it is the same act as acquiring it.

Draining is destructive (callbacks are popped), so a second drain is a
no-op — ``Engine.shutdown`` is idempotent for free. A callback that
raises is logged and skipped so one bad hook can't strand the rest.

Register ``release`` to be called (once) at the next :meth:`drain`.

``name`` is for diagnostics only. Registering the same name twice
is allowed — both run — so idempotent releases (e.g. setting a
global back to ``None``) are safe to register from multiple
acquire sites.

Run every registered callback in reverse order, then clear.

Reverse order so resources release in the opposite order they were
acquired (a later acquire may depend on an earlier one). Each
callback is best-effort: an exception is logged and the drain
continues. Idempotent — a second call finds an empty registry.

Row tiling: spend a stated number of bytes, not however many the batch has.

THE INVERSION THIS EXISTS FOR. A row-independent operation written the obvious
way materialises every intermediate at the full row count, so its peak is an
OUTPUT of whatever the batch happened to be — and a reserve for it is then a
measurement of that batch, correct until the batch changes. Tiling inverts it:
state the bytes the operation may spend, divide by the measured bytes per row,
and the row count becomes an OUTPUT of the reserve. The peak then stops
scaling with ``max_batched_tokens`` (or ``chunk_prefill``, or whatever else
sets the batch), and the reserve becomes a policy number instead of a
consequence.

THE PRECONDITION IS THE WHOLE CONTRACT: **no cross-row reduction.** Every
output row must depend on its own input row and nothing else. Where that
holds, tiling is EXACT — not a reassociation, not a precision change, so the
receipt is bitwise equality against the untiled path rather than a tolerance.
Where it does not hold, tiling is simply wrong and no budget makes it right.
Callers assert it against their own operation; :mod:`arbi_serve.spec_decode.
dflash`'s ``project_context_stacked`` qualifies because its only reduction is
the K-RMSNorm's, over ``head_dim``.

WHAT A CALLER STILL OWES. This module divides; it cannot supply the numerator,
and a caller holding the right shape around the wrong budget still gets a peak
pinned to whatever that budget was picked at. Two callers are in that state and
the difference between them matters:

  * the load-time quantizers (``weight_quant/awq/rtn.py``'s
    ``_quant_row_chunk``, ``weight_quant/awq/marlin``'s
    ``_unpack_grid_n_block`` / ``_repack_n_block``) have the right SHAPE — a
    byte budget divided by a per-row size derived from ``in_features`` — over a
    on every card; what it stands in for is free VRAM at load time, which the
    quantize path already reads elsewhere to refuse a head that will not fit.
  * ``engine/logprobs.py``'s ``_tile_rows`` takes its width from a runtime flag
    and so does not reach :func:`rows_per_tile` at all — only :func:`row_tiles`
    consumes it. A knob is not a budget: it says how many rows, never how many
    bytes. Inverting it is a behaviour change, and the ``scratch.logprobs``
    reserve is itself computed FROM the knob, so reading that reserve back as
    the budget returns the knob unchanged — a division that measures nothing.

``_repack_n_block`` additionally needs its result rounded DOWN to a multiple of
Marlin's N tile, a granularity :func:`rows_per_tile` has no concept of (its
``min_rows`` is a floor, not an alignment), so it keeps its own division.

Rows one tile may take, from a byte budget and a measured per-row size.

``budget_bytes`` must come from something real — a reserve line the caller
already declares, a pool's own size — never a number chosen because it
looked reasonable. A constant here recreates exactly the defect tiling is
being used to fix, one level up: the peak would then be pinned to whatever
batch the constant was picked at.

``bytes_per_row`` is the operation's peak at ONE row. Measured or
closed-form, but it must track the operation's real geometry, because it is
the only thing that makes the tile follow a different model or dtype.

``min_rows`` is a floor the tile may not cross for reasons that are not
about memory — a kernel that changes dispatch leg on the row count, a
hardware tile width. It is applied AFTER the division, so a budget too
small to reach it yields ``min_rows`` and overspends knowingly rather than
silently changing which kernel runs.

``row_bound`` caps the result: a call already inside the budget is returned
whole, which is what lets callers keep a single-shot, fixed-shape path for
the small row counts a captured graph covers, and tile only above it.

Returns ``row_bound`` unchanged whenever the tile cannot be priced (a
missing budget or per-row size), so an unpriceable tile degrades to the
untiled behaviour rather than to a guess.

``(lo, hi)`` half-open row spans covering ``rows`` in ``tile``-wide steps.

One span covering everything when ``tile`` is zero, negative or at least
``rows`` — so a caller that took the single-shot path and a caller that
tiled run the SAME loop body, and the untiled path is the one-iteration
case of the tiled one rather than a second copy of it.

Round ``rows`` UP to a multiple of ``granularity``.

A private ``MemPool`` retains one block per DISTINCT allocation size it has
served, so a caller that projects at its exact row count teaches the pool a
new size class on every new width. The reserve the serving floor must hold
for that pool is the UNION of those classes, which is why it keeps growing
boot over boot: each boot observes a wider union than the last, the next
boot reserves it, and the KV pool pays.

Rounding the projection width to a multiple collapses an unbounded set of
widths to ``ceil(max_rows / granularity)`` classes. The cost is the padded
rows' compute, which is discarded -- sound only because the projection has
NO cross-row reduction (``fc``, the norms, the K/V projections and RoPE are
all per-row), so slicing the result back to the true width is exact rather
than an approximation.

WHAT THIS DOES NOT DO, because the difference decides whether it is worth
calling. A bounded CLASS COUNT is not a bounded RESERVE. The union is a
SUM over the classes the pool has served, not a max over them, and the
allocator only reuses a freed block for a request the block can hold: a
class that is the running WIDEST when it first appears takes a segment of
its own, and every narrower class that follows is served from elsewhere,
so that segment stays mapped and idle for the life of the process. The
reserve this bounds is therefore ``classes x widest``, and shrinking it
further is a question about the ORDER widths arrive in, not about how many
there are. The construction that gives ONE class is the grown-on-demand
slab
(:meth:`~arbi_serve.models.qwen3_5.Qwen3_5Model._ensure_residual_overflow`):
allocate at the widest width, slice the true width back out, spend no
compute on padded rows. READ THE SECOND MEASUREMENT BELOW BEFORE REACHING
FOR IT -- the class collapse is real, and this pool does not care.

AND ROUNDING WIDTHS UP DOES NOT SHRINK THIS POOL, which is what decides
whether to reach for a coarser value. Collapsing the class count is not
free: every allocation is rounded UP, so the added bytes are paid on EVERY
request, while the classes the rounding removes were being served out of
split blocks the allocator was already reusing. Replay a captured arena
allocation trace back through the same allocator with progressively coarser
width rounding applied to what it asks for, and the class count falls by
more than an order of magnitude while the pool's reserved high-water rises
MONOTONICALLY -- the collapse returns less than the rounding costs.

SCOPE THAT RESULT HONESTLY before carrying it to this knob: that replay
rounds the arena's WHOLE request stream, and this granularity governs only
the observe projection's share of it, which is a small fraction of the
bytes the pool serves (the trace says which -- read it, do not assume this
sentence). So the measurement says the ARENA is not shrunk by rounding
widths up; it does not license a claim about what this particular value
reclaims in either direction. What it does settle is that a deployment
wanting the reserve smaller should not be reaching for a bigger bucket at
all. Nor for the slab, which is the second measurement.

NOR DOES THE SLAB, AND THAT IS THE MORE SURPRISING HALF. It was built and
A/B'd on the arena's single largest allocation site -- the EXL3 GEMM's
per-call output and Hadamard scratch -- allocated at the row bound the
engine already declares for that linear and sliced back to the served
width. The receipts say the mechanism did exactly what it was designed to
do: the site's distinct-size count fell by an order of magnitude and every
padded call was counted. Against a null control that reproduced the
unchanged arm's reserved high-water, stranded bytes and segment counts
EXACTLY, the pool's reserved high-water moved by ONE cuMem granule -- the
smallest change it can represent. Padding the FEW-HUGE sites instead is
refused on the same evidence.

AND EVERY WAY OF GATING IT PAYS SOMEWHERE ELSE, which is worth knowing
before rebuilding it. A pad must be refused wherever the padded bytes would
outlive the call. Refusing only while CAPTURING lets the compile/capture
warmup through -- it runs eager, inside ``capture.*``'s own ``use()``, whose
pools are fixed reservations -- and that arm booted with a materially
smaller KV pool. Refusing unless the arena is the active pool fixes the
memory exactly (KV lands on the null control) and costs DECODE instead,
past the median floor, because the gate now runs a pool lookup on every
call and books a refusal on the ~28% of calls that are outside the arena.

ONE MECHANISM EXPLAINS ALL OF IT, and it is not size classes. The reserve
tracks what is LIVE AT THE PEAK, and a size class is not evidence of
residency: decompose the recorded trace at the instant its peak occurs and
every site these levers converted holds ZERO. Padding an allocation that is
already dead when the peak happens cannot move the peak, however many
classes it removes. The ``classes x widest`` model that this docstring used
to rest on is also wrong by an order of magnitude against that trace -- the
caching allocator SPLITS a large free block to serve a smaller request, so
the classes were already being served out of shared segments.

THE QUESTION TO ASK OF A CANDIDATE SITE is therefore "is it resident at the
peak", never "how many classes does it mint". And be careful which
instrument answers it: a fresh-pool replay -- synthetic or trace-driven --
measures how far a pool grows FROM EMPTY, not what a warm engine's
committed high-water will give back, and on the one lever where both
numbers exist the two disagreed by 100x.
``docs/arena-size-class-collapse-2026-08-31.md`` is the A/B, its null
controls, the gate comparison and that instrument check.

``granularity <= 1`` returns ``rows`` unchanged: the opt-out is the
identity, not a special case in the caller.

cuMem arena that maps many slabs' leading sentinel row onto ONE granule.

A per-request row pool reserves slab row 0 as a permanently-zero
sentinel (see :mod:`arbi_serve.cache.recurrent_pool`). Row 0 is real
VRAM in every per-layer slab and holds nothing but zeros, so the cost
of the safety net is one row's worth of state per layer — the same
bytes, repeated, all zero.

Virtual memory removes the repetition without touching the tensors.
The arena reserves ONE VA span, packs the slabs into it back-to-back,
and maps physical granule by granule: every granule that lies ENTIRELY
inside some slab's sentinel row is mapped to a single shared physical
granule, and every other granule gets private physical. Each slab is
still one contiguous tensor with a stable ``data_ptr`` and the ordinary
``(rows, *geom)`` shape, so kernels, row indexing and CUDA-graph
capture are unchanged — the only thing that changed is which physical
page a sentinel-row address resolves to.

Two properties make this safe rather than clever:

  * **Round DOWN.** Only granules wholly contained in the sentinel row
    are aliased. The partial granules at either end — which also cover
    live row 1, or the previous slab — keep private physical. Rounding
    the other way would alias live state.
  * **No caching allocator.** The arena's pages never enter torch's
    caching allocator: the slabs are non-owning views over the arena's
    VA (:func:`torch.as_tensor` on ``__cuda_array_interface__``). Torch
    therefore has no block to free and no block to recycle, so an
    aliased range can never be handed to an unrelated tensor.

Shared pages are keyed by name. Slabs whose sentinel row must READ ZERO
share the ``"zero"`` page; slabs whose sentinel row is merely never read
(rollback snapshot frames) share a different page, so a stray write into
one class cannot pollute the other.

Sleep / wake mirrors :class:`~arbi_serve.runtime.cumem_allocator.GrowableRegion`:
:meth:`sleep` offloads the PRIVATE pages to pinned host and drops all
physical behind a stable VA; :meth:`wake` re-creates the shared pages
(zeroed), re-applies every alias, and restores the private pages. The
sentinel therefore comes back zero by construction rather than by
whatever happened to be dumped.

The alias quantum on ``device``, or 0 when there is no cuMem driver.

The budget predictors need the same number the arena will use, before any
arena exists, so they can price the sentinel rows the alias removes.

Physical the alias removes for slabs with these sentinel-row sizes.

``Σ floor(row / g) * g`` (aliasing rounds DOWN — see
:func:`plan_arena_layout`) less one granule for the shared page the
aliased granules resolve to. 0 when nothing reaches a full granule, in
which case no arena is built at all.

Assumes each slab's sentinel row starts on a granule boundary, which is
what the planner's align-when-it-gains rule delivers; a layout that could
not align a slab recovers one granule less, so this is an upper bound and
the measured ledger row is the truth.

One slab to place in the arena.

``nbytes`` is the slab's full byte size; ``sentinel_bytes`` is the
leading byte range (row 0) eligible for aliasing, and ``page`` names
the shared physical page its aliased granules resolve to.

Place ``requests`` and pick the granules safe to alias.

Slabs are packed back-to-back. A slab at offset ``o`` with a sentinel
row of ``s`` bytes contributes the granules wholly inside ``[o, o+s)``
— straddling granules also cover live row 1 (or the neighbouring slab)
and stay private.

A slab is padded up to a granule boundary ONLY when that strictly
increases its aliased granule count. The padding costs less than one
granule and the gain is a whole granule, so the rule never loses; a
slab whose sentinel row is smaller than the granularity is never
padded, because nothing can be aliased for it either way.

Pure arithmetic: no CUDA, no driver, directly testable.

A committed cuMem arena whose sentinel rows share physical granules.

Build it in two steps so the layout is decided before any driver work::

    arena = SentinelAliasArena(device=0, label="state.gdn_recurrent")
    i = arena.plan(slab_bytes, sentinel_bytes=row_bytes)
    arena.commit()
    slab = arena.tensor(i, shape, dtype)

Every slab comes back zeroed. The arena's VA lives as long as the arena
does, and the slabs it hands out hold a reference back to it, so the
reservation outlives every tensor into it — the same guarantee the caching
allocator gives an ordinary tensor. :meth:`close` releases it early for a
caller that knows it is finished.

``index`` as a ``shape``/``dtype`` tensor.

A non-owning view over the arena's VA: torch's caching allocator
never sees these bytes, so the tensor is neither freeable nor
recyclable and its ``data_ptr`` is stable for the arena's lifetime.

Physical the next :meth:`wake` will map, measured now (0 while awake).

The layout is fixed at plan time and a sleep keeps it, so the arena
re-maps exactly the private chunks plus one granule per shared page it
held — the same quantity :meth:`mapped_bytes` reports when awake. A
caller that sizes against free VRAM before the wake remaps (the
growable-KV re-provision) must hold this out, and asking the arena is
the only way to get the number without predicting it.

Drop all physical, keeping the VA. Returns bytes released.

Private pages are optionally dumped to pinned host first. Shared
pages are NOT dumped: a zero page is restored by zeroing, and an
unread page has no contents worth carrying.

Re-map physical at the SAME VA. Returns bytes mapped.

Shared pages are created fresh and ZEROED, then re-aliased, so the
sentinel is zero by construction on the far side of a sleep rather
than by trusting what was dumped.

Sub-1s sleep / wake via stable-VA release / resume.

Every "sleepable" GPU tensor — model params + KV pool slabs + TKV
scratch + the persistent buffers captured cudagraphs reference — is
allocated against a pre-reserved virtual-address range using the CUDA
driver API (cuMemAddressReserve / cuMemCreate / cuMemMap). On release,
physical pages are dumped to a pinned host buffer and unmapped
(cuMemUnmap), freeing VRAM but keeping the VA reservation. On wake,
fresh physical pages are mapped at the SAME VAs and the pinned-host
contents are copied back. Captured cudagraphs survive untouched
because their data_ptrs are still valid, so no re-capture is needed;
VRAM during sleep is ~0 (true release).

This stable-VA path is the only production sleep backend. It is gated
on cuda-python being importable (true on every production build); the
engine's ``_check_phase2_driver_available`` boot guard raises if the
bindings are missing rather than silently degrading.

CPU / no-GPU stub: when cuda-python is not importable OR the pool is
constructed on a CPU device (unit tests, metrics-contract tests, the
subprocess graph-pool measure), the pool runs as an inert stub — no
cuMem reserver, no pinned host staging, register/release/resume are
no-ops over the entry list. This is NOT a second production backend;
it exists only so the metrics gauges + public-surface contracts can
be exercised without a GPU. Real boots never reach the stub because
the driver guard fails loud first.

Out of scope:
  - TP > 1 (single-rank only).
  - Cross-process sleep (this file is the same-process sub-1s path).
  - MoE per-expert weight pinning.

Module layout: the stable-VA reservation primitive + tensor wrapping
live in :mod:`arbi_serve.runtime.sleep_va`, the pinned host staging pool in
:mod:`arbi_serve.runtime.sleep_staging`, the writeback helpers + entry
record in :mod:`arbi_serve.runtime.sleep_pool_base`, and the
registration surface in :mod:`arbi_serve.runtime.sleep_pool_registration`.
Every public name (``SleepableTensorPool``, ``StableVAReserver``,
``PinnedHostStagingPool``, ``driver_available``, ``tensor_from_va``) is
re-exported here so this module stays the single import surface.

Track every GPU tensor that needs to survive a sleep / wake cycle.

Owned by the engine. Two register surfaces:

  - :meth:`register_param` — a model parameter (``nn.Parameter``)
    whose ``.data`` is pinned to a stable VA so cudagraphs survive.

  - :meth:`register_buffer` — a GPU tensor stored on some host
    object (e.g. ``pool.slab``, ``tq_pool.indptr``) reachable by
    ``setattr(host, attr, new_tensor)``.

On release: every entry's contents are DtoH-dumped to pinned host
RAM and its VA is unmapped.
On resume: every entry's VA is remapped + its contents HtoD-restored.

Stable-VA is the only production backend. When cuda-python is not
importable OR the pool is constructed on a CPU device the pool runs
as an inert stub (no reserver / staging; register/release/resume are
no-ops). The stub exists only for CPU-lane unit + metrics-contract
tests — production boots fail loud on the missing driver via the
engine's ``_check_phase2_driver_available`` guard before ever
constructing a stub pool.

The tensor-registration surface (``register_*`` /
``migrate_buffer_pre_capture`` / ``bind_entry_host`` and their
pinned-host migration helpers) is inherited from
:class:`_SleepPoolRegistrationMixin`.

Active backend: ``"phase2"`` (stable-VA, production) or
``"stub"`` (CPU / no-GPU; register/release/resume are no-ops).

Exposed for the sleep-mode metrics gauge and for the
``release_memory_occupation`` / ``resume_memory_occupation``
guards, which refuse to run on a stub pool.

True when the real stable-VA (cuMem) backend is in use.

False in the CPU/no-GPU stub. Capture + build paths key the
pre-capture VA migration + sleepable-state registration off
this — there is nothing to migrate in the stub.

Free the pinned host staging slabs. DESTROYS any parked contents.

The staging pool is deliberately allocate-once-reuse-forever: a single
engine sleeps and wakes repeatedly and re-pinning host memory each time
would be slow. Under stable-VA residency that intent inverts — every
member owns its own pool, and a member that is parked and never woken
holds its slabs (the captured-graph I/O set: per shape, ``logits_out``
rows×vocab, ``block_table``, and the rest) in unswappable host RAM for
the process lifetime. Nothing else drops them short of
:meth:`shutdown`, which also tears down the VA reserver.

So this is for a member being DESTROYED — a rolled-back build, a
dropped variant — where the slabs back state no :meth:`resume` will
ever read. Calling it on a member that will be woken loses that
member's parked tensor contents: this pool cannot tell the two cases
apart, so the caller owns the distinction. Returns the host bytes
freed.

Physical the next :meth:`resume` will re-map, measured now.

Every region's VA reservation is granularity-rounded once and never
moves, and :meth:`resume` re-binds exactly the size each unmapped
region already reserved — so reserved minus mapped IS the wake's
allocation, not a forecast of it. A wake seam that sizes anything
against free VRAM BEFORE the per-tensor restore runs (the growable-KV
re-provision does) has to hold these bytes out, and this is where it
reads them: ~0 while awake, the whole registered set while parked.

Stable-VA release: DtoH-dump every mapped region into pinned
host RAM, then unmap so physical pages return to the pool while
the VA reservation lives on (so cudagraph data_ptrs stay valid
across the cycle).

Re-bind physical memory for every registered tensor.

Re-binds physical pages at the same VAs, then HtoD copies the
pinned-host parked contents back via N worker threads × M-MiB
chunks across the pool's own CUDA streams. ``data_ptr`` of every
registered tensor is unchanged so captured cudagraphs are still
callable.

When :meth:`start_preload` was called and its future is
already complete, the cached result is consumed in place of a
fresh restore — the H2D bytes are already on device.

In the CPU/no-GPU stub this is a no-op that just clears the
released flag.

Greedy bin-pack the entry list across ``_wake_workers``.

Each partition is a list of ``(va, nbytes, slab_id)`` tuples
sorted by va so the worker walks sequential pinned-host
memory. ``slab_id`` keys back into ``self._staging`` for the
pinned-host source.

One worker: ping-pong the partition's (va, slab) pairs through
``stream`` in ``chunk_bytes``-sized chunks.

Returns total bytes copied. The worker records an event at the
end of each launch and waits on it before reusing the stream
for the next chunk so dispatch stays ahead of the device.

Multi-stream chunked H2D restore for every registered entry.

Caller must have remapped physical pages at every VA before
calling. Returns total bytes copied. Synchronizes every worker
stream against the current default stream before returning so
downstream kernels see the new contents.

Re-attach Python-side tensor refs to the VA-backed tensors.

Defensive: we keep each entry's ``target`` alive across the
cycle so this is normally a no-op. The check exists to catch
future paths that swap ``param.data`` out from under us.

Synchronous wake: remap + chunked-MT H2D + rebind.

Used by both :meth:`resume` (sync caller) and
:meth:`wake_async` (via ``asyncio.to_thread``). When a
:meth:`start_preload` future is already complete, this short-
circuits to consuming that result instead of redoing the H2D.

Kick off the chunked-MT H2D restore in the background.

Returns a future that resolves to ``bytes_restored`` once the
H2D pipeline finishes. Idempotent: if a preload is already in
flight (or already complete and unconsumed), the existing
future is returned.

Stub pools have nothing to preload, so this returns an
immediately-resolved future with 0.

The future is consumed by the next :meth:`resume` /
:meth:`wake_async` call. :meth:`release` cancels any pending
future before re-dumping device state.

Async counterpart to :meth:`resume`.

Routes the chunked-MT H2D restore through ``asyncio.to_thread``
so the event loop never blocks on the multi-second H2D pipeline.
Consumes any outstanding :meth:`start_preload` future in place
of starting fresh.

Release every VA reservation + drop the pinned host pool.

Safe to call multiple times. After the first call the pool is
marked closed: any further ``register_*`` raises. Releasing a
closed pool is a no-op (release/resume bail on idempotence).

Shared state for the sleepable-tensor pool: writeback + entry record.

Holds the dict-aware writeback helpers (:func:`_wb_set` / :func:`_wb_get`)
and the :class:`_SleepEntry` record, shared by the registration mixin and
the pool's release/resume path so neither has to import the other.

Write ``value`` to ``host.attr`` for sleep/wake re-attach.

Prefers a dict-aware ``set_persistent_buffer(name, tensor)`` hook when
the host exposes one (``CapturedGraph`` keeps its recurrent
``state_indices_{kind}`` buffers in a dict, not a settable slot, so a
plain ``setattr`` raises ``AttributeError`` on that slotted dataclass).
Everything else (modules, adapters) takes the ``setattr`` path.

One registered tensor + its stable-VA release / resume state.

On release, the tensor's contents are copied to a pinned-host
staging slab and the VA is ``cuMemUnmap``-ped. On resume, a fresh
handle is ``cuMemMap``-ped at the SAME VA and the staged contents
are copied back. Captured cudagraphs survive because ``data_ptr``
(the VA) is unchanged across the cycle.

In the CPU / no-GPU stub there is no ``region`` — register/release/
resume are no-ops over the entry list (see the module docstring).

Registration surface for :class:`SleepableTensorPool` (mixin).

:class:`_SleepPoolRegistrationMixin` carries every ``register_*`` /
``migrate_*`` surface plus their pinned-host migration helpers. It is
mixed into ``SleepableTensorPool``; the methods rely on the instance
attributes wired by that class's ``__init__`` (``_reserver``,
``_staging``, ``_entries``, ``_stub``, ``_closed``, ``device``).

Guard the pool is open and build a ``_SleepEntry`` from
``tensor``'s metadata (nbytes / shape / dtype), with ``target``
pointing at ``tensor``. Shared prologue for every ``register_*``
/ migrate surface; the caller wires the VA region + writeback.

Reserve a stable VA, copy ``src``'s contents in, record the
region on ``entry`` and return the VA-backed tensor.

The caller performs the write-back (``param.data = ...`` for a
parameter, ``_wb_set`` + ``entry.target`` for a buffer).

Register an ``nn.Parameter`` for sleep/wake.

``host_object`` + ``host_attr`` is how we re-attach the
parameter on resume — typically ``(parent_module, param_name)``
but any setattr-able pair works.

Allocates a stable VA + copies the param's contents in.
**Holds peak GPU memory at 2× the param's size** during the call
(original storage + new VA-backed clone live concurrently until
the function returns and the original is dereffed). For bulk
model registration where this 2× holdover summed across every
param dominates VRAM, prefer :meth:`register_param_in_place`
which goes via pinned host RAM and keeps peak at ~1× the
per-tensor footprint.

In the CPU/no-GPU stub this just records the entry — no VA, no
copy.

Register a non-Parameter GPU tensor (pool slab, scratch, ...).

Returns the (possibly new) tensor that ``host_object.host_attr``
now references — the caller MUST use the returned tensor in
any further references because we swap it for a VA-backed clone.

Same 2×-peak warning as :meth:`register_param`. For bulk
registration use :meth:`register_buffer_in_place`.

In the CPU/no-GPU stub this just records the entry and returns
the original tensor unchanged.

Walk one tensor through the pinned-host migration sequence.

Returns ``(region, new_va_backed_tensor)``. Caller is responsible
for swapping the live reference (``param.data``,
``setattr(host, attr, ...)``) to ``new_va_backed_tensor`` and
releasing all references to the original ``tensor``.

Only the bytes are copied; the new tensor's logical layout is
contiguous in the requested shape/dtype. This is correct iff the
source ``tensor`` is contiguous (every model param + KV slab
meets that, but assert defensively).

In-place stable-VA migration of an ``nn.Parameter``.

Same final state as :meth:`register_param` (param.data wraps a
stable-VA buffer, registered for sleep/wake) but peak GPU
memory during the call is ~1× the param size instead of 2×.

Caller MUST drop every external reference to the original
``param.data`` after this returns and run
``torch.cuda.empty_cache()`` periodically (or after the bulk
loop) for the caching allocator to release the original
segments back to the driver. The engine's
``_register_sleepable_state`` does both.

In the CPU/no-GPU stub this just records the entry — no VA, no
migration.

In-place stable-VA migration of a non-Parameter GPU tensor.

Same final state + caller contract as :meth:`register_buffer`.
Returns the new VA-backed tensor (production) or the original
(CPU/no-GPU stub).

WARNING — cudagraph persistent buffers. On a real GPU this
MIGRATES the tensor's bytes to a fresh stable VA and **changes
its ``data_ptr``**. That is fatal for a buffer a captured
cudagraph ALREADY references: the graph baked the pre-migration
address into its kernel args, so post-capture migration leaves
the graph reading the now-stale original address while the live
replay path writes the new VA (manifests as
``vectorized_gather_kernel: index out of bounds`` on the first
replay — the captured embedding / page-metadata gather reads a
garbage index buffer). For capture-persistent buffers use
:meth:`migrate_buffer_pre_capture` BEFORE the graph is captured
so the graph bakes the stable VA from the start.

Migrate a soon-to-be-captured buffer to a stable VA *now*.

Identical byte-for-byte to :meth:`register_buffer_in_place`'s
migration step, but the sleep-pool entry's ``host_attr`` is left
UNBOUND so the caller can attach the real host (typically the
:class:`CapturedGraph` that doesn't exist yet at this point)
afterwards via :meth:`bind_entry_host`.

Call this AFTER allocating the persistent buffer but BEFORE the
cudagraph capture region. The capture then bakes the returned
VA-backed tensor's stable ``data_ptr`` into its kernel args, so
the live replay path (which writes the same tensor) and the
captured graph (which reads it) agree — and release/resume keeps
that ``data_ptr`` stable.

Returns ``(va_backed_tensor, entry)``. In the CPU/no-GPU stub
this is a no-op registration: the original tensor is returned
and the entry carries no VA. The caller still binds the host.

Attach the writeback target for an entry created by
:meth:`migrate_buffer_pre_capture`.

Idempotent re-binding: updates ``entry.host_attr`` so the
defensive rebind in :meth:`_rebind_targets` routes through the
right host object. Does NOT move any memory — the migration
already happened in :meth:`migrate_buffer_pre_capture`.

Pinned host staging pool used to park VRAM contents during sleep.

The :class:`PinnedHostStagingPool` allocates page-locked host RAM once
per process and reuses it across sleep/wake cycles at full PCIe
bandwidth. Re-exported from ``sleep_allocator`` for the existing public
surface.

Pre-allocated pinned host RAM used to dump VRAM during sleep.

Allocated ONCE at boot (or lazily on first release) and reused
across cycles. Sized by the registered total tensor bytes — i.e.
just enough to park every sleepable tensor at once.

Pinned host alloc is slow, but we pay it once per process.
Subsequent sleep cycles reuse the same buffer at full PCIe
bandwidth.

Get-or-allocate a pinned slab keyed by ``key`` of ``nbytes``.

Every slab is booked through the host ledger BEFORE it is page-locked
(:mod:`arbi_serve.runtime.pinned_host_budget`), so a sleep that would
park more VRAM in pinned host RAM than the host can hold is REFUSED
with a :class:`~arbi_serve.runtime.pinned_host_budget.PinnedHostBudgetError`
naming the bytes. There is nothing to shrink here — a partial park is
not a park — so refusing to sleep is the whole of the alternative to
being killed by the kernel while asleep.

Stable-VA reservation primitive + zero-copy tensor wrapping.

Owns the CUDA virtual-memory driver-API gate (``_DRIVER_AVAILABLE`` /
:func:`driver_available`), the :class:`StableVAReserver` that reserves VAs
and binds/unbinds physical pages on demand, and :func:`tensor_from_va` /
:func:`_byte_view` which wrap a stable-VA range as a zero-copy torch
tensor. Re-exported from ``sleep_allocator`` for the existing public
surface.

Pre-reserve VAs and bind/unbind physical pages on demand.

Each :meth:`reserve_and_map` returns a :class:`_VARegion` whose
``va`` is stable across release/resume cycles — :meth:`unmap_all`
drops physical pages but keeps every VA, :meth:`map_all` re-binds
fresh pages at the SAME VAs. Any ``torch.Tensor`` constructed
against ``region.va`` (via :func:`tensor_from_va`) keeps the same
``data_ptr`` across cycles, which is the contract captured CUDA
graphs depend on.

Single-device / single-rank only — TP>1 needs rank-aware
reservation.

Trivial holder for the ``__cuda_array_interface__`` dict.

``torch.as_tensor`` walks ``__cuda_array_interface__`` to wrap
arbitrary GPU pointers as zero-copy tensors. Used to bind a
stable-VA region to a torch tensor without going through the
caching allocator.

Wrap a stable-VA range as a zero-copy torch tensor.

The returned tensor shares storage with the VA region (no copy);
its ``data_ptr()`` equals ``va``. Caller is responsible for keeping
the underlying VA reservation alive at least as long as the tensor.

bf16 has no numpy typestr — we bind as opaque 2-byte and ``view()``
to bf16 after wrapping. Same for any other torch dtype not in the
numpy table.

Reserve a VA range and back it with fresh physical pages.

``requested_size`` is rounded up to the device granularity.
Returns a :class:`_VARegion`; the caller is responsible for
wrapping it as a torch tensor via :func:`tensor_from_va`.

Split a MIXED decode+prefill step into a captured-decode replay + an
eager-prefill forward (``ARBI_SPLIT_MIXED_DECODE_PREFILL``).

THE BUG this closes
-------------------
At concurrency the continuous-batching scheduler co-admits a few PREFILL
rows into the same slate as the N steady-state DECODE rows. In
``batch_build.build_batch`` the per-step ``is_prefill`` flag is the
aggregate ``any(per_row_is_prefill)`` — so ONE prefill row flips the
WHOLE batch to ``is_prefill=True``. Every captured-decode-graph lookup in
:mod:`arbi_serve.runtime.captured_lookup` then ``return None``s on
``batch.is_prefill``, and the entire step (including the N decode rows
whose B=N graph IS captured) falls to the eager per-layer launch path —
throwing away the captured decode graph on every co-admit step. That
inflates decode TPOT and wastes GPU at concurrency.

THE FIX
-------
When the slate is MIXED (carries BOTH single-token decode rows and
prefill rows), split the step into two forwards over DISJOINT row sets:

  1. a decode-only sub-batch (the decode rows) → replay the captured
     ``B=Ndecode`` whole-forward graph (or decode-pad to the nearest
     captured rung). This is the fast path currently being skipped.
  2. a prefill-only sub-batch (the prefill rows) → the whole-forward
     captured prefill graph when one covers it (``ARBI_SPLIT_PREFILL_CAPTURE``,
     default OFF — the same replay the non-split path uses), else the
     normal eager forward.

The per-row logits are then scattered back into the original slate order
so the sampler sees one ``(num_seqs, vocab)`` tensor.

Correctness
-----------
Each request's recurrent (GDN / Mamba) state is indexed by
``recurrent_state_indices`` (its allocated slab row) and its KV is paged
per-request, so the decode rows and prefill rows touch DISJOINT state and
KV — running them as two forwards is correct: a prefill row's forward
cannot read or write a decode row's state, and vice versa.

The merge is greedy-argmax-identical, NOT bitwise, to the legacy fused
eager forward. The only numeric difference is the GEMM batch dimension
``M`` (``Ndecode`` for the decode sub-forward vs the fused ``B``), whose
cuBLAS split-K reduction order shifts the last bf16 bits — the SAME
intrinsic nondeterminism the decode-pad path (:mod:`decode_pad`) already
documents and the engine already accepts between any two batch sizes.

The whole path is gated behind ``cfg.split_mixed_decode_prefill`` (env
``ARBI_SPLIT_MIXED_DECODE_PREFILL``); default ON (greedy argmax-identical;
the only difference is the bf16-LSB cuBLAS reduction-order nondeterminism the
decode-pad path already ships default-on). Set
``ARBI_SPLIT_MIXED_DECODE_PREFILL=0`` for the legacy byte-identical fused path.

How the sub-batches are built
-----------------------------
The full ``ScheduledBatch`` is materialized ONCE up front (so the page
table / slot allocation / recurrent-row resolution happen exactly as
today — those are NOT idempotent and must not be re-run). The two
sub-batches are then derived by ROW-SLICING the materialized flat
tensors:

  * per-row tensors (``seq_lens`` / ``block_table`` / ``state_indices``)
    index by slate row ``i``;
  * the flat token tensors (``input_ids`` / ``positions`` /
    ``slot_mapping``) index by the token span
    ``[cu_seqlens_q[i], cu_seqlens_q[i+1])`` of row ``i``;
  * ``cu_seqlens_q`` / ``cu_seqlens_k`` are rebuilt for the sub-batch.

Decode rows contribute exactly one token each (``n == 1``), so the decode
sub-batch is dense (Ndecode tokens). Prefill rows contribute ``n`` tokens
each. The block_table page-width is kept at the full batch's width so
both sub-batches' attention kernels read the same padded table the fused
path would have.

Row count the DECODE sub-batch would carry for ``slate``.

The cost gate's independent variable. Counted from the slate's own
host-side per-row facts, so this reads nothing off the device and
inserts no sync. A "decode row" is the same thing
:func:`_mixed_composition_reason` counts: non-prefill, one token.

Largest row count the boot activation profile measures on the
PREFILL-class path, for the serving widths this engine is configured at.

ASKED OF THE PROFILER, not restated. The value is the widest ``num_seqs``
over the prefill-class shapes
:func:`~arbi_serve.runtime.activation_profile._profile_shape_specs`
actually emits, so the guard cannot drift from the measurement it stands
for. The previous form of this function hard-coded ``1 + max_batch // 2``,
which was the mixed probe's row count only while
``max_batched_tokens <= chunk_prefill + max_batch // 2`` — true at the
shipped defaults and false the moment an operator widens the step budget.

Every input the profiled shapes are a function of is a parameter here and
all are REQUIRED: a call site that could omit ``chunk_prefill`` or
``max_batched_tokens`` would silently model a different engine's envelope,
and the failure mode is a capacity guard that reads too wide — which is
the guard failing OPEN, on the path that OOMed the CUDA context.

Returns ``0`` (guard disabled) when any width is not a usable int — an
unresolved ``"auto"`` must never take down the per-step hot path.

Row count the PREFILL forward would carry if this step does NOT split.

Declining the split folds the decode rows into the prefill forward,
so the un-split prefill forward carries the WHOLE slate. This is the
capacity guard's independent variable — see
:func:`_profiled_prefill_row_envelope`.

Return True when this step may take the split decode/prefill path.

Three conditions, all required:

* the slate is genuinely MIXED (:func:`is_mixed_decode_prefill`),
* the batch carries **no media plan**, and
* the decode sub-batch is BIG enough that splitting still PAYS
  (:func:`decode_sub_batch_rows` ``>= min_decode_rows``).

The split's price is a fixed second read of the model weights; its
return is per decode row and grows with the row count, so it is the
row count — not the prefill chunk size — that decides. See
:data:`~arbi_serve.runtime_flags.SPLIT_MIXED_MIN_DECODE_ROWS_DEFAULT`
for the measurement, including the pair that refutes the
prefill-chunk shape.

``min_decode_rows`` is keyword-only and REQUIRED — the caller resolves
it through
:func:`~arbi_serve.runtime_flags.split_mixed_effective_min_decode_rows`,
which is where an operator's explicit setting is honoured. No default
is offered on purpose: a defaulted threshold would let a new call site
silently reintroduce the ungated split, which is precisely the
regression the gate exists to prevent.
:data:`~arbi_serve.runtime_flags.SPLIT_MIXED_MIN_DECODE_ROWS_UNGATED`
(0) is the explicit way to spell "no cost gate".

The media refusal is a CORRECTNESS gate, not an optimization choice.
``attach_multimodal`` builds ``batch.mm`` (placeholder ``token_mask`` +
the tower's encoded ``embeds``) and the M-RoPE ``mrope_positions`` grid
against the WHOLE batch's flat token stream. :func:`_slice_sub_batch`
constructs a fresh :class:`ScheduledBatch` and carries NEITHER field,
so a split multimodal step hands the model a prefill sub-batch with
``mm = None``: :func:`merge_multimodal_embeddings` is then skipped
entirely, the placeholder rows keep their raw ``<image>`` text
embeddings, and the model prefills a prompt that never contained the
image. The merge's count guard cannot catch it — the merge does not
run — so the request silently prefills garbage and decode samples EOS
on the first step (empty transcription).

This composition only arises at concurrency >= 2 (a media prefill row
co-admitted with decode rows).

Carrying a correctly SLICED ``mm`` / ``mrope_positions`` into the
sub-batch would let multimodal steps keep the split's captured-decode
replay; that is a perf follow-up, not a correctness one.

Every ``False`` is attributed on the ``split_mixed_forward`` flag-truth
counter with the gate that declined, so "the split never ran" is
diagnosable per reason instead of a silent eager fall-through.

Return why this step may NOT split, or ``None`` when it may.

The single decision core behind :func:`is_splittable`. Order is
deliberate: the media refusal is a CORRECTNESS gate and comes first,
then the mixed-composition walk (:func:`_mixed_composition_reason`),
and only then the COST gate — which is the one refusal that says "this
step could have split and chose not to", so it must not mask a
correctness or composition refusal above it in the attribution.

Return why ``slate`` is NOT "decode rows + a prefill row", or ``None``.

A multi-token non-prefill row (verify / MTP) seen before both classes
are present bails; both classes present ⇒ mixed. The decision core
shared by :func:`is_mixed_decode_prefill` and
:func:`_split_refusal_reason`.

Return True when ``slate`` carries BOTH a single-token decode row
and a prefill row — the only composition the split helps.

A "decode row" is a non-prefill row with ``n == 1`` (the captured
decode graphs are all ``S=1``). A row with ``n > 1`` that is not
flagged prefill (verify / spec) is treated as non-decodable here and
forces the legacy fused path (returns False) so we never mis-route a
verify step. Decision core: :func:`_mixed_composition_reason`.

True when ``rows`` are CONSECUTIVE slate indices (``i, i+1, …``).

Empty / single-row lists are trivially adjacent. ``rows`` comes from
:func:`partition_rows` in ascending slate order, so adjacency ⇔ the
kept rows form one contiguous slate range — and therefore their
flat-token spans ``[cu_q[i], cu_q[i+1])`` tile ONE contiguous range.

Build a standalone ``ScheduledBatch`` from the ``rows`` of ``batch``.

``rows`` is the list of slate-row indices (into the full batch) to
keep, in their original order. ``cu_q`` is the full batch's host-side
``cu_seqlens_q`` list (token-boundary cumsum, length ``B+1``) used to
locate each row's flat-token span.

Returns a fresh batch whose flat token tensors are the concatenation
of the kept rows' spans and whose per-row tensors are the kept rows.
``cu_seqlens_q`` / ``cu_seqlens_k`` are rebuilt; ``block_table`` keeps
the full page-width. All scheduling metadata (``*_meta``) starts empty
— the caller rebuilds it for the eager sub-batch and the captured
replay does not consult it.

``eng`` supplies the FLA host-twin decision. The parent's
``cu_seqlens_q_cpu`` describes the parent's boundaries and CANNOT be
carried onto a sub-batch whose ``cu_seqlens_q`` is rebuilt — so the
twin is rebuilt here too, from the same host cumsum that sources the
device tensor, through the one seam every batch construction shares
(:func:`~arbi_serve.runtime._batch_build_helpers.host_cu_seqlens_q_twin`).
Leaving it unset instead is not "a field the split does not need": the
compiled decoder layer GUARDS on its None-ness
(``_gdn_fla_forward._resolve_cu_seqlens``, inside the ``fullgraph=True``
region), so a sub-batch that drops it presents an uncompiled branch and
the prefill sub-forward JIT-compiles INSIDE a live request
(``jit_compile_serving``).

Hot-path note. When the kept rows are ADJACENT slate rows (the common
split shapes — e.g. the co-admitted prefill rows are one contiguous
slate run), every sliced tensor is a contiguous range of the source,
so we return zero-copy ``narrow`` VIEWS: no per-token Python index
list, no blocking H2D, no gather kernel. Non-adjacent rows fall back
to ``index_select`` with the index built on CPU (vectorized
``arange``/cat) and shipped in one
``non_blocking`` H2D — value-identical to the old list-built gather.
``max_seq_len`` is taken from the full batch (a safe upper bound used
only for kernel/metadata sizing) instead of a per-call
``seq_lens.max().item()`` D2H sync.

Rebuild ``lora_state`` on a sliced sub-batch. No-op without LoRA.

:func:`_slice_sub_batch` carries ``lora_assignments`` (a per-row list) but
CANNOT carry the parent's ``lora_state``: that object is built against the
parent's flat token layout — its per-token BGMV routing indices are indexed
by the parent's ``cu_seqlens_q``. Handing it to a sub-batch whose token
count and row order differ would mis-route every token.

So it must be REBUILT from the sub-batch's own ``lora_assignments`` +
``cu_seqlens_q``, exactly as :func:`batch_build.build_batch` does for the
fused path and :meth:`model_runner._reconstruct_verify_batch` does for the
verify path.

Leaving it ``None`` is silently WRONG twice over: the model reads ``batch.lora_state is None`` and short-circuits every LoRA
linear to the base weight, AND ``captured_lookup.lora_bucket_for`` reads the
same field, so it returns bucket 0 and hands the decode sub-batch the
**no-LoRA** captured graph.

``use_capture_pool`` is set only for the decode sub-batch, whose captured
graph replay requires the state's tensors to live at the pool's persistent
``data_ptr``s. The prefill sub-batch never replays a captured graph while a
LoRA is in flight (``lookup_captured_graph_for_prefill`` refuses on
``req.lora_id``), so it builds fresh per-step tensors — which also keeps it
from overwriting the pool buffers the decode replay above just consumed.

Run the decode rows through their captured graph and the prefill
rows eager, then merge per-row logits in slate order.

Returns a ``(num_seqs, vocab)`` logits tensor row-aligned to ``slate``
(so the existing :func:`forward_exec.sample` / ``_sample_async``
indexing is unchanged). The caller has already materialized ``batch``
(page-table / slot / recurrent-row resolution done) and flushed the
per-step zero-clears / savepoint resumes.

Reuses the existing captured-graph lookup + decode-pad machinery for
the decode sub-batch (the fast path the fused ``is_prefill`` flag was
skipping) and the canonical ``_build_metadata`` + ``_run_model_forward``
for the prefill sub-batch.

Write ``src`` (``(len(rows), vocab)``) into ``out`` at row indices
``rows`` — value-identical to ``out[rows] = src`` (rows are unique).

Adjacent rows (the common split shapes) use a zero-copy ``narrow`` +
``copy_`` (no index tensor at all); non-adjacent rows use
``index_copy_`` with the index built on CPU and shipped in one async
H2D.

Multi-model stable-VA residency — recapture-free model switching.

This is the **distinct-VA-range-per-model** scheme (``docs/memory-accounting.md``,
"Release, wake, swap"). It keeps the parked model's **VA
reservation + captured graph objects** resident in the live process and
swaps only the **physical** VRAM backing.

The crux insight: a captured CUDA graph bakes the VIRTUAL
addresses of the tensors it touches, not the physical pages. If every model's
working set lives in a DISJOINT **no recapture**.

Layering
--------
* :class:`VaArena` — a fixed-base, non-overlapping VA reservation for one
  model's whole working set (built on the cuMem ``reserve_va_at`` primitive). It
  is GPU-touching only at :meth:`VaArena.reserve` / :meth:`VaArena.free`; the
  registry below is pure bookkeeping + injected mechanics so it is CPU-testable.
* :class:`StableVaResidencyRegistry` — holds one :class:`ResidentModelRecord`
  per prepared model, enforces the single-physical-at-a-time invariant, and
  drives ``park_fn`` / ``wake_fn`` (injected mechanics) so the whole state
  machine is unit-testable on CPU with no GPU.

The DEFAULT single-model path never constructs this — it is additive / opt-in.

A fixed-base, non-overlapping VA reservation for one model's working set.

Reserves ``size`` bytes of VA at the deterministic base ``base_va`` (no
physical) so every allocation the model makes inside it bakes addresses in
``[base_va, base_va+size)``. Sub-allocation inside the arena uses the
existing cuMem allocator / :class:`GrowableRegion` machinery unchanged — the
only new thing is WHERE the arena sits.

Mechanism only: :meth:`reserve` / :meth:`free` touch the cuMem driver; the
object is otherwise plain bookkeeping (so the registry that owns arenas is
CPU-testable).

One prepared model's residency record.

Holds the model's VA arena + the phase-1 artifacts that are RETAINED while
the model is parked (so a switch back is recapture-free): the captured graph
objects, the pool/metadata bookkeeping, and the weights source for flat-load
on wake. ``mapped`` tracks whether the model's physical VRAM is currently
mapped (exactly one record is ``mapped`` at a time — the single-GPU
invariant).

Registry of per-model VA-arena residency records + the switch state
machine.

Enforces the single-physical-at-a-time invariant: at most one record is
``mapped`` (physical VRAM backing its arena). A switch parks the active
record (unmap physical, keep arena + graphs) then wakes the target (map
physical + flat-load weights). Both records' VA reservations + graph objects
coexist (they cost ~no VRAM) — that is what makes the switch back
recapture-free.

``park_fn`` / ``wake_fn`` are injected so the state machine is CPU-testable
with no GPU. The registry NEVER touches
the GPU itself — it owns ordering, arena placement + disjointness, and the
active pointer.

Reserve the arena's VA range at its fixed base (idempotent).

Returns the granted base (== ``base_va`` by construction — the cuMem
shim fails the reservation if the driver grants a different base). The
range costs NO physical VRAM; physical is mapped into sub-ranges by the
model's pools / GrowableRegions later. Raises if the fixed base is taken
(the caller must pick another base / stride).

Create (and register) a distinct, non-overlapping VA arena for a new
model ``key``.

Places the arena at the next free slot's deterministic base. Asserts
disjointness against every existing arena (a defensive invariant — the
slot stride guarantees it by construction, but we verify so a
mis-sized arena_bytes can never silently alias). Does NOT reserve the VA
with the driver yet (call :meth:`VaArena.reserve` when bringing the
model up) — this is pure placement so it is CPU-testable. Raises if
``key`` is already registered.

Wake model ``key``: map physical at its arena + flat-load its
weights. Its graphs are immediately valid (their baked VAs never moved).
Runs ``wake_fn``. The single-GPU invariant requires no OTHER record be
mapped first — :meth:`switch_to` enforces park-then-wake; calling
:meth:`wake` directly while another is mapped raises.

Make model ``key`` the active (physically resident) model.

Strict park-then-wake (single-GPU): the active model's physical is
unmapped BEFORE the target's is mapped, so two models' physical never
coexist. The target's graphs are recapture-free because its VA arena +
graph objects were retained while parked. Returns a result dict the
caller can stamp wall-time + a recaptured-flag onto.

Result::

    {
        "target": key,
        "prev": <prior active key | None>,
        "noop": bool,           # target already active
        "parked": <key | None>, # what we parked (None if nothing)
        "recaptured": False,    # this path NEVER recaptures, by design
    }

De-register a record: free its VA arena and forget it.

Used to abandon a freshly-registered member whose build the caller
is rolling back (e.g. the TP>1 build-required confirmation came back
negative). The record's PHYSICAL must already be released by the
caller (the engine's orphan teardown); this only drops the
bookkeeping + the VA reservation. Refuses to drop the active record
(park/wake the replacement first) so the registry never points
``_active`` at a freed arena.

The STEP CLASS: how many tokens each SEQUENCE contributes to this step.

One token per sequence is decode. More than one is either a prefill chunk or
a speculative verify slate. The distinction is a property of the REQUEST — its
remaining prompt, its drafter's ``K`` — never of how many other requests happen
to be in flight, which is what makes it usable as a numerics-routing key.

WHY A CLASS AND NOT A ROW COUNT. A kernel picked from the step's TOTAL row
count makes one sequence's arithmetic a function of the server's load. The
EXL3 GEMM does exactly that today and the cost is recorded at the
``run_model_forward`` marking site: the same prompt takes the trellis leg
alone and reconstruct+hgemm alongside three others, which moves 57.5% of a
narrow linear's output. ``exl3_prefill_row_invariant`` exists to flag that
defect; routing on the class removes it, because ``B``, ``B*4`` and ``B*32``
rows of the same class all resolve here to the same answer.

THE SAME SHAPE THE ATTENTION SIDE ALREADY SHIPS. turbo-attn's
``decode_attend._declines_verify_to_prefill_split`` hands a verify chunk to the
Turbo prefill kernel — the PREFILL attention kernel — at ``block_m >= 2``,
i.e. at ``tokens_per_seq``,
and ``TKV_MTP_PREFILL_SPLIT`` is default-on, so essentially all MTP/DFlash verify
attention already runs on the prefill kernel. ``block_m`` there is
:func:`tokens_per_seq` here.

WHY A MIXED STEP CANNOT DEFEAT IT. ``batch.is_prefill`` is the aggregate
``any(per_row_is_prefill)``, so one prefill row would otherwise flip a whole
mixed step to PREFILL while its decode rows rode along. The engine does not
run such a step as one forward: :mod:`arbi_serve.runtime.split_mixed` issues
two forwards over DISJOINT row sets, each of which classifies here on its own.
A batch reaching this function therefore never contains both classes.

RELATION TO :func:`arbi_serve.runtime.step_mix_stats.classify_eager_forward`.
That one answers "why did this forward miss a captured graph" and is
debug-only; it reads the same three host-side facts (``is_prefill``,
``max_query_len``, ``mtp_meta``) because there is only one way to tell the
classes apart without a device read. This module is the PRODUCTION classifier
— it is on the default forward path, so it stays free of the capture
vocabulary and of anything that could sync.

NO DEVICE READS. Every fact is host-side and already materialised on the
batch, so classification cannot insert the first-op-after-forward sync the
SPMD loop was purged of, and every TP rank derives the same answer from the
batch it built identically — no broadcast.

What kind of step is running, by tokens-per-sequence.

``DECODE``
    Every sequence contributes exactly one token. The class a
    row-count-keyed dispatch cannot keep stable, because at batch 256 a
    decode step is 256 rows and crosses every row threshold a prefill
    chunk does.
``PREFILL``
    At least one row is a prompt chunk. Taken from ``batch.is_prefill``
    rather than from the chunk width, so a request whose FINAL chunk is
    one token stays in this class — a chunk width is set by how the
    scheduler split the token budget, which is the concurrency dependence
    this whole module exists to remove.
``VERIFY``
    A speculative slate: ``B x (K+1)`` rows, ``K`` from the request's own
    drafter config. Not ``is_prefill`` — verify rows are decoding rows —
    which is exactly why ``is_prefill`` alone is not a sufficient key.

Tokens this step asks of its WIDEST sequence.

``max_query_len`` is a running max over the slate's per-row token counts,
computed on the host while the batch is built. 1 on plain decode, ``K+1``
on a verify slate, the chunk width on prefill.

Defensive default of 1: a batch-like object that never set the field is
treated as decode, which is the fail-safe direction — decode is the class
whose arithmetic every other class is compared against.

Rows this forward hands every linear — the step's flat token count.

ONE NUMBER PER FORWARD, not per linear: every linear in the model sees the
same ``(N_tokens, in_features)`` activation, so the width a leg decision
is taken on is a property of the STEP. On a verify slate it is
``B x (K+1)``, which is what ``exl3_int8_verify=auto`` compares against
its threshold.

HOST-SIDE AND SYNC-FREE, which is the whole constraint. ``input_ids`` is
already materialised when the batch is built and a tensor's ``shape`` is
host metadata, so this reads no device memory and allocates nothing — the
same discipline :func:`tokens_per_seq` keeps, for the same reason: this is
on the production forward path and a sync here would land in the decode
loop the SPMD path was purged of.

Defensive 0 for a batch-like object with no ``input_ids``: 0 is below
every threshold, so an unreadable width refuses the new leg rather than
taking it — the fail-safe direction, since the leg being refused is the
one the shipped dispatch already serves.

Classify one step. Pure, host-side, and free of concurrency.

The PREFILL test comes first and is deliberately ``is_prefill`` rather
than a width: a one-token final prompt chunk is still prefill, and letting
its width decide would put the same prompt in different classes depending
on how the budget was split that step.

Per-forward captured-vs-eager step-mix accounting (debug-only).

Answers, for one serving run, the two questions the per-shape
``capture_hist`` histograms cannot: WHY a forward fell to eager, and how
much GPU time each dispatch class consumed. Populated only when
``ARBI_DEBUG_CAPTURE_LOOKUP=1`` (the same flag that gates the existing
capture histograms); the production path never constructs an event or
touches a dict.

Design constraints:

  - **Zero host↔device syncs.** GPU time rides paired CUDA events
    recorded on the current stream around each forward; elapsed times
    are harvested opportunistically on later calls via ``Event.query()``
    (never ``synchronize``). A bounded in-flight pool caps event churn;
    when the pool is exhausted the call is still COUNTED, just not
    timed (``counts`` is exact, ``gpu_ns`` is a floor).
  - **Rank-symmetric and read-only.** Recording never branches on
    anything a peer rank could disagree on — it observes the dispatch
    decision, it never influences it.
  - **No device reads in classification.** The eager-cause classifier
    uses only host-side batch facts (shapes, ``is_prefill``,
    ``max_query_len``, ``mtp_meta`` presence) — reading tensor VALUES
    here would insert the exact first-op-after-forward sync the SPMD
    loop was purged of.

Surfaced via ``GET /v1/admin/capture_hist`` (rank 0) and the per-rank
periodic step-mix log line in the SPMD loop (every rank).

Name the reason a ``forward`` call fell through to the live path.

Host-side facts only (see module docstring). The S>1 causes mirror
the refusal order of ``lookup_captured_graph_for_forward`` /
``.._padded``: missing ``mtp_meta`` and a ragged (non-uniform-K)
flat batch each refuse BOTH the exact and the pad-up lookup, so
they are terminal causes; a uniform S>1 batch that still missed is
off-ladder (no exact rung, no pad candidate above it). An S==1
batch that missed the exact rung is off-ladder by definition — the
pad-up lookup refuses ``S<=1``.

Counts + CUDA-event GPU timing per forward dispatch class.

One instance per :class:`EagerModelRunner`. All methods are cheap
enough for the per-forward hot path of a DEBUG run; the production
path (``debug_capture_lookup`` off) never calls them.

Attribute every host<->device sync on the serving path to its call site.

``torch.cuda.set_sync_debug_mode("warn")`` reports a sync as a bare
``UserWarning`` naming the aten op that blocked. That says WHAT synchronized
and not WHO asked, and the two are rarely in the same file: the op is
``aten::item`` and the caller is a scheduler accounting loop six frames up.
Without the caller a census cannot key its counts on anything actionable.

So the reporter here answers the second question. It emits ONE line per sync
carrying the frames that led to it, oldest first, and hands every other
warning to the reporter it replaced.

One line per sync is deliberate rather than terse. A census is a count per
chunk (prefill) or per step (decode) keyed on a call site, which is
``sort | uniq -c`` over the run's log; a multi-line traceback per sync does
not survive that, and a per-site aggregate held in the process cannot be read
back without an endpoint built to read it. The line is the record.

Diagnostic only: nothing installs this unless ``ARBI_SYNC_DEBUG`` is set, and
``set_sync_debug_mode`` itself is what makes the warnings exist at all.

Route torch's sync warnings through a call-site-naming reporter.

Idempotent: a second call replaces nothing, so a re-entered boot path
cannot chain reporters into a recursion. Every non-sync warning goes to
the reporter that was installed before, unchanged.

Per-request timing capture: raw timestamps + component durations.

Goal: let benchmarks see exactly what each phase of a request cost,
without the server pre-baking any "clean / dirty" interpretation into
the schema. The frontend composes any subtracted / steady-state metric
itself from the raw signals below.

The data model is a single dataclass :class:`RequestTiming` (lazily
attached to a :class:`Request` at admission) plus a
list of :class:`StepRecord` entries — one per engine step that touched
the request. Computed values exposed by ``to_payload()`` are RAW
component durations and timestamps only; per-label first-call costs
are emitted as discrete labeled items so the consumer can sum / filter
to taste.

First-call costs are recorded via the :func:`mark_first_call` context
manager: any engine block known to potentially be cold-only on its
first call wraps the block, the manager measures wall delta on entry /
exit and attributes it to the step's ``first_call_ms`` under a label.
The labels are stable strings (``"cudagraph_capture"``, ``"jit_init"``,
``"autotune"``, ``"metadata_first_build"`` …) so OTEL emits them as
``arbi.first_call.label`` attributes.

Capture has two tiers, split by what each costs on the hot path:

  - **Lifecycle stamps** (``detailed=False``) — the once-only
    admit / first-scheduled / first-token / complete timestamps plus a
    per-token ``stamp_last_token``. A handful of ``time.time()`` calls
    and one small object per request; nothing accumulates with output
    length. This tier feeds the request-timeline ring behind
    ``GET /v1/admin/request_timeline`` and the per-request latency
    histograms, so it is attached to every request.
  - **Per-step records** (``detailed=True``) — one :class:`StepRecord`
    per engine step that touched the request, plus first-call cost
    attribution. This is what ``usage.timing`` reports, and it is the
    tier that ALLOCATES per step and RETAINS per output token, so it is
    opt-in: ``cfg.timing_debug`` or the request's ``return_timing``.

:meth:`RequestTiming.begin_step` returns ``None`` on the lifecycle tier,
which is the single place the two are told apart.

Consumer guidance — what each commonly-asked metric is in this schema:
  - "TTFT" = ``components.first_token_s`` (raw — includes any first-
    call cost the request happened to absorb).
  - "TTFT minus first-call costs" = ``first_token_s -
    sum(first_call_costs_ms.values()) / 1000`` (frontend computes; the
    server emits raw signals only).
  - "TPOT" = ``decode_window_s / max(1, n_output_tokens - 1)``.
  - "Steady-state TPOT" = the same formula on requests where
    ``first_call_costs_ms`` is empty.

One engine ``_step()`` invocation that touched the owning request.

Times are wall-clock seconds (Unix epoch via :func:`time.time`).
``first_call_ms`` is an additive accumulator: a single step may pay
several distinct first-call events (cudagraph capture *and* a JIT
first-call), and the breakdown lives in ``first_call_costs_ms``.

Full per-request timeline + raw component durations.

Stamped at canonical points in the engine hot path (admission,
grammar attach, first scheduled step, first / last token emitted,
finish). The ``steps`` list grows as the request progresses; on
finish, the API layer calls :meth:`to_payload` to materialise the
summary block returned to the client.

All timestamps are wall-clock seconds (Unix epoch). 0.0 means
"not yet stamped" — the ``to_payload()`` serialisation maps that to
JSON ``null`` so consumers can distinguish unstamped from
stamped-at-the-unix-epoch.

The Python attribute names keep the ``t_`` prefix (``t_admit``,
``t_first_token_emitted``, etc.) for back-compat with internal call
sites; the JSON dict keys emitted by ``to_payload()`` drop the
prefix because the parent ``timestamps`` block already conveys
"these are timestamps".

Attribute the wall-delta of a block as a first-call cost on ``step_record``.

Used at engine call sites that may pay one-shot init costs on first
call (cudagraph capture for an unseen shape, first metadata-builder
invocation for a backend, first kernel-module load).

Both ``timing`` and ``step_record`` are accepted as None — the
common (warm) path through the lifecycle tier, where neither
object exists, must be a single attribute compare with no
measurement overhead.

Implementation note: on the lifecycle tier the body still
runs (we MUST execute the underlying init), but we skip the
:func:`time.time` calls entirely — the None-check is a single
cheap attribute compare.

Factory used at request admission — every request gets one.

``detailed`` selects the per-step tier (``cfg.timing_debug`` or the
request's ``return_timing`` opt-in); without it the request carries
lifecycle stamps only. Centralised so the tier decision lives in
exactly one place.

Open a new :class:`StepRecord` at the start of ``_step``.

``None`` on the lifecycle tier (``detailed=False``) — the caller
then threads no record through the step body, which is what keeps
the always-on tier free of per-step allocation. Otherwise the
engine threads the returned record through the step body and
stamps forward / sample / commit on it.

Serialize for the ``usage.timing`` response extension.

Schema (intentionally raw — frontend composes derived metrics):

::

  {
    "timestamps": {
      "admit": <unix_seconds_float | null>,
      "grammar_attached": <float | null>,
      "first_scheduled_step": <float | null>,
      "first_token_emitted": <float | null>,
      "last_token_emitted": <float | null>,
      "complete": <float | null>,
    },
    "components": {
      "queue_wait_s": <float>,                  # first_scheduled_step - admit
      "grammar_compile_s": <float | null>,      # grammar_attached - admit
      "first_token_s": <float>,                 # first_token_emitted - admit (TTFT)
      "decode_window_s": <float>,               # last_token_emitted - first_token_emitted
      "total_s": <float>,                       # complete - admit
      "prefill_compute_s": <float>,             # sum of prefill step durations
      "decode_compute_s": <float>,              # sum of decode step durations
      "n_steps": <int>,
      "n_output_tokens": <int>,
    },
    "first_call_costs_ms": {                    # OMITTED when empty
      "<label>": <ms_float>, ...
    }
  }

Timestamp dict keys drop the ``t_`` prefix (callers know they're
timestamps because they're under ``timestamps``). Unstamped
values serialize as ``null`` so consumers can distinguish "not
stamped" from "stamped at the unix epoch".

Single source of truth for all ``ARBI_*`` / ``TKV_*`` runtime flags.

This module is the enforced registry for every arbi-serve runtime config
env var. The contract (policed by
``tests/test_runtime_flags_single_source.py``):

  * Every ``ARBI_*`` / ``TKV_*`` env knob arbi-serve reads is a field on
    :class:`RuntimeFlags`, with a typed default, allowed values where it
    is an enum, and a one-line what-it-does/impact doc.
  * Call sites read ``runtime_flags().<field>`` — never
    ``os.environ.get("ARBI_…")`` directly. The guard test fails the
    build if a raw read sneaks back in outside this module.

Why:

  1. **Discoverability.** ``python -m arbi_serve.runtime_flags --doc``
     prints every flag (name / type / default / allowed values / impact),
     generated from the dataclass — no hand-maintained copy to drift.
  2. **Type safety.** Each flag has a typed parser (``bool``, ``int``,
     ``str``, ``tuple[int, ...]``, choice). Call sites don't reimplement
     ``"1" in {"1","true",…}`` boilerplate.
  3. **Behaviour preservation.** Each field's default matches the
     effective default already in effect at the call site it replaces.

Usage:

    >>> from arbi_serve.runtime_flags import runtime_flags
    >>> if runtime_flags().compile_on:
    ...     ...

``runtime_flags()`` re-reads the environment on every call so the
standard pytest ``monkeypatch.setenv`` pattern keeps working. Callers
that want a boot-time snapshot do ``flags = runtime_flags()`` and pass
the frozen dataclass through.

Adding a flag:

  1. Add a field to :class:`RuntimeFlags` with a typed default and a
     ``#:`` doc comment (env-var name + what-it-does/impact).
     :meth:`RuntimeFlags.from_env` derives the parser from the field's TYPE
     automatically (``bool`` / ``int`` / ``str`` / ``str | None`` /
     ``tuple[int, ...]`` / ``tuple[int, ...] | None``), so there is no
     per-field parser line to add or keep in sync.
  2. For a choice/enum field, register its allowed values in
     :data:`_ALLOWED_VALUES`. For a field whose parsing does not fit the
     type-driven scheme, add a ``raw -> value`` entry to
     :data:`_SPECIAL_PARSERS`.
  3. Register the env-var name in :data:`_ENV_NAMES` if it doesn't follow
     the ``ARBI_<FIELD>`` convention.
  4. At the call site, read ``runtime_flags().<field>``.

Frozen snapshot of every ``ARBI_*`` / ``TKV_*`` runtime flag that
arbi-serve reads.

Each field's default matches the effective default already in effect at
the call site it replaces — switching to the registry is a behaviour
no-op. ``#:`` comments document the env-var name, what the flag does,
and its perf/behaviour impact.

Refuse an out-of-range ``async_output_depth`` (fail-loud, both entry points).

Returns ``depth`` unchanged when it is supported
(``0 <= depth <= MAX_SUPPORTED_ASYNC_OUTPUT_DEPTH``, currently 1); raises ``ValueError``
naming the setting and the supported range otherwise. Called at engine boot
(``engine.__init__``) AND on the live ``/v1/admin/config_override`` path (via
the ``async_output_depth`` param parser), so an unsupported depth can be
entered through neither door.

Resolve the split-mixed cost gate's threshold, honouring provenance.

Returns the decode-row floor the per-step gate should apply, or
:data:`SPLIT_MIXED_MIN_DECODE_ROWS_UNGATED` (0) for "do not gate".

The whole subtlety is what ``ARBI_SPLIT_MIXED_DECODE_PREFILL=1`` means
once the flag ships default-ON *and* a cost gate exists behind it. The
shipped default ON means "the split is ELIGIBLE; let the cost model
pick per step". An operator typing that same ``1`` means something
stronger — they are running an A/B arm, and an arm that silently
declines to split at the size they are probing is a bogus arm, not a
clever engine. So an EXPLICIT ``split_mixed_decode_prefill`` (env var
present, or a live ``/v1/admin/config_override`` in the active
overlay) suppresses the gate.

Explicitness is read with :func:`flag_explicitly_set` — env/overlay
PRESENCE, never ``value != default``. A value-vs-default test would
answer "did the operator choose this?" wrongly the moment any recipe
or preset writes the field's own default back into an operator-facing
slot: the write leaves no trace, the comparison sees the default, and
every explicitness check downstream silently flips to "no".

The threshold knob itself wins when BOTH are set: naming a number is
strictly more specific than naming a direction, so
``ARBI_SPLIT_MIXED_DECODE_PREFILL=1
ARBI_SPLIT_MIXED_MIN_DECODE_ROWS=8`` gates at 8 rows rather than
ignoring the number the operator just typed.

``split_mixed_decode_prefill=False`` never reaches here — the caller
short-circuits on it, so an explicit OFF stays an absolute OFF.

Rows at which a routed-expert layer switches to the expert-major kernel.

The crossover is a property of the topology, not of the engine: at
``expert_shard_count == 1`` every expert is resident and the filter is
inert, so the two legs trade at a different row count than they do under
sharding. One shipped number cannot be right for both.

An explicit ``ARBI_MOE_EXPERT_MAJOR_MIN_ROWS`` wins at every topology —
naming a number is more specific than any default. Explicitness is
:func:`flag_explicitly_set` (env/overlay PRESENCE), never
``value != default``: a recipe that writes the field's own default back
into the slot would otherwise erase the operator's intent.

Install the active member's RuntimeFlags overlay (live config override).

Pass the variant's ``{field: value}`` deltas; ``None``/empty clears it
(back to the env-derived snapshot). Invalidates the cache so the next
:func:`runtime_flags` reflects the new overlay immediately. Called by the
stable-VA controller on switch + around a member's build.

Temporarily install ``overlay`` as the active-member flag overlay.

Used around a member's ``build()`` so the snapshotted-at-build runtime
flags (split_mixed_decode_prefill / async_output_depth / … ) are read at
the member's values during ITS capture. Restores the prior overlay on exit.

Return the process-local :class:`RuntimeFlags` snapshot.

The snapshot is built lazily from ``os.environ`` on first call and
cached — hot-path callers (one per kernel launch) pay a single dict
lookup, not a full ~100-field re-parse of the environment.

The cache is invalidated by :func:`reset_runtime_flags_cache`. Any
in-code env mutation that must be visible to subsequent
``runtime_flags()`` reads (e.g. :func:`force_compile_off`) calls that
hook. Production env is fixed at boot, so the cache never goes stale
there. Under tests the cache is DISABLED wholesale (the autouse
fixture wraps every test in :func:`no_runtime_flags_cache`), so a test
that ``monkeypatch.setenv``s or pokes ``os.environ`` mid-body always
sees its change with no manual reset — the stale-snapshot footgun
can't be re-introduced.

When an ACTIVE-MEMBER overlay is installed (live config override), it is
applied on top of the env-derived snapshot so the returned flags reflect
the currently-active prepared config variant. No overlay ⇒ byte-identical
to the historical env-only path.

Disable the :func:`runtime_flags` cache so every call re-reads env.

Test-only seam. While active, ``runtime_flags()`` re-parses
``os.environ`` on every call, so a mid-test env mutation
(``monkeypatch.setenv`` or a raw ``os.environ[...] = ...``) is visible
immediately without a manual :func:`reset_runtime_flags_cache`. The
autouse fixture in ``tests/conftest.py`` wraps the whole suite in this,
making the stale-snapshot footgun structurally impossible in tests.
Production never enters this block, keeping the hot-path cache.

Compiled dataclass default for a :class:`RuntimeFlags` field.

Raises ``KeyError`` on an unknown field so a typo can't silently read
as some default. Used by flag-truth value-knob counters to fire only
on an override-from-default (a value equal to the compiled default is
inert — firing on it would certify a knob that changed nothing).

True iff the operator set ``field_name`` — env var or live override.

Distinguishes "the operator asked for this value" (the env var is present
in the environment, or a live ``/v1/admin/config_override`` installed it in
the active-member overlay) from "this is just the shipped dataclass
default".

This is what a default-on flag needs to log honestly. When a flag ships Keying the log level on this instead
keeps the loud line pointed at the case that is actually suspicious — a
deliberate request that silently did nothing (a bogus A/B arm) — while an
ordinary boot that legitimately cannot take the path stays quiet.

Raises on an unknown field so a typo can't silently read as "not set".

Parse the raw env value for one :class:`RuntimeFlags` field.

Selects the parser from the field itself, in precedence order:
a :data:`_SPECIAL_PARSERS` entry, then a :data:`_ALLOWED_VALUES` choice
set, then the field's type via :data:`_TYPE_PARSERS`. Raises ``TypeError``
(loud, at build time) for a field whose type has no registered parser, so a
newly added field can never silently skip reading its env var.

Temporarily set ``ARBI_COMPILE_OFF=1`` for the duration of the block.

Used by the activation-profile / model-runner profile forwards, which
must run eager regardless of the operator's compile setting. Keeps the
env write+restore in the registry so call sites don't poke os.environ
directly. Downstream code reads ``runtime_flags().compile_off``.

Extract the ``#:`` doc comment for ``field_name`` from its defining class.

The dataclass uses ``#:`` comment lines preceding each field as the
what-it-does/impact doc. We parse the source rather than duplicate the
text — single source.

The fields live on the frozen-dataclass base classes of
:class:`RuntimeFlags` (``_RuntimeFlagsCore`` / ``_RuntimeFlagsBackend`` /
``_RuntimeFlagsSpec``, split out only to keep each source file under the
OSS <=1000-line rule), so we walk the MRO and read whichever base actually
defines the field. ``inspect.getsource`` on just ``RuntimeFlags`` would see
an empty class body and find nothing.

Render the full flag reference as Markdown, from the dataclass.

For every field: env-var name, type, default, allowed values (if any),
and the what-it-does/impact doc. Rendered ON DEMAND via
``python -m arbi_serve.runtime_flags --doc``; nothing is committed, so
the dataclass is the only source and cannot drift from a checked-in copy.

Resolve the tri-state ``serve_enable_traces``.

Unset (``None``) follows ``serve_enable_otel`` — per-request
tracing rides the same OTLP push the metrics/logs export uses,
so turning OTEL on turns traces on by default. An explicit
``ARBI_SERVE_ENABLE_TRACES=0`` opts out of spans while keeping
the rest of the OTEL channel.

Build a :class:`RuntimeFlags` from the current ``os.environ``.

Iterates the dataclass fields and applies each field's parser, which
is selected from the field itself (its type annotation, its choice
registration in :data:`_ALLOWED_VALUES`, or a non-uniform entry in
:data:`_SPECIAL_PARSERS`) by :func:`_parse_env_field`. One field ⇒ one
parser, derived from the field, so the field list and the parsing logic
cannot drift the way a parallel hand-written parser block could.

Behaviour is byte-identical to the per-field form it replaces: same
env-var name (:func:`env_name_for`), same default (:data:`_FIELD_DEFAULTS`,
the field literal), same coercion/truthiness/unset-vs-empty handling,
and the same loud ``ValueError`` on a malformed bool/choice/int-tuple.

Backend/kernel runtime flags (backend, fp8/awq, sleep/wake, cache, observability, debug, profiling) — a :class:`~arbi_serve.runtime_flags.RuntimeFlags` field group.

This module defines one frozen-dataclass base class carrying a subset of the
``RuntimeFlags`` fields verbatim (names, types, defaults, and their ``#:``
env-doc comments), keeping this source file under the OSS <=1000-line rule.
The concrete ``RuntimeFlags`` inherits every such base, so
``dataclasses.fields(RuntimeFlags)`` yields all fields in their original order
(MRO reverse gives core -> backend -> spec), and ``_doc_for_field`` walks the
MRO to read each field's ``#:`` comment from the base that defines it.

``from __future__ import annotations`` is required here: it keeps each field's
``.type`` in string form (``"bool"`` / ``"str | None"`` / ``"tuple[int, ...]"``),
which is exactly what the type-driven ``from_env`` parser dispatch keys on.

Core runtime flags (compile/inductor, boot, capture, decode, scheduler, TP, memory) — a :class:`~arbi_serve.runtime_flags.RuntimeFlags` field group.

This module defines one frozen-dataclass base class carrying a subset of the
``RuntimeFlags`` fields verbatim (names, types, defaults, and their ``#:``
env-doc comments), keeping this source file under the OSS <=1000-line rule.
The concrete ``RuntimeFlags`` inherits every such base, so
``dataclasses.fields(RuntimeFlags)`` yields all fields in their original order
(MRO reverse gives core -> backend -> spec), and ``_doc_for_field`` walks the
MRO to read each field's ``#:`` comment from the base that defines it.

``from __future__ import annotations`` is required here: it keeps each field's
``.type`` in string form (``"bool"`` / ``"str | None"`` / ``"tuple[int, ...]"``),
which is exactly what the type-driven ``from_env`` parser dispatch keys on.

Speculative-decode + TKV runtime flags (MTP, TKV codec-core) — a :class:`~arbi_serve.runtime_flags.RuntimeFlags` field group.

This module defines one frozen-dataclass base class carrying a subset of the
``RuntimeFlags`` fields verbatim (names, types, defaults, and their ``#:``
env-doc comments), keeping this source file under the OSS <=1000-line rule.
The concrete ``RuntimeFlags`` inherits every such base, so
``dataclasses.fields(RuntimeFlags)`` yields all fields in their original order
(MRO reverse gives core -> backend -> spec), and ``_doc_for_field`` walks the
MRO to read each field's ``#:`` comment from the base that defines it.

Each ``#:`` block is the operator-facing description of the flag: what turning
it on does, then only what changes a decision — which path it affects, the
trade-off, when to want it, and any hard prerequisite. Where a number matters,
name the counter or metric that reports it rather than freezing a measurement
here; ``tests/test_flag_doc_prose.py`` gates that.

``from __future__ import annotations`` is required here: it keeps each field's
``.type`` in string form (``"bool"`` / ``"str | None"`` / ``"tuple[int, ...]"``),
which is exactly what the type-driven ``from_env`` parser dispatch keys on.

Typed env-var parsers and choice/allowed-value sets for RuntimeFlags.

These are pure functions/constants with no dependency on ``RuntimeFlags``
itself, kept in their own module to keep :mod:`arbi_serve.runtime_flags`
under the OSS <=1000-line rule; ``runtime_flags`` re-imports them so
``arbi_serve.runtime_flags.<name>`` stays importable.

Preserve the unset-vs-empty distinction.

Unset → ``default``; a literal empty string stays ``""`` (some call
sites use ``ARBI_FOO=`` to override an earlier export back to empty,
which is distinct from "never set"). Used by ``ARBI_MTP`` and the
grammar / budget cache dirs.

Parse a comma-separated int tuple with an explicit UNSET sentinel.

Unlike :func:`_parse_int_tuple`, an unset/blank env var returns ``None``
(preserve the config default) rather than a baked default tuple, while
``off``/``none``/``0`` returns ``()`` (an explicit empty set — e.g. a single
full-width KV-page bucket). A CSV list returns a sorted-dedup positive-int
tuple. Used by knobs whose config default lives on the dataclass, where the
env var must distinguish "leave the default" from "explicitly empty".

The row counts an explicit ``ARBI_EXL3_PIN_SELECT_ROWS`` names.

Sorted and deduplicated: the value is a SET of rows the kernel families are
raced over, and it lands in the pin-cache and shipped-table match keys, so
two spellings of one set must not key two entries.

REFUSES rather than defaults. An empty list leaves the race nothing to
time, and a row below one is not a shape any linear presents; falling back
to the shipped set on either would race a family over rows the operator did
not ask for and then record it under the key they did.

Canonicalise ``ARBI_EXL3_PIN_SELECT_ROWS``; a value it cannot read raises.

Three spellings: ``""`` (the shipped row set), ``auto`` (derived from the
served config at the pin seam) and a comma list. Validated HERE, at parse
time, because the flag decides which kernel every EXL3 linear runs and the
pin is frozen long before the first request — a value silently read as
"shipped" would serve numerics the operator did not ask for with no way to
tell from the outside.

Parse the tri-state ``ARBI_TRUE_STOCHASTIC_DRAFT`` mode.

Unset / empty → ``default`` (the dataclass literal, currently ``auto`` —
per-row temperature routing). ``default`` is threaded in rather than
written here twice: this parser is what every env-derived boot goes
through, so a hardcoded fallback would silently outrank the field literal
and a change to the shipped default would be a no-op on every real server
while still passing a dataclass-level unit test.

EMPTY is treated as UNSET, not as "off". Compose passes an unset knob
through as an empty string (``ARBI_TRUE_STOCHASTIC_DRAFT=``), so mapping
empty to ``0`` would disable the shipped default on exactly the
deployments that never opted out of it. Explicit ``false`` / ``off`` /
``no`` still mean ``0``.

Legacy boolean spellings map to ``1``/``0``; anything else must be one of
:data:`_TRUE_STOCHASTIC_DRAFT_MODES` (fail-fast on misconfigured boots,
matching :func:`_parse_choice`).

Parse the tri-state ``ARBI_EXL3_INT8_VERIFY`` mode.

Same shape as :func:`_parse_true_stochastic_draft`, and for the same three
reasons: ``default`` is threaded in so the field literal stays the single
source, EMPTY is UNSET rather than off (compose passes an unset knob
through as ``ARBI_EXL3_INT8_VERIFY=``, and mapping that to ``0`` would
disable the shipped default on exactly the deployments that never opted
out), and an unrecognised value fails the boot loudly rather than
resolving to some leg nobody asked for.

Parse ``ARBI_EXL3_INT8_DECODE_MIN_ROWS``: ``auto`` or a positive int.

Same three rules as :func:`_parse_exl3_int8_verify`: the default is
threaded in so the field literal stays the single source, EMPTY is UNSET,
and anything else fails the boot loudly -- a row threshold that parsed
``0`` or a typo as "serve every batch" would route decode onto a leg
nobody asked for.

Parse ``ARBI_RECURRENT_PREFILL_CHUNK``: empty = OFF, else a positive int.

Refuses the boot on any other value. The flag is a scheduler cap: a
value that silently coerced to OFF would run a different config than
the operator asked for while the env still displays the asked-for
value — the scheduler consumer (``scheduler.schedule``) does a plain
``int()`` on the non-empty string and relies on this parse having
already rejected garbage.

Parse ``ARBI_BOOT_HEARTBEAT_S``: ``0`` = silent, else a cadence.

Deliberately LENIENT where the other float flags are loud. This one
only sets the cadence of a progress line, so a malformed, negative,
or NaN value falls back to the default cadence instead of refusing
the boot — a knob that governs logging must never be able to fail
the thing it is reporting on (see
:mod:`arbi_serve.engine.boot_heartbeat`). ``0`` is an explicit
opt-out (silent), matching ``ARBI_JIT_LOCK_PROGRESS_S``.

An interval is CLAMPED to ``threading.TIMEOUT_MAX``. ``inf`` — and
any absurdly large finite value — otherwise reaches ``Event.wait``
and raises ``OverflowError: timestamp out of range for platform
time_t`` on the heartbeat thread, which is both silent (no progress
line ever) and noisy (an unhandled thread exception at boot). The
clamp preserves the intent ("effectively never") without the crash.

Return ``tkv.config.TKV_FLAGS[field].default``.

Codec-core ``TKV_*`` flags (calibration-file / kernel knobs) keep
their canonical default in tkv's registry so the value can't drift
between tkv (vllm/sglang) and arbi-serve. ``from_env`` reads the
default from there instead of redeclaring a literal.

Imported lazily — tkv is heavy and importing it at module load risks
circular imports. ``runtime_flags()`` re-reads per call; the cost is
one dict lookup after tkv is first imported.

Return ``rows`` if it is a usable split-mixed cost threshold.

Accepts ``0`` (:data:`SPLIT_MIXED_MIN_DECODE_ROWS_UNGATED`, no gate)
through :data:`SPLIT_MIXED_MIN_DECODE_ROWS_CEILING`. Raises
``ValueError`` naming the setting and the accepted range otherwise.

Called at engine boot AND on the live ``/v1/admin/config_override``
path, so an unusable threshold can be entered through neither door.
The loud refusal is the point: a threshold silently clamped into range
would keep serving while the operator believed a different gate was in
force.

Sampling utilities operating on ``(K, num_seqs, V)`` logits.

The :class:`Sampler` + :class:`LogitsProcessor` chain is the single
production sampling path. ``XGrammarLogitsProcessor`` is exposed via
lazy attribute access — its underlying ``import xgrammar`` pulls
``transformers``, and structured-generation code paths are only
touched once the engine encounters its first
``response_format``-bearing request.

Triton ``@triton.jit`` kernels for the counter-based Gumbel-max sampler.

These kernels back the public wrappers in
:mod:`arbi_serve.sampler.gumbel_argmax_triton`; they live in a sibling
module purely to keep each first-party file under the size cap. Import
the public API from ``gumbel_argmax_triton`` — not from here.

See that module's docstring for the sampler's design (single-pass
counter-based Gumbel-max, Philox4x32-10 on-the-fly noise, cross-rank /
SPMD determinism).

One Philox4x32-10 draw → a fp32 uniform in the open interval (0, 1).

``c0``/``c1``/``c2`` are uint32 tensors seeding counter lanes 0-2
(lane 3 is zero); ``key_lo/key_hi`` are the Philox key (a 64-bit seed
split into two uint32). Ten Feistel-style rounds with per-round key
bumps — the standard Philox4x32-10 used by cuRAND / torch. Returns
lane-0 of the output mapped to (0, 1). All arithmetic is uint32
(wraps), so the result is bit-identical on every device and every
rank. Distinct ``(key, c0, c1, c2)`` tuples yield independent draws —
callers spread (vocab index, row, draw-class) over the counter lanes
to get decorrelated streams from one seed.

Philox4x32-10 uniform keyed on ``(row_seed, vocab_index)`` only.

``counter`` seeds counter lane 0 (the vocab index); lanes 1-2 are
zero. The per-row decode sampler's draw — see
:func:`_philox4x32_lane0` for the full-counter variant.

Partial single-pass Gumbel-max argmax over a contiguous vocab chunk.

Grid is ``(BATCH_SIZE * SPLIT,)``. Program ``pid`` owns row
``pid // SPLIT`` and the contiguous vocab chunk ``pid % SPLIT`` (so
a B=1 decode fans a single row across up to ``SPLIT`` SMs — the whole
point, since one SM cannot saturate HBM over a 248k-wide row). Each
program writes its chunk-best ``(score, idx)`` to ``PART_*`` and a
trivial host-side ``argmax`` over the ``SPLIT`` partials finishes the
reduction. Contiguous (not strided) chunks keep the per-program tiles
coalesced AND make lower vocab indices land in lower ``split_id``s, so
the finishing ``argmax`` inherits a lowest-index tie-break.

Partial single-pass Gumbel-max argmax keyed on a DEVICE step seed.

Like :func:`_gumbel_partial_kernel` but the Philox key is the shared
step seed read INSIDE the kernel (``tl.load(SEED_PTR)`` — no host
readback, capture-replay reads whatever ``copy_()`` last wrote) and
the per-row / per-draw-class decorrelation rides the counter lanes:
``(vocab_index, row, DRAW_OFFSET)`` — the same namespacing
:func:`_residual_gumbel_partial_kernel` uses (verify recovery/bonus
own ``DRAW_OFFSET = OFFSET_RECOVERY``; the drafter chain passes
``OFFSET_DRAFTER + 17 * step``). ``-inf`` score entries (masked-out /
zero-probability tokens) stay ``-inf`` after the Gumbel add and can
never win, so the drawn support is exactly the finite-score set.

Grid, chunking, and tie-break follow :func:`_gumbel_partial_kernel`:
``(N_ROWS * SPLIT,)`` programs, contiguous chunks, strictly-greater
running max ⇒ lowest-index tie-break.

Partial single-pass Gumbel-max over the MTP verify residual / bonus rows.

Flat row layout over ``(K+1, B)``: row ``r = k * B + b``. Rows
``r < KB`` (``k < K``) are RESIDUAL rows — the Gumbel-max runs on
``log((p - q)+)``, an exact categorical draw from the normalised
Leviathan-2023 recovery distribution ``(p_target - q_draft)+ / Z``
(the ``1/Z`` is a per-row constant that drops out of the argmax).
Rows ``r >= KB`` (``k == K``) are BONUS rows — ``q`` is not read
(residual = ``p``), giving an exact draw from ``p_target[K]``.
Zero-residual entries score ``-inf`` and can never win, so the
drawn support is exactly the positive-residual set.

``Q_IS_ONEHOT`` (constexpr — the dense specialization compiles the
branch out entirely, so the dense path's generated code is unchanged)
selects the INDEX-FORM point-mass proposal: instead of loading a
materialized one-hot ``Q`` row, ``q`` is synthesized per lane as
``1.0 where offs == DRAFT[k, b] else 0.0`` — exactly the values the
dense one-hot load would produce, so ``res = max(p - q, 0)`` is
bit-identical while the ``(K, B, V)`` fp32 q tensor is never built
or read (the 27B-c8 copy/index lever). ``DRAFT`` is only
dereferenced on residual rows in one-hot mode.

``IN_LOG_SPACE`` (constexpr — the dense-probs specialization compiles
the branch out, generated code unchanged) selects the FUSED softmax:
``P`` then holds temperature-scaled (optionally masked) LOGITS and the
target prob is computed per lane as ``p = exp(logit - LSE[row])`` using
the precomputed row ``logsumexp`` — so no dense ``(K+1, B, V)`` probs
tensor is materialized and no separate full-vocab ``softmax`` launch
runs before this kernel. A masked-out logit is ``-inf`` (top_k/top_p),
giving ``p = 0`` exactly, and an out-of-range lane loads ``-inf`` too,
so the residual support and the degenerate-fallback ``argmax(p)`` are
identical to the probs path (``exp`` is monotone). Bit-for-bit the
softmax differs from a materialized ``torch.softmax`` only by float32
reduction order, which the residual sampler is lossless to
(Leviathan-2023 Thm 1) and which cannot move the argmax draw.

The Gumbel noise is generated ON THE FLY from Philox4x32-10 with the
step seed as the key and ``(vocab_index, row, DRAW_OFFSET)`` on the
counter lanes — no materialized ``(K+1, B, V)`` noise tensor, and a
pure function of ``(seed, k, b, v)`` so every SPMD rank draws the
byte-identical token. ``DRAW_OFFSET`` namespaces this draw class away
from other Philox consumers keyed on the same seed (the decode
sampler uses zero counter lanes 1-2).

Alongside the Gumbel partial, each program tracks its chunk's plain
``argmax(p)`` (``PART_PVAL`` / ``PART_PIDX``) — the finishing
reduction's fallback for a degenerate row whose residual is
everywhere zero (draft ⊇ target support), matching the eager
sampler's degenerate contract.

Grid, chunking, and tie-break follow :func:`_gumbel_partial_kernel`:
``(N_ROWS * SPLIT,)`` programs, contiguous chunks, strictly-greater
running max ⇒ lowest-index tie-break.

Gumbel(0,1) noise AT arbitrary global vocab indices, one row/program.

For candidate ``(row, c)`` with global vocab id ``v = IDX[row, c]``,
writes ``-log(-log(u))`` where ``u`` is the SAME Philox4x32-10 draw the
full-vocab recovery/bonus kernel (:func:`_residual_gumbel_partial_kernel`)
computes at that token: counter lanes ``(c0=v, c1=row, c2=DRAW_OFFSET)``,
key = the shared step seed. So the noise a compacted row sees at a
candidate is byte-identical to the noise the gathered full-vocab path
would have generated at the SAME token — the property that makes the
shard-resident recovery/bonus draw match the gathered draw.

``row`` is the FLAT verify-row index ``k * B + b`` (the caller flattens
``(K+1, B, C)`` → ``((K+1)*B, C)``), matching the full-vocab kernel's
``row = k*B + b`` counter lane. A negative (pad) index still produces a
well-defined noise value (uint32-cast wraps); pad candidates carry a
``-inf`` residual downstream and are never selected, so the value is
inert.

Per-chunk online-softmax partials: chunk max + chunk sum-exp
(relative to the chunk max). Grid ``(N_ROWS * SPLIT,)``, contiguous
chunks — the finishing host-side reduce combines them into the global
``(max, Z)`` per row.

Per-chunk kept-probability mass for min_p, probs-space (matching
``processors.min_p_mask``): keep ``probs ≥ min_p · max(probs)`` where
``max(probs) = div_rn(1, Z)`` (the global-max token is always kept by
top-k / top-p). Rows with ``min_p == 0`` write mass 0 (unused).

Write ``q`` and draw the token in ONE pass over the masked logits.

``q = div_rn(exp(x − max), Z)`` (min_p-renormalised per row when
enabled + ``min_p > 0``); the draw is a counter-based Philox4x32-10
Gumbel-max over ``log(q)`` keyed on ``(seed, row, DRAW_OFFSET)`` — the
IDENTICAL noise + arithmetic ``gumbel_argmax_from_seed_ptr`` uses — with
a strictly-greater running max (lowest-index tie-break). Contiguous
chunks make the host finishing ``argmax`` inherit the lowest-index
tie-break.

``PER_ROW_SEED`` selects the keyed-watermark noise instead: one seed per
row with counter lanes 1-2 zeroed, the layout
:func:`gumbel_argmax_from_seeds` and the model-free detector
(``philox_uniform_host``) recompute. Same ``q``, same arithmetic — only
the noise source differs, so the two modes share this one program rather
than one masked-softmax tail per draw class.

Fused single-pass Triton kernel for the counter-based rejection sampler.

Backs :func:`arbi_serve.sampler.rejection_sampler.rejection_sample_kernel`.
One program per row (grid ``(B,)``). Each program:

  1. computes the row normalizer ``(max, Z=sum exp(x-max))`` in ONE online
     pass — NO ``(B, V)`` CDF is ever materialized;
  2. runs the rejection loop IN-KERNEL: draw a candidate token from the
     SAME counter-based Philox stream keyed on ``(row_seed, draw_index)``
     the torch reference uses, locate it by inverse-CDF (one accumulation
     pass over the vocab), test ONLY that candidate against the survivor
     threshold ``tau`` (``logit >= tau`` — the top_k/top_p/min_p kept set
     reduces to a per-row logit threshold), accept or redraw up to the
     bounded budget; on exhaustion commit the deterministic in-support
     ``fallback``.

The whole draw drops Gumbel-max's per-element cost (a full Philox4x32-10
draw + a double ``log`` at EVERY vocab entry): here Philox fires only
ONCE per candidate (typically once total — the no-filter draw always
accepts round 0), and the per-element work is a single ``exp``.

Determinism anchor. The Philox uniform is the 24-bit
``((c0 >> 8) + 0.5) * 2^-24`` form (see
:func:`~arbi_serve.sampler.rejection_sampler.philox_uniforms`), computed
here in fp32 exactly as the torch reference computes it — so the fused
kernel draws the byte-identical uniform, and (gap ≫ fp32 rounding) the
byte-identical token. Rank-symmetry follows: every rank runs this kernel
on the rank-identical logits with the rank-agreed seed and threshold.

One Philox4x32-10 draw -> a 24-bit fp32 uniform in ``(0, 1)``.

Scalar counter ``(c0=draw_index, c1=0, c2=DRAW_OFFSET, c3=0)`` keyed on
``(key_lo, key_hi)``. Returns ``((c0 >> 8) + 0.5) * 2^-24`` in fp32 —
the SAME arithmetic
:func:`~arbi_serve.sampler.rejection_sampler.philox_uniforms` uses, so
the drawn uniform is bit-identical to the torch reference.

Device kernels for the fanned-out top_k/top_p/min_p threshold mask.

Two launches, both gridded over ``(rows x vocab-chunk)`` so a ``B = 1``
decode row occupies every SM. :mod:`arbi_serve.sampler.topc_mask` holds the
host-side dispatch and the rationale; nothing here is imported directly.

Per-chunk exact top-``CAP`` — the only pass that reads all of ``V``.

Grid is ``(N_ROWS * SPLIT,)``: program ``pid`` owns row ``pid // SPLIT``
and the contiguous vocab chunk ``pid % SPLIT``, so a single decode row
fans across ``SPLIT`` SMs instead of crawling one.

A chunk wider than one tile is folded by a running merge — the top-``CAP``
of ``(running top-CAP) || (this tile's top-CAP)`` is the top-``CAP`` of
everything seen, by the same union argument the two levels rest on — so
``BLOCK_SIZE`` bounds the register footprint independently of how wide
the vocabulary is. At the served width ``N_TILES == 1`` and the loop
compiles away entirely.

``WRITE_MASS`` adds the chunk's softmax mass in its own max's units, which
is what lets a nucleus be taken against the WHOLE row rather than against
the candidates: a ``top_p`` with no ``top_k`` renormalises over every
token, and no candidate set can supply that denominator. It rides this
pass because the denominator needs exactly the bytes the selection is
already reading — a separate mass pass would double the only full-``V``
read the design has.

Row keep-threshold ``T`` and row max, from the level-1 partials.

The union of the per-chunk top-``CAP`` sets contains the row's global
top-``CAP`` (a chunk can contribute at most ``CAP`` members to it), so
one more ``tl.topk`` over the ``SPLIT * CAP`` partials is the EXACT
global top-``CAP`` — no approximation anywhere in the selection.

The whole filter chain then collapses to a single scalar: top_k, top_p
and min_p each cut the row at a value, ``keep iff logit >= T`` with
``T`` the largest of the three cuts. Each cut is placed AT a real logit
value rather than strictly between two, which is what makes the kept
set match the pivot kernel's on an exact tie: both keep the whole tie
group at the cut.

Cheap enough (a ``CAP``-wide reduction over one ``PART_BLOCK`` tile)
that every masking program recomputes it rather than reading it from a
grid-``(N_ROWS,)`` kernel — which would reintroduce the single-CTA step
this whole design exists to remove.

``NUCLEUS`` covers the row shape a candidate set cannot bound on its own:
``top_p`` active with ``top_k`` DISABLED, where the nucleus renormalises
over the whole vocabulary and may run past the capacity. Two things make
it exact rather than approximate. The denominator comes from the per-chunk
masses, so it is the row's true softmax normaliser and not the candidates'
— and the row is SERVED only when the candidates provably contain the
boundary, with margin. ``handled = 0`` says they do not, and the caller's
escalation re-derives that row from the row itself.

The second return, ``handled``, is what keeps this honest under fp32. The
boundary is where a cumulative mass crosses ``p``; two implementations
summing the same terms in different orders disagree about that crossing
when the boundary token's own mass is comparable to the summation error.
``MARGIN`` is the relative slack the crossing must clear on BOTH sides —
below it, no arithmetic that differs by less than ``MARGIN`` can pick a
different boundary, so the served kept set is the same set the row-serial
search would have produced, and above it the row is not served at all.

Apply the row threshold to one vocab chunk; optionally sum for the LSE.

Same ``(N_ROWS * SPLIT,)`` grid as the partials pass. Only the DROPPED
lanes are written — kept lanes already hold their original fp32 value,
so a kept logit is bit-identical to its input by construction.

Under ``NUCLEUS`` the kernel also publishes ``OUT_P``: the ``top_p`` each
row still owes. A served row owes nothing and gets ``1.0`` — the DISABLED
encoding — so the escalation launch that follows walks straight past it,
and an unservable row gets its ``top_p`` back unchanged and untouched
logits to apply it to. Routing by a DEVICE value is what keeps the pair
capturable: the launch shapes never move, so the same recorded graph
serves whatever the next ``copy_()`` puts in the parameter buffers.

Fused Top-K + Top-P Triton ``@jit`` kernels — vendored from vLLM.

The device-side kernels for the combined top-k / top-p sampler live here;
:mod:`arbi_serve.sampler.topk_topp_triton` holds the host-side dispatch,
caches, and the public ``apply_top_k_top_p_triton`` /
``sample_top_k_top_p_triton`` wrappers, keeping each file small.
See that module's docstring for the full vendoring rationale, the
arbi-serve epilogues (``SAMPLE`` / ``MINP_ENABLED`` / ``WRITE_LSE``) and
provenance. Nothing here is re-exported directly; import the wrappers.

Exact top_p nucleus boundary, by bisection on the IEEE754 ordered key.

WHY THIS EXISTS. The vendored Qrita kernel finds its ``top_p`` cut with a
ternary search over the CONTINUOUS probability range, capped at 18 iterations
with a ``1e-9`` convergence floor, falling back to the interval MIDPOINT when
neither probe satisfies its acceptance test in budget. Ternary search narrows
by ~1/3 per step, so 18 steps from ``[0, 1]`` bottoms out at ~``1.4e-9`` —
exactly the floor. When two adjacent probabilities near the nucleus boundary
are closer than that, the search cannot separate them and the midpoint lands
on whichever side the interval happened to straddle. OR too large, both observed, on rows whose
boundary sits inside a plateau of near-identical logits — which is what a
real ``lm_head`` produces over a 248k vocabulary.

The fix is not more iterations. More iterations makes the search exact IN
PRACTICE, which is a claim about inputs rather than about the algorithm, and
it still carries a tolerance constant that has to be argued about.

WHAT THIS DOES INSTEAD. IEEE754 floats, reinterpreted through the standard
order-preserving transform, are a CONTIGUOUS RANGE OF INTEGERS. Bisecting
that integer range is bisecting the set of representable values, so it
terminates on an exact float in at most 32 steps with **no tolerance
constant anywhere** — the loop ends when ``lo == hi``, not when an interval
gets small enough.

The search runs in LOGIT space, never probability space, for the same reason:
the mask compares logits, so a boundary found among probabilities would have
to come back through ``log(p * Z) + max``, and that round-trip reintroduces a
ULP of slop at precisely the boundary being fixed. Bisecting logits directly
yields a threshold that IS a representable logit, and the returned pivot is
its exact predecessor, so ``logit > pivot`` reproduces ``logit >= threshold``
with no arithmetic in between.

COST, and the mistake worth recording. A first version bisected the key range
unconditionally and MEASURED 44% SLOWER on top_p-only rows. The pass-count
argument behind it — "24-32 cheap passes against ~36 expensive ones" — was
wrong about the baseline: the upstream ternary search is DATA-ADAPTIVE and
exits the moment its acceptance test can prove a pivot, which on ordinary
rows happens in a handful of iterations, nowhere near its 18-iteration cap.
Comparing worst case against worst case hid that the common case was cheap.

So the fix is
not to be unconditionally exact, it is to keep the early exit and make it
sound. The loop below carries the same proof-of-boundary test, in one pass
instead of the upstream's two, and only falls through to full bisection on
the rows where the old search would have given up and guessed.

Map fp32 to a uint32 whose integer order matches float order.

Non-negative floats keep their bit pattern with the sign bit set;
negative floats are complemented. ``^`` rather than ``~`` because Triton's
interpreter rejects the negative intermediate that ``~`` produces on an
unsigned type.

Pivot such that ``logit > pivot`` is exactly the ``top_p`` nucleus.

The nucleus is the smallest set of highest-probability tokens whose mass
reaches ``p``, with the whole tie group at the boundary included. Mass is
monotone in the threshold, so the boundary is the LARGEST logit ``L``
whose kept mass still reaches ``target``, and bisection on the ordered key
finds it exactly.

``target`` is passed rather than derived from the candidates because the
candidate set is often a SUBSET of the row (the outlier buffer). The
nucleus is defined against the full distribution's mass, so deriving it
locally would renormalise against the subset and cut in the wrong place.

Returns the predecessor of ``L``, so the caller's existing
``logit > pivot`` mask keeps ``logit >= L`` — every tie at ``L``
included, with no duplicate bookkeeping.

Shared host-side helpers for the sampler Triton kernels.

Small pure-Python utilities used by both :mod:`topk_topp_triton` and
:mod:`gumbel_argmax_triton` to size Triton launch grids. Kept in one
place so the two kernel files don't each carry a private copy.

Triton kernels for the watermark context cache's per-step device work.

The eager torch expression of :meth:`WatermarkContextCache.row_seeds` /
``advance`` (ring gather, SplitMix64 hash chain, Bloom check+insert,
seed select, ring shift) is ~20 tiny elementwise/gather/scatter launches
per decode step. Each is microseconds of device work but ~5-10 us of
launch overhead — measurable against a small model's ~2 ms step. These
kernels collapse the whole path to TWO launches (seed+Bloom, then the
post-sample ring advance), bit-identical to the torch path.

All 64-bit arithmetic runs in uint64 (native logical shifts, wrapping
mul/add) and matches the pure-Python reference in
:mod:`arbi_serve.sampler.watermark` — parity-covered by the same tests
that pin the torch path.

Two-probe Bloom membership on context hash ``h`` for cache ``row``.

Returns the pre-insert "seen" verdict; inserts the probes only when
``do_insert`` is true (predicated store — the check side is free of
writes so speculative contexts never poison the filter).

Watermark seeds for ``D_COUNT`` depths from ``D_START``. Grid (B,).

Depth ``d``'s context is the last ``H`` of (committed ring ++
``DRAFT[0:d]``) — the residual draw at the first rejected depth ``d``
sees exactly this context because the accepted prefix IS the drafted
prefix, and the bonus draw (depth ``K``) sees all ``K`` drafted
tokens. A depth's hash reads only tokens, never another depth's hash
or seed, so any contiguous depth range is computable on its own: the
verify tail takes ``D_START=0, D_COUNT=K+1``, a drafter's slot-``s``
seeds take ``D_START=s, D_COUNT=1``. Bloom is CHECK-ONLY here (a
masked depth falls back to the row's UNKEYED seed at that absolute
position: ``splitmix64(salt ^ splitmix64(pos + d + REMIX))``, the same
value the drafter that proposed the token derived); inserts happen at
commit, for the depths that were actually consumed.

``OUT`` is indexed by the depth's OFFSET in the range, while the
fallback keys on the ABSOLUTE depth — so a sliced range yields
exactly the rows a full sweep would have put at ``d * B + b``.

Commit an MTP verify step into the watermark state. Grid (B,).

Bloom-inserts the contexts actually consumed (depths ``0..n`` — the
accepted-drafted draws plus the recovery/bonus draw), then advances
the ring by the ``n+1`` committed tokens. Contexts for depths
``<= n`` are windows over (ring ++ committed prefix), which for the
accepted prefix is identical to the drafted prefix the seed kernel
hashed — so the inserted hashes match the seeds that were served.

Stop an in-flight request by forcing its stop token.

THE STOP IS ADDRESSED, NOT INFERRED FROM A CONNECTION. Cancelling by transport
teardown asks a chain of layers to notice that a socket died, and each link
answers late or not at all: a streaming generator learns of a disconnect only
when it is resumed (so never during prefill), Starlette's
``Request.is_disconnected()`` polls ``receive()`` inside an already-cancelled
scope, and closing an HTTP stream from a ``finally`` that is running *because*
of a ``CancelledError`` is itself cancelled before it reaches the socket.
Measured against this engine, a run whose natural length was 52.3 s and whose
client was killed at 10 s finished at 51.2 s: the cancel arrived after the work
it was cancelling.

So this does what the thinking budget already does for ``</think>`` — it forces
a token. The row is masked to a stop id on the next sampling step, the terminal
check sees a stop, and the request ENDS THE WAY IT WOULD HAVE ENDED ANYWAY:
a real ``finish_reason``, usage reported, KV released down the ordinary path,
the trace written. A torn-down connection gives none of that.

WHAT IT CANNOT DO, stated because the gap is structural rather than a
shortcoming to be fixed here: a logits processor runs only where there are
logits, and a mid-prefill chunk samples nothing (``step_variants``: "a prefill
row on its FINAL chunk … a mid-prefill chunk samples None"). A request that has
not yet produced its first token therefore cannot be stopped this way, and the
disconnect path in :mod:`arbi_serve.server.routes._request_ctx` is what covers
that window.

The id this row must emit this step, or ``None``.

Duck-typed on ``force_stop_now`` so an engine test double without the field
reads as "not stopping" rather than raising inside the sampler.

Mask verify position ``p == 0`` to the forced stop id.

A free function for the same reason the thinking budget's is: the mask can
then be driven against a slate directly in a test, without a sampler.

Under a vocab-parallel ``window`` every rank bans its whole slice — the ban
is the vocabulary minus one token — and only the rank that OWNS the forced
id reopens it, so the finite value lands exactly once across the group.

Logits processor forcing a stop token for rows an operator stopped.

Mirrors :class:`~arbi_serve.sampler.thinking_budget.ThinkingBudgetGuard`:
the row is masked to ``-inf`` and only the forced id reopened, at a FINITE
constant rather than its original logit — a model that assigned the stop id
``-inf`` would otherwise leave an all-``-inf`` row, where argmax and softmax
are both undefined.

One bool read per row, and the common case (nobody stopping) returns the
logits untouched, so an unused guard costs a comparison per row per step.

Ask one in-flight request to stop generating and finish.

Arms :attr:`~arbi_serve.engine.request.Request.force_stop_now`; the guard
forces a stop id on the next sampling step and the terminal check honours
it — including under ``ignore_eos``, which declares the MODEL's stop
uninteresting rather than declaring the request unstoppable.

ONE ENDPOINT FOR BOTH PHASES. The caller sends one request and the engine
chooses: a row that has produced a token is stopped by forcing one, a row
that has not is cancelled by id. ``how`` reports which happened, so the
outcome is legible without the caller ever having had to pick.

Refuses rather than no-ops, and says which:

  * the tokenizer publishes no stop id, so there is no token to force;
  * the request is not running (matched by the caller's own
    ``X-Request-ID`` first, then the engine's per-process counter);
  * it has not produced a token yet, so no sampling step exists to carry
    the forced id — the disconnect path covers that window, and reporting
    success here would promise a stop that lands only when prefill ends.

The slate form does not read the draft tokens.

:meth:`apply_to_verify_slate` constrains ``p == 0`` only and deletes
its ``drafts_per_row`` argument unread, so answering the base
protocol's default here would ask the verify path to pull every row's
drafts off the device — a blocking host copy on every decode step, for
a value this processor never looks at. This guard is attached whenever
the tokenizer has a stop id, which is every real boot, so the default
is the one answer it must not inherit.

Force the stop id at the IMMEDIATE verify position.

WITHOUT THIS THE STOP DOES NOT REACH A SPECULATING BOOT. The
shard-resident verify draw does not call ``__call__``; it asks each
processor whether it constrains the slate and then calls this. A guard
that answers True here and cannot write is worse than one that never
claimed to — and the whole shipped configuration on this box speculates,
so the plain path is the one that would almost never run.

Only ``p == 0`` is constrained: it is the position whose context is the
committed history, so it is where the stop must land. The row's drafts
at ``p > 0`` are then rejected against it and exactly the stop is
emitted.

Whether any column is being forced to stop this step.

Keyed on the same read the mask makes, so the verify slate and the mask
cannot disagree about which rows are constrained.

Single-pass counter-based Gumbel-max sampler (no-filter stochastic decode).

The default decode shape (``temperature > 0``, ``top_p = 1``, ``top_k``
unset, ``min_p = 0``) is a plain categorical draw from
``softmax(logits)``. The classic way to draw it is Gumbel-max:

    token = argmax_v( logit_v + g_v ),   g_v = -log(-log(u_v)),  u_v ~ U(0,1)

which is EXACTLY a categorical draw from ``softmax(logit)`` — no softmax
normalization, no sort, no top-k/top-p pivot search needed (the ``-max``
and ``/Z`` of softmax are constants that drop out of the argmax).

This module fuses that into ONE pass over the vocab:

  * the per-element Gumbel noise ``g_v`` is generated ON THE FLY from a
    **counter-based Philox4x32 RNG** keyed on ``(row_seed, vocab_index)``
    — so there is NO materialized ``(B, V)`` noise tensor and NO separate
    softmax pass;
  * a running argmax over the vocab tiles emits one token per row.

Cross-rank determinism (SPMD / TP). Because the noise is a pure function
of ``(row_seed, vocab_index)`` and ``row_seed`` is rank-agreed (derived
from the request's ``(seed, position)`` — the same source
``rank_symmetric.gumbel_exp_noise`` used), every TP rank generates
byte-identical noise and therefore selects the byte-identical token with
NO noise tensor exchanged between ranks. The single-process (c=1) path
uses the same kernel with a per-step seed drawn from the sampler-private
generator.

Distribution note. This is exact Gumbel-max = exact categorical from
``softmax(logits)``; it is the SAME distribution as the vendored
``_topk_topp_kernel`` Gumbel tail, not a bit-identical reproduction of a
specific ``torch`` noise draw. Temperature is folded into ``logits`` by
``apply_temperature`` upstream, so this kernel sees pre-scaled logits.

The ``@triton.jit`` kernels backing these wrappers live in the sibling
module :mod:`arbi_serve.sampler._gumbel_argmax_kernel` (split out purely
to keep each first-party file under the size cap); import the public API
from here.

Single-pass counter-based Gumbel-max sample → ``(B,)`` int64 tokens.

Draws one token per row as ``argmax_v(logit_v + Gumbel(row_seed, v))``
with the Gumbel noise generated on the fly (no materialized noise
tensor, no separate softmax). ``row_seeds`` is the per-row Philox key:
supply the rank-agreed ``(seed, position)``-derived seed for SPMD
cross-rank determinism, or a per-step private-generator draw for the
single-process path.

``logits`` is read-only (unlike the in-place top-k/top-p kernels).

Single-pass counter-Gumbel sample keyed on a device seed → ``(B,)`` int64.

Draws ``argmax_v(score_v + Gumbel(seed, row, v, draw_offset))`` per
row — an exact categorical draw from ``softmax(scores)`` (pass
per-row log-probs or masked logits; ``-inf`` entries are excluded
from the support). Unlike :func:`gumbel_argmax_from_seeds` the Philox
key is a single SHARED step seed read in-kernel from ``seed_dev``
(no per-row seed tensor, no host readback — capture-safe), with
``(row, draw_offset)`` on the counter lanes for decorrelation. The
noise is a pure function of ``(seed, row, v, draw_offset)``, so every
SPMD/TP rank draws the byte-identical token from the rank-agreed
seed with no exchange.

``draw_offset`` namespaces the draw class on counter lane 2 —
distinct consumers sharing a step seed pass distinct offsets (the
verify recovery/bonus kernel owns ``OFFSET_RECOVERY``; the drafter
chain passes ``OFFSET_DRAFTER + 17 * step`` so the K chained calls'
draws don't alias). Callers guarantee each row has at least one
finite score (an all ``-inf`` row returns index 0).

``scores`` is read-only.

Single-pass counter-based recovery + bonus draws → ``(K+1, B)`` int64.

Row ``k < K`` of the result is an exact categorical draw from the
normalised Leviathan-2023 residual
``(p_target[k] - draft_probs[k])+ / Z`` (falling back to
``argmax(p_target[k])`` when the residual is everywhere zero); row
``K`` is an exact draw from ``p_target[K]``. The Gumbel noise is
generated in-kernel from Philox keyed on ``(seed, row, vocab_index)``
— no materialized ``(K+1, B, V)`` noise tensor, no separate softmax
or full-vocab argmax pass, and byte-identical across ranks for the
same seed (SPMD determinism without a noise exchange).

``draft_probs=None`` selects the INDEX-FORM point-mass proposal:
``q[k, b, v] = 1`` iff ``v == draft_tokens[k, b]``, synthesized per
lane inside the kernel (constexpr-specialized; the dense
specialization's generated code is unchanged). Bit-identical residual
values to loading a densified one-hot; ``draft_tokens`` is required
then and ignored in dense mode.

``seed_dev`` is read INSIDE the kernel (``tl.load``), so no host
readback of a device seed and — once the surrounding eager tail is
captured — a graph replay reads whatever ``copy_()`` last wrote into
the buffer. ``draw_offset`` namespaces the draw class on a counter
lane; distinct consumers sharing a seed pass distinct offsets.

``in_log_space=True`` fuses the softmax: ``p_target`` then holds the
temperature-scaled (optionally top_k/top_p-masked) verify LOGITS and
``lse`` the precomputed ``(K+1, B)`` row ``logsumexp``, so the kernel
computes ``p = exp(logit - lse)`` per lane — no dense ``(K+1, B, V)``
probs tensor and no separate ``softmax`` launch. The drawn tokens are
distributionally identical (softmax reduction-order aside, which the
argmax draw is invariant to).

``p_target`` / ``draft_probs`` may be non-contiguous leading-dim
slices of a persistent scratch pool (unit stride required only on the
vocab dim). All operands are read-only.

Full-vocab-identical Gumbel noise sampled AT candidate indices.

Returns a ``(K+1, B, C)`` fp32 tensor whose ``[k, b, c]`` entry is the
Gumbel(0,1) noise the counter-based full-vocab recovery/bonus kernel
(:func:`residual_bonus_gumbel_from_seed`) generates at token
``cand_idx[k, b, c]`` for flat row ``k*B + b`` — so the shard-resident
recovery/bonus draw over the compacted candidates draws from the SAME
per-token noise realization as the gathered path, byte-for-byte, with
no materialized ``(K+1, B, V)`` noise tensor and no host readback.

``draw_offset`` namespaces the draw class on counter lane 2 (the
verify recovery/bonus tail passes ``OFFSET_RECOVERY``). ``seed_dev`` is
read in-kernel, so a captured verify tail replays against whatever
``copy_()`` last wrote.

Fused multi-SM drafter tail → ``(tokens (B,), q (B, V))``.

Fans the (typically B=1) decode row across every SM: per-chunk
online-softmax partials → global ``(max, Z)`` reduce → optional min_p
kept-mass reduce → one q-write + counter-Gumbel draw pass → finishing
argmax. Byte-matches the unfused ``torch.softmax`` [+ ``min_p_mask``] +
``gumbel_argmax_from_seed_ptr`` tail on the drawn TOKEN + kept SUPPORT
(q within a few fp32 ULPs; lossless for the residual sampler). All
reduces are capturable torch ops (no host readback); ``seed_dev`` is
read in-kernel, so the captured chain replays against whatever
``copy_()`` last wrote.

``row_seeds`` selects the KEYED-WATERMARK draw instead of the step seed:
one seed per row with zero counter lanes, the layout
:func:`gumbel_argmax_from_seeds` and the model-free detector recompute.
``q`` does not depend on which mode ran — the keyed draw is this tail
with a different Gumbel stream, not a different filtered distribution.

The mode is its OWN argument rather than a shape test on ``seed_dev``,
because ``B == 1`` is the common decode width: a ``(1,)`` keyed seed and
a ``(1,)`` step seed are indistinguishable by shape, and reading one as
the other would silently draw from the wrong noise stream on exactly the
slate the drafter runs most.

Boot-time pre-compile of the stochastic sampler's Triton kernels.

The captured decode path warms the GREEDY sampler at boot (argmax rides
the captured graphs), but the stochastic tail — the Qrita top-k/top-p
mask kernel and the counter-based Gumbel-max kernel — compiles lazily on
the FIRST ``temperature > 0`` request, which on a production chat serve
(Wise preset: temp 1.0 / top_p 0.95 / top_k 20) is the very first
request.

:func:`warm_sampler_kernels` drives the SAME wrapper entries the serving
sampler dispatches into (``apply_top_k_top_p_triton``,
``gumbel_argmax_from_seeds``, ``sample_top_k_top_p_triton``,
``apply_penalties_fused``) over the bounded specialization set the
served config can reach:

* batch rows ``B``: Triton value-specializes integer scalars into three
  classes — ``== 1``, ``% 16 == 0``, other. The Gumbel kernel's cache
  key adds its ``split`` scalar (derived from ``num_SMs / B``) and the
  split-dependent ``BLOCK_SIZE`` constexpr, so the warm set is computed
  DEVICE-ADAPTIVELY: walk every ``B`` in ``1..max_batch``, mirror the
  wrapper's ``split``/``BLOCK_SIZE`` derivation, and warm one
  representative ``B`` per distinct specialization-class tuple. Those
  rows also cover every ``class(B)`` — the only ``B`` axis the mask and
  penalty kernels key on.
* mask-kernel AXIS ENABLEMENT (:data:`_TOPK_TOPP_AXIS_CLASSES`): each of
  ``k``/``p`` present-or-``None`` is a SEPARATE compiled kernel, twice
  over — it flips the ``TOPK_ENABLED``/``TOPP_ENABLED`` constexpr AND
  changes the disabled argument's POINTER DTYPE (``_launch_topk_topp``
  passes ``logits`` itself as the dummy pointer, so a disabled ``p``
  turns the ``*fp32`` P slot into ``*bf16``). The served sampler emits
  all three enabled combinations, because
  :func:`~arbi_serve.sampler.sampling_masks._fused_topk_topp_args`
  returns ``None`` for an axis no row asks for and ``top_p >= 1.0``
  counts as "not asked for" (``sampling_metadata``: ``any_tp`` is set
  only by ``top_p < 1.0``). A ``top_p=1.00`` model card — Qwen3.5-0.8B's
  documented non-thinking preset — therefore serves the ``(k, None)``
  kernel, which a ``(k, p)``-only warm never compiles. Each combination
  is warmed for BOTH the mask-only entry and the fused mask+sample one,
  plus the sample-only (``k=None, p=None``) variant of the SPMD
  noise-supplied path.
* penalty knobs (:data:`_PENALTY_KNOB_CLASSES`): ``HAS_REP`` /
  ``HAS_FREQ`` / ``HAS_PRES`` are constexprs on the fused penalty
  kernel, so each non-empty subset of the three knobs is its own
  compiled kernel. Which subsets and which COUNT SOURCE are reachable
  depends on whether the engine installed the accumulator pool — see
  :func:`_warm_penalty_kernels`. ``B`` is a runtime scalar there too
  (three classes), and the counts buffer's dtype and ROW STRIDE are part
  of the key.

min-p and the LSE-emitting variant stay lazy: min-p is not in any served
preset, and the MTP rejection sampler's LSE fusion is warmed by the MTP
verify boot path where MTP is enabled.

This enumerated set is the A-PRIORI half of warmup coverage: it is what
makes the very FIRST boot at a build/arch warm, before anything has been
observed. It is not the whole of coverage, and it is not where a newly
discovered gap belongs. A specialization that escapes it is named by
:mod:`arbi_serve.jit_detector` and recorded by
:mod:`arbi_serve.jit_replay`, which pre-compiles it at the next boot from
Triton's own serialized specialization — no shapes to reverse-engineer and
no warm code to keep in sync with the kernel. Add to this module only what
a first boot must have; let the record cover the rest.

Failure here is logged loud and non-fatal: an unwarmed kernel costs the
first matching request its compile latency (which the detector reports);
it never makes a result wrong.

One ``B`` per distinct Gumbel-kernel specialization tuple.

Mirrors :func:`~arbi_serve.sampler.gumbel_argmax_triton.
gumbel_argmax_from_seeds`'s ``split``/``block_size`` derivation for
every servable ``B`` and keeps the first ``B`` producing each new
``(class(B), class(split), block_size)`` tuple — the axes the
compiled-kernel cache is keyed on. Bounded: at most a handful of
rows regardless of ``max_batch``.

One ``(K+1, B)`` slate per distinct residual/bonus specialization tuple.

The verify slate is ``(K+1, B)`` FLAT rows, so
:func:`~arbi_serve.sampler.gumbel_argmax_triton.
residual_bonus_gumbel_from_seed` derives its ``split`` /
``BLOCK_SIZE`` from ``n_rows = (K+1)·B``, not from ``B`` — a
different ladder from :func:`_representative_rows`.

BOTH slate dims are walked, because both are servable and both move
the key. ``K`` is not fixed at the drafter's ``max_k``: a slate
accepts at the step's UNIFORM ``K_step``, which a per-row ``mtp_k``,
a drafter fallback, dynamic-``K`` or a depth gate can shorten to any
``1..max_k``. ``split = num_SMs // n_rows`` then lands on a different
value, and its divisibility-by-16 class is a cache-key axis — so
warming only the deepest chain leaves every shallower one cold.

Keeps the first ``(K+1, B)`` producing each new ``(class(B),
class(n_rows), class(K·B), class(split), block_size)`` tuple, the
axes the compiled-kernel cache is keyed on. The pointer strides are
NOT axes here: the warm slices the same persistent
``(max_k+1, max_batch, V)`` / ``(max_k, max_batch, V)`` /
``(max_k, max_batch)`` verify buffers the serving path slices, whose
leading strides are slate-invariant. Bounded: a handful of slates
regardless of ``max_k`` or ``max_batch``.

Empty below ``k_plus_1 == 2``: a ``K == 0`` step drafts nothing, so
there is no residual to draw and no launch to warm.

Pre-compile the sampler's Triton kernels for the served shapes.

Covers the stochastic tail (top-k/top-p mask, fused mask+sample,
counter-Gumbel) AND the fused penalty kernel — everything the
per-step sampler chain can dispatch outside the captured graphs.

Returns the number of warm launches issued (0 when CUDA/Triton is
unavailable). Allocations are small transients ((B, V) logits at
``B <= 16`` — :func:`_representative_rows` tops out there even at
``max_batch=64``) released when this returns.

Pre-compile the fused repetition/frequency/presence penalty kernel.

The per-step sampler chain dispatches it for any request that sets a
penalty, outside every captured graph.

Which variants are REACHABLE depends on the accumulator pool, so this
reads the same hook :func:`~arbi_serve.sampler.penalty_accumulator.
get_penalty_accumulator` dispatches on and warms only what the served
config can launch. The boot order guarantees the answer is final:
``_install_penalty_accumulator`` runs in the model-load phase, long
before the serve-kernel warmup.

* POOL INSTALLED (the production boot). Presence / frequency without
  repetition never reach a kernel at all — they are one subtraction of
  the accumulator's delta — so only the four ``HAS_REP`` subsets are
  warmed. ``counts`` is the accumulator's persistent occurrence
  buffer: INT32, contiguous ``(B, V)``, row stride ``V``.
* NO POOL. Every request builds the fallback matrix, so all seven
  subsets are reachable. ``counts`` is built there as
  ``torch.zeros(B, V + 1)[:, :V]`` — the trailing column is the
  scatter's pad sink, so its ROW STRIDE is ``V + 1``, not ``V``.

Both details are load-bearing: Triton keys the compiled kernel on the
pointer dtype and specializes integer arguments on divisibility-by-16,
so warming the wrong count source leaves the served one cold.
``rep``/``freq``/``pres`` are cast to the LOGITS dtype at the call
site, so the warm casts them the same way.

One representative row count per Triton integer-scalar class.

An ``n``-row launch is compiled per :func:`_spec_class` of ``n``, so
walking ``1..max_rows`` and keeping the first of each class covers
every reachable row count with at most three launches.

Pre-compile xgrammar's token-bitmask kernel for the served shapes.

Any request may carry a ``response_format``, so the structured-output
mask is always reachable at serving; it runs OUTSIDE every captured
graph, one fused ``apply_token_bitmask_inplace`` launch per step over
the constrained rows of the flat ``(K*B, V)`` logits view.

The compiled kernel is keyed on the LOGITS DTYPE (which also picks
``num_warps``, so each dtype is its own compile) and on the
integer-scalar class of ``num_rows`` — the count of CONSTRAINED rows,
which is 1 for a lone constrained request and up to the flat slate
width when several are. ``vocab_size`` and both row strides are
``B``-invariant, so the warm set is ``dtypes x class(num_rows)``:
at most six launches, independent of ``max_batch``.

Returns the number of warm launches (0 when CUDA or xgrammar is
unavailable).

Pre-compile the SAMPLED-draft tail — the drafter half of a non-greedy step.

Under ``ARBI_TRUE_STOCHASTIC_DRAFT`` (default ``auto``) a non-greedy
row does not draft by argmax: the drafter filters its own logits and
DRAWS from them, then reports the realized ``q``. That tail runs
outside every captured graph — boot capture records the GREEDY chain,
because capture has no request in hand and the sampled route is a
per-row decision — so on a stock boot its kernels compile inside the
first non-greedy request, next to the verify tail's.

Driven through :meth:`~arbi_serve.spec_decode.rejection_sampler.
GraphSafeRejectionSampler.sample_drafter_token`, the entry the drafter
chain calls, so the mask dispatch (fanned top-C vs pivot), the
parameter encoding and the draw offsets are the serving code's own.

Two independent axes, walked separately rather than as a product
because they key different kernels:

* the MASK (:func:`~arbi_serve.spec_decode.rejection_sampler_ops.
  mask_spec_logits_`) keys on the ``(top_k, top_p)`` shape — the same
  three served classes :data:`_TOPK_TOPP_AXIS_CLASSES` enumerates, as
  VALUES here (the drafter's encoder never passes ``None``: an
  unasked axis arrives as ``top_k = V`` / ``top_p = 1.0``);
* the fused filter+draw's ``DRAW_OFFSET`` is a constexpr the chain
  steps through — ``OFFSET_DRAFTER + 17 · step`` for ``step`` in
  ``0..K-1`` — so each drafted depth is its own compiled kernel.

Both ride the ``B`` ladder of :func:`_representative_rows`: the fused
draw derives ``split`` / ``BLOCK_SIZE`` from the row count exactly as
``gumbel_argmax_from_seeds`` does.

``keyed_watermark`` (a watermark key is configured) adds a THIRD axis:
the coupled draw reads a per-row context seed with zero counter lanes,
which is a ``PER_ROW_SEED`` constexpr on the same fused tail — a
different compiled kernel from the step-seeded one, and the only one a
watermarking boot's sampled drafts ever launch. It needs no draw-offset
ladder: the keyed mode does not read the offset, and the wrapper pins it
to 0 so every depth shares the one program.

min_p stays lazy, the same rule the rest of this module follows: it is
in no served preset, and the drafter's encoder only builds a min_p
tensor when a request asks for one.

Allocations are transients — one ``(B, V)`` fp32 logits block per row
class, plus whatever the drafter tail itself builds — released when
this returns. Returns the number of warm draws (0 without CUDA).

Driving the real entry also fires its flag-truth counters
(``drafter_fused_draw``, ``topc_mask`` / ``topc_mask_pivot_fallback``)
at Both are
declared ``Inapplicable`` — nothing asserts a count — and a bench that
brackets its window diffs the scrape rather than reading it absolute.

Refuse a verify buffer too small for the served slate.

A buffer the served ``(K+1, max_batch, V)`` / ``(K, max_batch)`` slate
does not fit inside means the warm and the serving path disagree on
the slate bound; slicing it anyway would warm a SHORTER slate — a
different specialization from the one traffic launches, i.e. exactly
the silent miscoverage the warm exists to prevent.

One ``(K+1, B)`` slate per distinct watermark-state specialization.

A different ladder from :func:`_verify_slates`: the state kernels are
keyed on ``KP1`` and on the integer class of the slate width ``B``
(which is also the committed/draft tensor's leading stride), not on
the residual draw's row-count and split axes.

Pre-compile the watermark ring/Bloom state kernels.

The keyed DRAWS are warmed alongside the verify sampler; the state
kernels that COMMIT them are reached by no sampler launch, so without
this they compile under the first live request.

``_wm_mtp_advance_kernel`` specializes on ``KP1``, so a boot whose
served draft length changes needs its own warm; it additionally
specializes on ``INSERT_BLOOMS``, the ring-only variant a
greedy-routed slate takes. ``_wm_mtp_seed_kernel`` specializes on its
depth range: the verify tail launches ``(0, KP1)``, and a stochastic
drafter launches ``(s, 1)`` once per slot ``s``. Every one of those is
servable, so every one is warmed. Warms against private scratch ring
and Bloom buffers, never the live cache.

Pre-compile the MTP verify tail's sampler kernels.

The verify rejection path runs OUTSIDE the captured graphs and
dispatches Triton entries the plain-decode warm never reaches:

* the slate mask, over ``(K+1)·B`` full-vocab rows and across the same
  :data:`_TOPK_TOPP_AXIS_CLASSES` the plain-decode warm walks (the
  verify path builds its ``(k, p)`` tensors from the SAME per-request
  metadata, so a ``top_p >= 1.0`` preset serves the ``p=None`` kernel
  here too), with and without the LSE epilogue (``lse_out`` — the
  Philox verify epilogue on the GATHERED path; ``WRITE_LSE`` is a
  constexpr, so the two are different kernels). Driven through
  :func:`~arbi_serve.spec_decode.rejection_sampler_ops.
  mask_spec_logits_`, the serving DISPATCH, not through one of the two
  kernels beneath it: that dispatch chooses the fanned top-C mask or
  the vendored pivot kernel from the HOST ``(top_k, top_p)`` values and
  the logits dtype, so a warm that calls either kernel directly leaves
  whichever the served values actually select cold;
* the SHARDED path's candidate-width mask: ``compacted_p_target``
  masks ``(rows, C)`` logits where ``C = tp_size · n_cand`` (the
  all-gathered per-shard top-``n_cand`` candidates) — a row width
  three orders of magnitude below ``V``, i.e. a DISTINCT kernel
  specialization from any full-vocab warm;
* :func:`~arbi_serve.sampler.gumbel_argmax_triton.
  gumbel_noise_at_indices` at the merged candidate shape
  ``(K+1, B, C)`` with the recovery draw-offset class;
* ``mtp_sample_residual_philox`` — the DEFAULT (non-fused) Philox
  verify tail, whose single-pass counter-based recovery+bonus draw
  is :func:`~arbi_serve.sampler.gumbel_argmax_triton.
  residual_bonus_gumbel_from_seed`. Warmed through the OP, the entry
  ``GraphSafeRejectionSampler.sample_batched`` dispatches, so the
  wrapper arguments the specialization is keyed on are built by the
  serving code rather than restated here. The kernel compiles per
  ``(Q_IS_ONEHOT, IN_LOG_SPACE, DRAW_OFFSET, VOCAB_SIZE, BLOCK_SIZE)``
  constexpr class plus the integer-scalar specialization of its
  ``(K+1)·B`` rows, ``K·B``, ``B`` and ``split``, so the warm walks
  BOTH slate dims (:func:`_verify_slates`) and BOTH proposal
  encodings:

  - ``draft_probs=None`` — the point-mass proposal a greedy-argmax
    drafter reports (``POINT_MASS_Q``);
  - ``draft_probs=verify_draft_probs[:K, :B]`` — the dense ``q`` a
    boot that can realize a SAMPLED draft hands the sampler.
    ``ARBI_TRUE_STOCHASTIC_DRAFT`` defaults to ``auto``, so on the
    served config this is the encoding the FIRST non-greedy request
    reaches; ``Q_IS_ONEHOT`` being a constexpr makes it a different
    compiled kernel from the point-mass one.

  The ``in_log_space=True`` fused variant (``mtp_fused_rejection``,
  default off) stays lazy — the "only warm what this boot can reach"
  rule the top-k/top-p / min-p warm follows; the detector names it if
  a request ever gets there.

``keyed_watermark`` (a watermark key is configured) adds the
EXACT-MATCH verify draw ``_sample_batched_keyed_exact`` dispatches:
the same residual/bonus kernel under ``PER_ROW_SEED=True`` (a
``(K+1)·B`` seed grid instead of the scalar step seed) over a
freshly-built all-``-1`` draft tensor — a distinct compiled kernel
from the scalar-seed one, and the ONLY verify draw a watermarked
slate reaches. Warmed off the key alone, not off the live apply gate:
``watermark_enabled`` is flippable on a standing server, so a boot
that resolves the key must cover the draw the flip turns on.

``verify_p_target`` / ``verify_draft_tokens`` / ``verify_draft_probs``
are the ENGINE's persistent verify buffers —
``VerifyBuffers.sc_p_target`` ``(max_k+1, max_batch, V)``,
``VerifyBuffers.draft_tensor`` ``(max_k, max_batch)`` and
``VerifyBuffers.sc_draft_probs`` ``(max_k, max_batch, V)``. The
residual warm slices them exactly as the serving path does
(``[:K+1, :B]`` / ``[:K, :B]``) instead of allocating fresh
contiguous tensors, because Triton keys the compiled kernel on every
integer argument's specialization class and the slate's LEADING
STRIDES are arguments: a ``[:, :B]`` slice of a ``max_batch``-wide
slab carries row stride ``max_batch`` (``max_batch · V`` for the
vocab-scale ones), which a fresh ``(K, B)`` / ``(K+1, B, V)``
allocation does not reproduce for ``B < max_batch``. Warming on the
serving buffers makes every stride, dtype and alignment match by
construction, and costs no memory — the buffers already exist and the
KV budget already books them, so this warm adds nothing to the
forecast the pool was sized from. All three are read-only here.

``verify_draft_probs`` doubles as the dense-``q`` REACHABILITY answer:
``VerifyBuffers`` allocates ``sc_draft_probs`` on exactly
:func:`~arbi_serve.spec_decode._mtp_driver_ops.dense_draft_q_possible`,
the predicate that decides whether a slate can carry a dense
proposal, so ``None`` means the dense kernel is unreachable and
warming it would be a boot tax for a kernel nothing launches.

``None`` for any of them (no verify buffers — a vocab-less pool, CPU
tests) falls back to fresh contiguous tensors, except
``verify_draft_probs`` whose absence skips the dense class.

Parameters are the RESOLVED served values (drafter ``K+1``, the
``sharded_verify_n_cand`` budget, the TP world size,
``OFFSET_RECOVERY``, the verify buffers) — passed in by the boot
phase so this module reads no engine state. Returns the number of
warm launches (0 without CUDA).

Sliding-window no-repeat-n-gram logits processor.

Blocks the token that would complete an n-gram whose first ``n−1``
tokens match the sequence's current ``n−1``-token suffix, when that
n-gram already occurred inside the lookback window. OCR models
(Unlimited-OCR) rely on this to stop coordinate-token loops — greedy
decoding over dense layout/coordinate vocabularies otherwise repeats a
long span verbatim. ``window == 0`` degenerates to the
unbounded-lookback ``no_repeat_ngram_size`` behaviour.

Per-request opt-in via ``SamplingParams.no_repeat_ngram_size`` /
``no_repeat_ngram_window`` (both surfaced on the OpenAI request
bodies). Stateless — history is re-read from the request each step, so
there is no per-request state to keep in sync.

The history comes from
:func:`~arbi_serve.sampler.penalty_history.sequence_history_parts`, the
same accessor the repetition/frequency/presence penalties use, so
"the tokens this request is scored against for repetition" is ONE
definition rather than two that can drift. For every ordinary request
that is ``prompt_token_ids + output_token_ids``, exactly as before. For
the frame-lockstep duplex text channel it is the connection's real
SPOKEN tokens (design doc §7.53/§7.59) — which is what makes this
processor usable there at all: a duplex row's ``output_token_ids`` is
one entry per 80 ms tick and ~95% PAD, so n-grams over it would be
runs of PAD and the blocker would ban the very token that holds a turn
open. Filtering PAD/BOS/EOS out of the history is the whole carve-out:
every banned id is by construction a member of the history, so a token
absent from it can never be banned.

Refuse an n-gram-blocked row on a speculative draw.

The blocker bans the tokens that would complete an n-gram already seen
in the row's history, so its ban list is keyed on the row's current
``n-1``-token suffix. Every position of a speculative slate has a
DIFFERENT suffix, and the drafts that produce those suffixes are exactly
what the draw is deciding whether to keep. Masking all positions against
the committed suffix bans tokens that are legal there and misses the
ones that are not.

A request that set ``no_repeat_ngram_size`` therefore takes the ``K=1``
decode path (``mtp_k=0``, resolved at admission): the block it asked for,
without the speculation speedup.

Mask tokens that would repeat an n-gram within the lookback window.

Implements the :class:`arbi_serve.sampler.processors.LogitsProcessor`
protocol. Rows whose request has ``no_repeat_ngram_size == 0`` pass
through untouched; a step with no opted-in request returns the
input tensor unchanged (zero-copy).

Refuse the slate for every row that asked for the n-gram block.

The ban list is a function of the row's ``n-1``-token suffix, which
differs at every speculative position, so there is no per-position
form of this constraint to apply — see :func:`_refuse_speculative`.
Rows at ``no_repeat_ngram_size == 0`` are untouched, so a slate that
contains none is a no-op.

``mtp_k`` is forced to 0 for an opted-in request at admission, which
is what keeps this refusal off the served path.

Tokens that would complete an already-seen n-gram.

The scan considers n-grams STARTING in
``[len(seq) − window, len(seq) − n]`` (whole sequence when
``window <= 0``) and bans the final token of every n-gram whose
first ``n−1`` tokens equal the current suffix.

Omni output guard — mask speech-code logits off speech turns.

On an omni model every token id is reachable, but a request that did
NOT ask for speech output (`modalities` without ``"audio"``) should
never emit speech codes or tts control markers: the demux would swallow
them (no text corruption), but they waste KV pages / token budget and
can derail the turn. This processor applies a precomputed ``-inf`` bias
over the speech-code range and the control-id set to every row whose
request has an :class:`OmniStreamState` with speech DISABLED; rows that
seeded a speech turn are untouched (the model interleaves freely there
— its text tokens ARE the transcript, so no in-segment constraint is
imposed).

Same registration surface as the xgrammar processor
(``Sampler.set_logits_processors``); logits arrive ``(K, num_seqs, V)``
and rows map 1:1 to ``requests``.

Mask the speech ids across EVERY position of a guarded column.

The ban is a fixed per-request vocabulary set, not a function of the
row's history, so the same mask is correct at every speculative
position and the drafts are irrelevant. A guarded row therefore
cannot emit a speech code through the verify draw either.

A vocab-parallel ``window`` restricts the ban to the ids that rank
holds — a slice of the same boolean mask, since the ban is a set of
token ids and nothing else.

Persistent additive accumulator for the presence / frequency penalties.

Both penalties are ADDITIVE in logit space::

    penalty_add[b, t] = presence[b] * (count[b, t] > 0) + frequency[b] * count[b, t]

so one ``(rows, vocab)`` buffer can carry both, and applying them is a
single coalesced ``logits - penalty_add[:B]``.

Repetition does not share that term — it is multiplicative and
sign-dependent — but it does share the STATE. All three knobs are
functions of the same per-``(row, token)`` occurrence history: repetition
and presence need the SEEN set, frequency needs the counts themselves.
:attr:`PenaltyAccumulator._counts` already holds it, so a caller that
asks for ``need_counts`` gets the seen set out of the same persistent
buffer and applies its own arithmetic on top (see
:func:`~arbi_serve.sampler.processors.apply_penalties`).

Why a persistent buffer
-----------------------
The alternative is a ``(B, vocab+1)`` fp32 occurrence matrix rebuilt EVERY
decode step: allocate it, zero it, ``scatter_add_`` the whole
``B x history`` id matrix into it, then read it back in the penalty
kernel — per-step cost proportional to CONTEXT. The accumulator scatters
only the tokens EMITTED THIS STEP, which is O(B), and the step's device
work is the one subtraction.

Both the fold and the subtract are also free of any host-blocking device
read, which is the property that lets the engine keep running the next
step ahead of the GPU. See :meth:`PenaltyAccumulator.update`.

Bit-identity
------------
The accumulator holds the SAME value the counts chain would subtract, so
the result is bit-identical whenever the chain has exactly one additive
term:

* PRESENCE ONLY (the production knob): ``scatter_`` writes ``pres[b]``,
  an idempotent set, and the chain's ``pres_b * seen.to(dtype)`` is
  ``pres[b]`` exactly. Verified with ``torch.equal``.
* FREQUENCY ONLY: the accumulator entry is RECOMPUTED from the running
  count as ``freq[b] * count.to(dtype)`` — the chain's expression
  verbatim — not incremented by ``freq`` per occurrence. Accumulating
  per occurrence instead would drift (the chain rounds ONCE, an
  incremental sum rounds every step); recomputing costs one extra
  gather over the touched positions only, which is O(B).
* BOTH ACTIVE: the chain subtracts in TWO steps
  (``(logits - freq*counts) - pres*seen``), rounding the intermediate;
  the accumulator subtracts their sum in one, rounding once. This is the
  ONLY combination that is not bit-identical, and it is the only one
  that is not reachable from a single-knob request. The deviation is
  bounded at one ulp of the logits dtype and
  ``tests/test_sampler_penalty_accumulator.py`` asserts that bound.

Row indexing
------------
Row ``i`` holds slate position ``i``, so the apply is a plain slice with
no gather. When the slate reorders (a request finishes and the rest
compact up) the rows are PERMUTED to match, by row-copy through a single
scratch row — at most one ``vocab``-wide copy per moved row, and zero
copies in steady decode where the slate is stable. Rows whose tenant is
new are zeroed and re-seeded from whatever
:func:`~arbi_serve.sampler.penalty_history.penalty_history_parts` says
that row's knobs are scoped to — OUTPUT ONLY for presence / frequency
(the OpenAI contract, and what vLLM and sglang implement), ``prompt +
output`` only when the row asks for a repetition penalty.

This used to seed every row from the full ``prompt + output``, which made
presence and frequency cover the prompt. That is not a harmless superset:
at a 16k prompt with ``presence_penalty=1.5`` every distinct token in the
document is pushed down on every decode step, so the engine penalises the
vocabulary of the very text it was asked to summarise, and the distortion
grows with prompt length.

The seed is O(history) HOST work, ONCE per request, on the step that
admits it — the same work the history cache already does for a fresh row.
Steady decode appends only the tokens emitted since the last step.

FOLLOW-UP (not done): bound the history to a sliding window rather than
the whole output. It is arguably better behaviour — a penalty that never
forgets makes long generations progressively harder to continue — and it
also caps the fold at O(window) instead of O(output), so it saves compute
rather than costing it. Needs a decision on the window and on how it
interacts with the repetition scope before it can ship.

Per-row presence / frequency knobs as host floats.

Read off the requests when they carry sampling params — which every
engine caller does — so a device-resident knob tensor is never pulled
back. The tensor fallback is for direct callers that pass knobs
unattached to requests.

The ``(B,)`` fp32 knob vector, on the HOST.

The scattered value must be the knob taken ``fp64 -> fp32 -> dtype``,
which is the hop ``_build_sampling_metadata`` takes when it stores a
Python float into an fp32 tensor and uploads it. Building the same
fp32 vector here reproduces that bit pattern without reading the
device-resident copy back.

Host tensor -> device WITHOUT a pageable-H2D host stall.

``.to(device, non_blocking=True)`` off a PAGEABLE source is not async at
all: the CUDA runtime serializes it as ``cudaMemcpyAsync`` +
``cudaStreamSynchronize``, so the host blocks until the stream drains.
The penalty fold runs AFTER the step's forward is enqueued, which makes
that drain cost a whole forward rather than the copy's own microseconds
— the same mechanism, and the same fix, as the

Staging through the CUDA caching HOST allocator is safe for a temporary:
the allocator defers reuse of the pinned block until the recorded copy
event completes. Non-CUDA devices take the plain copy (``pin_memory``
requires CUDA) and land identical values.

Pinning is defined only for dense CPU memory, so a source that is ALREADY
device-resident takes a D2D copy (a no-op when it is already on ``device``)
instead. That source is not a production shape — ``_host_knobs`` builds the
knob vector on the host precisely so it never is — but the repair for one
is never ``.cpu()``: that is the D2H readback this helper exists to avoid,
and the one
``tests/test_sampler_penalty_no_device_readback_gpu.py`` forbids.

Slate-aligned ``(rows, vocab)`` additive penalty accumulator.

One instance per ``(device, vocab_size, dtype)``. Not thread-safe —
the sampler runs single-threaded inside the engine step.

``allocate`` is the hook that puts the buffer in a NAMED cuMem pool
(``scratch.penalty_accum``); it takes ``(shape, dtype)`` and returns a
zeroed tensor. The default allocates from whatever pool is current,
which is what the CPU tests and any non-engine caller want.

Engine boot hook: route accumulator allocations into a named pool.

Until this is called the accumulator is DISENGAGED on CUDA (see
:func:`get_penalty_accumulator`): a multi-tens-of-MB device buffer
allocated outside every named pool lands in
``unpooled.torch_default_pool``, which the VRAM ledger reports as
unbudgeted and which KV sizing has already spent.

Return the accumulator for this ``(device, vocab, dtype)``, or ``None``.

``None`` means the caller must take the counts path: on CUDA that is
the case whenever no engine has installed a pool hook, which keeps the
buffer from ever being allocated unbudgeted.

``device`` IS PART OF THE KEY, BY ITS STRING FORM. ``torch.device("cuda")``
and ``torch.device("cuda:0")`` are different keys and therefore different
accumulators, even though they name one GPU. The engine always builds the
indexed form (:mod:`arbi_serve.engine.engine_init`), so serving is
consistent; a harness or test that primes with the bare ``"cuda"`` and
then drives a path reading ``tensor.device`` gets a SECOND, unprimed
instance whose counts buffer was never materialized — which silently
routes :func:`~arbi_serve.spec_decode.verify_penalties._committed_counts`
down its padded-history fallback and measures the wrong path. Pass the
indexed device.

Forget the installed pool hook and every accumulator built through it.

The counterpart to :func:`install_penalty_accumulator_pool`, for the
stable-VA member-build seam
(:func:`arbi_serve.engine.member_scratch_retire.retire_member_scoped_device_caches`).
The installed closure captures the OUTGOING member's
``scratch.penalty_accum`` pool (and a carve of its persistent slab); once
that member is parked and its pool physical evicted, an allocation through
the stale closure lands in unmapped VA. Leaving it disengaged until the
incoming member's build re-installs its own is exactly the pre-install
behaviour (:func:`get_penalty_accumulator` returns ``None`` on CUDA), so
nothing is served from a dead pool in the window.

Allocate the FULL buffer set now; return the bytes taken.

Called once at engine boot. The bytes were already subtracted from
the KV budget (``predict_penalty_accum_pool_bytes``), so taking
them eagerly is what makes the reserve and the residency agree —
growing later would spend VRAM the KV pool was handed.

Materialize the occurrence-count buffer.

Returns the buffer; :attr:`_counts_seeded` records whether it was
already live, because a buffer created AFTER rows have folded
history in starts empty and every live row has to be reseeded.

The persistent ``(rows, vocab)`` additive penalty entry, or ``None``.

``None`` before :meth:`prime` (or any :meth:`update`) has
materialized the buffer. Row ``i`` is slate position ``i``, same as
:meth:`update`'s return value -- this IS that tensor, exposed for a
reader that wants the entry without driving a fold.

The drafter's penalty alignment
(:mod:`arbi_serve.spec_decode.draft_penalties`) holds this view
across a CUDA-graph capture, so the buffer's ADDRESS must outlive
every replay. It does exactly as long as :meth:`prime` ran first:
the boot reserve allocates ``max_rows`` up front and
:meth:`_ensure_rows` is then a no-op forever after.

The persistent ``(rows, vocab)`` int32 occurrence counts, or ``None``.

``None`` until some call has asked for them (``need_counts`` or a
frequency penalty). Row ``i`` is slate position ``i``, same as
:meth:`update`'s return, so the view needs no gather.

One-way copy ``src`` -> ``dst``; ``src``'s tenant is not preserved.

Used when the row being overwritten holds a request that has LEFT
the slate — there is nothing to save, so the move is one copy
instead of a swap's three.

Permute rows so row ``i`` holds ``desired[i]``.

Selection sort over at most ``B`` positions. A stable slate does
ZERO device work — the common case, and the whole reason rows are
slate-aligned rather than gathered at apply time.

The bound on the rare case is worth stating. A row whose current
occupant has left the slate is overwritten one-way (ONE row copy);
only a genuine reorder, where both rows' tenants are still live,
needs a three-copy swap. The dominant churn shape — one request
finishes and the rest compact up — is entirely one-way, so it
costs ``B - 1`` row copies, not ``3(B - 1)``. A pathological full
reorder is still bounded at ``3(B - 1)`` copies on the ONE step it
happens; a per-step counts matrix pays its full cost on EVERY step.

Fold this step's new tokens in; return the ``(B, vocab)`` delta.

``None`` when every row's history is empty — there is nothing to
subtract, and the caller returns its input untouched rather than
materialise a zero delta over the whole vocab.

``need_counts`` keeps the occurrence buffer live even when no row
asks for a frequency penalty, so a caller reading the seen set
(repetition) gets it from :meth:`counts_view` instead of
rebuilding one.

``want_accum=False`` says the caller will NOT read the returned
delta — the fused kernel derives presence and frequency from the
counts directly on a repetition batch — so the additive entry is
left unfolded and the rows are flagged for repair. The return value
is then only meaningful as ``None`` / not-``None``.

``pres`` / ``freq`` are the ``(B,)`` fp32 knob tensors. They are a
FALLBACK source of the per-row knob floats, for direct callers
whose requests carry no sampling params; every engine caller's
knobs are read off the requests instead. Nothing on this path
touches the tensors' storage, because ``_build_sampling_metadata``
builds them on the LOGITS device and any read of a device tensor —
``.tolist()``, or indexing it with a host index vector — is a D2H
copy that BLOCKS the host until the stream drains. Once per decode
step that costs the whole pipelining win, not the copy.

The value actually scattered is taken ``fp64 -> fp32 -> dtype``,
the same hop the counts chain's ``pres.to(device, dtype)`` takes,
so its bit pattern is unchanged; ``_host_knobs`` does the fp32 leg
on the host, and fp32 -> bf16/fp16 is round-to-nearest-even on
host and device alike. Casting from the raw Python float instead
would be a DIFFERENT rounding (fp64 -> dtype) and could differ in
the last bit.

Scatter this step's new tokens into the accumulator.

Indices are FLATTENED to ``row * vocab + token`` so rows with
different tail lengths need no padding column — the steady-state
tail is one token per row, so this is a ``B``-element H2D copy.

The knobs arrive as HOST floats and the per-element value vector is
built host-side, so the whole fold is host reads plus one H2D. The
obvious alternative — indexing the sampler's device-resident knob
tensor with the host row vector — is a CUDA gather with a CPU
index, which pulls the index across and BLOCKS the host until the
stream drains. See :meth:`update` for why that matters.

``additive`` is False when no row asks for presence or frequency —
a repetition-only batch. The accumulator entry is then zero
everywhere and stays that way, so only the occurrence counts are
folded and the value chain (gather + the entry scatter) is skipped
entirely.

Rebuild the additive entry of ``rows`` from the occurrence counts.

The entry is a pure function of the counts
(``pres * (count > 0) + freq * count``), so a row whose entry was
skipped while a repetition batch held the slate is recovered with
one device pass per row — no history walk, no host work
proportional to context. Bit-identical to what :meth:`_fold` would
have written: ``pres * 1`` is ``pres`` exactly, and the remaining
multiply-then-add rounds in the same order.

Device-resident, incrementally-appended token histories for the penalty path.

Why this exists
---------------
:func:`arbi_serve.sampler.processors.apply_penalties` needs, per decode
step, a padded ``(B, max_hist)`` tensor of every row's ``prompt +
output`` token ids so one ``scatter_add_`` can build the occurrence
counts. Rebuilding that from the Python lists on every step is O(total
context) HOST work per step, and it dominated the whole penalty path:

    B=64, V=248320, hist=16896, measured on a 4090 (sm_89)
        apply_penalties TOTAL         56.6 ms
        rebuild from Python lists     53.6 ms   (94.8%)
        everything on the device       0.9 ms   ( 1.7%)

The prompt never changes and the output grows by a handful of tokens per
step, so the history is APPEND-ONLY in the common case. This cache keeps
each row's ids in a pooled device buffer and appends only the new tail,
turning the per-step host cost from O(context) into O(tokens appended).
``apply_penalties`` at that shape goes from 56.6 ms to 0.435 ms.

End-to-end, Qwen3.5-0.8B TP1 bf16, 16k prompts, ``--ignore-eos`` pinning
both arms to exactly 512 output tokens, one server boot per row, ONLY
``presence_penalty`` varying (TPOT ms, two reps each)::

                    c16 pp=0      c16 pp=2    tax      c64 pp=0      c64 pp=2     tax
    before   13.92 / 13.77   20.13 / 20.22   +46%   21.41 / 21.51  40.62 / 41.92  +93%
    after    13.76 / 13.91   14.53 / 14.71   + 5%   21.41 / 21.49  23.22 / 22.79  + 7%

The ``pp=0`` control is unchanged to within noise, which is what makes
the ``pp=2`` column attributable. Note what it also says: arbi's
penalty-FREE decode (13.8 ms at c16) is already slower than vLLM's and
SGLang's penalty-INCLUSIVE decode (10.1 / 9.8 ms), so the residual
cross-engine gap at c16/c64 is baseline decode, not this path — do not
expect more of it from here.

Correctness
-----------
The returned tensor is bit-identical to the freshly rebuilt one: same
ids, same ``vocab_size`` pad sink, same dtype. A row is REBUILT from
scratch whenever the append-only assumption cannot be proven:

* the row is new;
* ``len(prompt_token_ids)`` changed — the recompute/preemption path
  rewrites the prompt (``scheduler/preemption.py`` clears
  ``output_token_ids``; ``distributed/spmd_mirror.py`` folds the emitted
  output back into the prompt);
* the total history got SHORTER than what is cached (preemption reset).

Within a fixed prompt length a longer-or-equal total can only come from
appends to ``output_token_ids``, which is append-only, so the cached
prefix is provably still valid.

That covers preemption/recompute exactly, and not by luck: BOTH rewrite
sites set ``new_prompt = prompt_token_ids[:prompt_consumed] +
output_token_ids`` and then clear the output, so the concatenated
``prompt + output`` the penalty path reads is a PREFIX of what it was.
It is unchanged when the prompt was fully consumed (``prompt_consumed ==
len(prompt)``, so the same ids merely move across the boundary) and
strictly shorter otherwise — which the length check catches.

Keyed by ``request_id`` (an ``int`` on both :class:`Request` and the SPMD
``_MirroredRequest`` view) rather than object identity, because the SPMD
path builds a FRESH per-row view object every step. Ids are unique only
WITHIN one engine instance — ``engine_init`` seeds a fresh
``itertools.count(1)`` — so a rebuilt engine would reuse them. The
process-wide caches are therefore dropped in ``Sampler.__init__``, which
runs once per engine build.

The ``(prompt, output)`` id sequences a penalty must count over.

SCOPE IS PER-KNOB, and this is the function that decides it:

* ``presence_penalty`` / ``frequency_penalty`` count the OUTPUT ONLY.
  That is the OpenAI contract and what vLLM and sglang implement.
* ``repetition_penalty`` is the one knob scoped to ``prompt + output``.

So the prompt is included here only when the row actually asks for a
repetition penalty. Counting the prompt for presence/frequency is not
a harmless superset: at a 16k prompt with the Qwen non-thinking
recipe's ``presence_penalty=1.5``, EVERY distinct token in the
document takes a -1.5 logit hit on every decode step, so the engine
penalises the exact vocabulary of the text it was asked to summarise.
The damage grows with prompt length, which is why it showed up as a
long-context regression before it showed up as a quality complaint.

KNOWN RESIDUAL: a row that sets repetition != 1.0 AND presence or
frequency != 0 gets the prompt counted for all three, because the
accumulator keeps ONE occurrence buffer per row and repetition needs
the prompt in it. Splitting that needs a second buffer (or a separate
prompt-seen mask threaded through the fused kernel) and is not worth
the VRAM for a combination no shipped profile uses — Qwen ships
``repetition_penalty=1.0`` in every profile.

A request may override the history entirely by setting
``penalty_token_ids``, in which case THAT list is the whole history
and the prompt contributes nothing.

The override exists for the frame-lockstep duplex text channel (design
doc §7.53), whose ``output_token_ids`` is one entry per 80 ms tick and
therefore ~95% PAD: scored raw, a repetition penalty would drive PAD —
the token that holds a turn open and closes it — down on every step,
and the turn machinery is built on PAD being freely re-emittable. The
duplex layer instead maintains a list of the connection's REAL spoken
text tokens, so PAD/BOS/EOS are absent from the history and thus can
never be penalized. Filtering the history IS the special-token
carve-out; no sampler branch is needed for one.

Both penalty caches (:class:`PenaltyHistoryCache` and
``PenaltyAccumulator``) assume an append-only history and detect a
rewrite by comparing lengths, so an override must only ever grow —
which the duplex list does, one entry per real text token.

The FULL ``(prompt, output)`` the request has seen and said.

Scope-free: this is "what tokens exist", with only the
``penalty_token_ids`` override applied. Consumers that must see the
prompt regardless of any penalty knob read this — notably
``no_repeat_ngram_size``, which has to block an n-gram that occurred
in the PROMPT and would silently stop doing so if it inherited
presence/frequency's output-only scope.

``penalty_token_ids`` is read as a plain attribute, so every
Request-shaped view the scoring path is handed must declare it. A view
that merely LACKED it would read as "no override" and be scored against
``prompt + output`` — which for the one request type that sets the
override is the PAD stream the override exists to keep out, and a wrong
history is not a thing that can be noticed downstream.

Return the ``(B, max_hist)`` int64 padded history, or ``None``.

``None`` means every row's history is empty (nothing to count).
The pad value is ``vocab_size``, matching the ``(B, V+1)`` counts
buffer's discarded sink column.

Single-pass fused repetition / frequency / presence penalty kernel.

The PyTorch chain in :func:`arbi_serve.sampler.processors.apply_penalties`
is vectorized but not fused: it materialises six-to-eight full
``(K, B, V)`` intermediates, each one a separate launch reading and
writing the whole vocab. This kernel does the whole chain in ONE pass:
read ``logits`` + ``counts``, write ``out``.

per call and the speedup::

     B   counts MiB   presence only        all three knobs

There is deliberately NO runtime flag gating this. The kernel is
bit-identical to the chain it replaces and measured a tie-or-win, never a
loss, so there is nothing for an operator to A/B; the unfused PyTorch
chain remains as the non-CUDA path, which keeps it exercised. Three
things to read off that table before changing this file:

* Presence-only at ``B <= 64`` is a TIE, not a win. The counts matrix
  are nearly free — a pure memory-traffic argument overestimates the
  win. Our Matrix A cell (B=64, presence only) sits exactly on that
  neutral point, which is why the end-to-end A/B cannot separate fused
  from unfused.
* The win is real and large wherever counts spills L2 (B >= 128) or
  wherever repetition/frequency are also active, because the unfused
  path then materialises the repetition branch as well. Fusion is what
  makes the penalty cost flat in WHICH knobs a tenant sets, rather than
* None of this is the headline win. The penalty path's cost was 94.8%
  HOST-side history rebuilding (see :mod:`penalty_history`); fusing the
  device tail is the last few percent, not the fix.

Bit-identity
------------
The result is BIT-IDENTICAL to the PyTorch chain, not merely close.
PyTorch's CUDA elementwise kernels evaluate bf16 operands in fp32
(``opmath_t = float``) and round to bf16 on store, so every intermediate
in the reference chain is a bf16 value. The kernel reproduces that by
keeping every intermediate in the LOGITS dtype and issuing one Triton op
per PyTorch op, which lowers to the same cvt-f32 / op / cvt.rn-bf16
sequence.

Two things had to be got right, both caught by
``tests/test_sampler_penalty_fusion.py`` and neither visible under a
tolerance-based comparison:

* Do NOT hand-roll the rounding as ``.to(dtype).to(tl.float32)`` around
  each step. The compiler folds the truncate/extend PAIR away, the
  intermediates silently keep fp32 precision, and the result then
  diverges from the reference at exact round-to-even ties.
* Triton's default fp32 divide is APPROXIMATE. ``tl.fdiv(...,
  ieee_rounding=True)`` is required for the repetition branch.

The reference chain being replicated, in order::

    counts3 = counts.to(dtype)  # (1, B, V)
    seen3 = counts3 > 0
    rep_adj = where(logits > 0, logits / rep, logits * rep)
    out = where(seen3 & (rep != 1.0), rep_adj, logits)
    out = out - freq * counts3
    out = out - pres * seen3.to(dtype)

Each of ``HAS_REP`` / ``HAS_FREQ`` / ``HAS_PRES`` is a ``constexpr``, so
a batch that asks for presence only compiles a kernel with the
repetition and frequency arithmetic removed entirely.

Two count sources
-----------------
``counts`` is either the per-step ``(B, V+1)[:, :V]`` fp32 matrix the
fallback path builds, or the persistent int32 occurrence buffer the
accumulator maintains. Same kernel, same arithmetic — they differ only in
the pointer dtype and the row stride, both of which Triton specializes on,
so both are in the boot warm set.

Logits processors operating on ``(K, num_seqs, vocab_size)`` tensors.

Every processor takes a 3-D logits tensor and returns one of the same
shape; ``K=1`` in the non-MTP path. Constant-K shape costs zero at
runtime and keeps the surface MTP-ready.

The :class:`LogitsProcessor` protocol is the seam structured-decoding
processors plug into. The free-function processors below
(``apply_penalties``, ``apply_logit_bias``, ``apply_temperature``)
are the built-in chain.

Vectorization. Each free-function processor short-circuits when no
request needs it (greedy / default knobs across the batch). When the
batch is mixed the per-request work is done with one batched tensor op
instead of a Python loop over rows: penalty counts via a single padded
``scatter_add_`` on ``(B, max_recent)`` history, temperature via an
in-place ``logits.div_(temp.unsqueeze(1))``, and logit-bias via a
single padded ``scatter_add_`` on ``(B, max_bias_keys)`` ids/vals.

Per-step logits transformer.

A processor reads ``(K, num_seqs, V)`` logits and the per-request
list and returns a same-shape tensor. It MAY mutate per-request
internal state (e.g. an xgrammar matcher) when invoked, but MUST
NOT mutate other requests' state — the multi-tenant invariant
requires processors to be reentrant across the request set.

Stateful processors keep their per-request state on
``Request.grammar_state`` (or another request-scoped slot), never
on the processor instance itself.

Return the padded ``(B, max_recent)`` int64 history, or ``None``.

``None`` when every request's history is empty.

The pad value is ``vocab_size`` so a downstream ``scatter_add_`` into
a ``(B, V+1)`` counts buffer hits a sink column we discard. This
avoids materialising a separate row-mask scatter.

Delegates to :class:`~arbi_serve.sampler.penalty_history.PenaltyHistoryCache`,
which keeps the ids resident on the device and appends only each
step's new tail. Rebuilding this from the Python lists every step was
94.8% of the whole penalty path at ``B=64, hist=16k`` (53.6 ms of
56.6 ms) — the histories are append-only, so re-reading O(context)
ints per step was pure waste.

Whether this step takes the single-pass Triton penalty kernel.

Not behind a runtime flag: the kernel is bit-identical to the chain
below (``tests/test_sampler_penalty_fusion.py`` asserts
``torch.equal`` across bf16/fp16/fp32, four shapes and four knob
combinations) and measured a tie-or-win, never a loss — so there is
no quality or perf question for an operator to A/B. The unfused chain
stays as the non-CUDA path, which is what keeps it exercised.

Triton is imported lazily so the CPU path never needs it.

Apply repetition / frequency / presence penalties per request.

``window`` restricts the write to one vocab-parallel rank's slice: the
per-``(row, token)`` occurrence state is a global-vocab quantity, so it
is read at full width and the WINDOW is what narrows — a view of the
resident accumulator, never a shard-side reimplementation of the
arithmetic. ``None`` is the full vocabulary and is byte-identical to
having no window parameter at all.

Scope is PER-KNOB and lives in
:func:`~arbi_serve.sampler.penalty_history.penalty_history_parts`:
presence / frequency count the OUTPUT ONLY, repetition counts
``prompt + output``. Returns
a new tensor of the same shape by default; in-place is not the default
because penalties may need to write to logits that are still backed by
a CUDAGraph output buffer.

``inplace`` writes the result INTO ``logits`` and returns it, for a
caller that was going to ``copy_`` the result back anyway and owns the
buffer. On the production path — presence and/or frequency with no
repetition — that turns three full-vocab passes (read, write a fresh
tensor, read it back, write the original) into ONE fused subtract.
Only worth asking for on a caller whose tensor is big enough for that
to matter: the MTP verify slate is ``(K+1, B, V)``, i.e. K+1 times the
K=1 sampler's own draw.

All three knobs are functions of ONE per-``(row, token)`` occurrence
history, which the persistent accumulator
(:mod:`arbi_serve.sampler.penalty_accumulator`) maintains
incrementally. Whenever it is available this function reads that state
and never builds a per-step counts matrix:

* with no repetition anywhere in the batch, presence + frequency are
  ADDITIVE, so they arrive pre-summed as the accumulator's ``(B, V)``
  delta and the step is one subtraction — no penalty kernel at all;
* a batch that DOES ask for repetition takes the fused kernel, reading
  the accumulator's int32 occurrence buffer where the fallback reads a
  freshly scattered fp32 matrix. Everything else about that tail is
  identical, which is what makes the two bit-identical.

The counts path is the fallback for callers with no accumulator (on
CUDA: no engine pool hook, or a vocab/dtype the boot reserve was not
computed for). It allocates and scatters a ``(B, V+1)`` matrix per
call.

``tests/test_sampler_penalty_accumulator.py`` drives both arms through
this entry point and compares with ``torch.equal``.

Add per-request logit bias to specific token IDs.

Vectorized: builds a padded ``(B, max_keys)`` ids tensor + matching
values tensor, then one ``scatter_add_`` into a ``(B, W+1)`` buffer
delivers every per-request bias in a single launch. The pad column
is dropped before broadcasting back over K.

``window`` restricts the write to one vocab-parallel rank's slice. A
bias is a per-``(row, token)`` constant, so the restriction is exact by
inspection: keep the keys the window carries, index them by
``id - start``, and no other rank's keys can reach these columns.

Scale logits by 1/temperature per request (skip greedy rows).

Vectorized: one in-place ``logits.div_(temp.unsqueeze(1))`` after
replacing greedy temperatures with 1.0 to keep their rows
untouched. Works in-place on a freshly cloned tensor so an
upstream CUDAGraph output buffer is not modified.

Inform the processor that ``token`` was committed for ``request``.

Called by the engine after sampling, before the next forward.
Stateless processors implement this as a no-op; stateful ones
(xgrammar) advance their per-request automaton here.
Implementations doing CPU-bound work should remain sync-fast
(one matcher step per token); otherwise dispatch via
``asyncio.to_thread``.

Whether :meth:`apply_to_verify_slate` reads ``drafts_per_row``
for any of ``requests``.

The verify path keeps the draft tokens on the device and pulls them
to the host only for a processor that answers True here. The default
is True: a processor that has not said otherwise is handed the
tokens, never a slate it cannot constrain.

Apply this processor's constraint to the MTP verify slate.

The speculative draw does not run ``__call__``: it samples a
``(K+1, B, V)`` slate whose column ``col``, position ``p`` draws the
continuation of ``drafts_per_row[col][:p]``, so a constraint derived
from a row's token history must be rebuilt PER POSITION from those
drafts rather than applied once against the committed history. This
method is that per-position form, and it is the reason the chain
cannot simply be replayed on the slate.

Every member of the chain implements it: the verify path reads it
directly off each processor, so a processor that does not have one
is an error and never a row served unconstrained. A processor whose
constraint has NO per-position form REFUSES the rows that asked for
it — a request must not receive a constraint on the ``K=1`` draw and
silently lose it the moment speculation turns on.

Position-independent constraints (a fixed per-request vocabulary
mask) apply unchanged across all ``K+1`` positions.

``window`` restricts the write to one vocab-parallel rank's slice of
the vocabulary, for the shard-resident verify draw. Every constraint
here is per-token, so the restriction is exact: keep the entries the
window carries and index them by ``global_id - start``. A processor
that cannot restrict its own form must RAISE on a window rather than
apply a full-vocab mask to a shard's columns, which would ban the
wrong tokens outright. ``None`` is the whole vocabulary.

Whether :meth:`apply_to_verify_slate` would write for ``requests``.

The shard-resident verify draw has to know BEFORE it lays out the
slate whether anything will be applied to it, because reaching the
appliers at all costs it a reorder of the vocab-wide local logits.
Asking each processor is the only form of that question that cannot
go stale: a caller-side list of "which flags mean constrained" would
answer for the processors that existed when it was written, and a
new one would be silently skipped on that path alone.

Must agree with :meth:`apply_to_verify_slate` — a ``True`` that
writes nothing is only wasted work, but a ``False`` that would have
written is a request served unconstrained.

Rank-symmetric Gumbel-max noise for SPMD stochastic sampling.

Under ``ARBI_SPMD_TP`` every rank holds the SAME post-all-reduce logits
and the SAME replicated request state, so the ONLY thing that can break
cross-rank determinism is the random draw. A stateful per-rank
``torch.Generator`` diverges (each rank advances its own offset). This
module replaces it with a counter-based draw keyed on
``(base_seed, position)`` — a value every rank agrees on without a
broadcast — so every rank produces byte-identical noise and therefore
byte-identical samples.

The draw goes through a CPU ``torch.Generator`` (same device-independent
contract as ``spec_decode.rejection_sampler._seeded_uniforms``): the
uniforms are bit-identical on CPU and CUDA, which is what makes the CPU
gloo parity test a faithful proxy for the GPU TP path.

The base seed one request's stochastic draws are keyed on.

The request's explicit ``sampling.seed`` when set, else its
``request_id``: stable for the request's whole life, distinct across
requests, and identical on every rank (the id is replicated).

Per-row Philox seeds for a ``(k, B)`` slate, flat row ``j * B + b``.

Row ``b`` at slot ``j`` draws with
``derive_row_seed(request_base_seed(req_b), request_position(req_b) + j)``
— the same function of the same ``(seed, position)`` the SPMD driver
derives for its workers, so a single-process server and a
tensor-parallel one sample the same request identically.

Derive a positive 63-bit seed from ``(base_seed, position)``.

Counter-based: distinct ``(base_seed, position)`` pairs scramble to
decorrelated seeds. Rank-independent — every rank passes the same
``base_seed`` (the request seed) and ``position`` (the request's
committed token count), so every rank derives the same seed.

Build the ``(B, V)`` Exp(1) Gumbel-max noise, one row per request.

Each row ``i`` is drawn from a CPU generator seeded with
``derive_row_seed(base_seeds[i], positions[i])`` — so the noise is a
pure function of the rank-agreed ``(seed, position)`` and is identical
on every rank. ``device`` only governs where the returned tensor
lands; the draw itself is always on a CPU generator (device-
independent values), matching ``_seeded_uniforms``.

The noise is ``Exp(1)`` (the ``probs / Exp(1)`` Gumbel-max trick the
sampler uses), clamped strictly positive so the downstream divide is
finite.

Counter-based (Philox-keyed) rejection sampler for stochastic decode.

An OPT-IN alternative to the single-pass full-vocab Gumbel-max sampler
(:mod:`arbi_serve.sampler.gumbel_argmax_triton`). The Gumbel path scans
every one of the ``V ~= 152k`` vocab entries per row, generating a fresh
Philox draw AND a double-``log`` per entry. This sampler instead draws a
SMALL, bounded number of candidate tokens — one counter-based Philox
uniform PER CANDIDATE, keyed on ``(row_seed, draw_index)`` — maps each
uniform through the proposal categorical by inverse-CDF, and accepts the
first candidate that passes the top-k / top-p / min-p filter, skipping
the full-vocab noise scan. Reusing the same rank-agreed counter-based
Philox stream arbi already uses for the Gumbel draw keeps this
DETERMINISTIC.

Why this is an exact draw from the filtered categorical. Let ``q`` be the
proposal ``softmax(logits)`` over the FULL vocab and ``S`` the surviving
(kept) set after top-k / top-p / min-p. The target filtered categorical
is ``p_i = q_i / Q_S`` for ``i in S`` (``Q_S = sum_{j in S} q_j`` the
retained mass) and ``0`` otherwise. Rejection sampling with ``q`` as the
proposal and envelope ``M = 1 / Q_S`` accepts candidate ``c`` with
probability ``p_c / (M q_c) = [c in S]`` — i.e. draw ``c ~ q`` and accept
iff ``c in S``. The accepted candidate is then distributed EXACTLY as
``p`` (textbook rejection sampling). Acceptance probability per draw is
``Q_S``; expected draws ``1 / Q_S`` (e.g. ``~1.25`` at ``top_p = 0.8``),
so a small ``max_draws`` bound makes an exhausted budget astronomically
unlikely — and when it does exhaust we commit the highest-probability
survivor (``argmax`` over the masked logits), which is deterministic and
always in ``S``.

Determinism / rank-symmetry. The candidate uniforms are a PURE function
of ``(row_seed, draw_index)`` via Philox4x32-10 (the SplitMix64-derived
``(seed, position)`` row seed the SPMD driver already agrees on — see
:mod:`arbi_serve.sampler.rank_symmetric`). The proposal CDF and the keep
mask are pure functions of the rank-identical post-all-reduce logits and
the per-request knobs. So every TP rank draws the byte-identical token
with NO noise tensor exchanged, and the same ``(seed, position)`` decodes
the byte-identical token on every replay — the same invariants the Gumbel
path holds. The whole draw is plain ``torch`` (softmax / cumsum /
searchsorted / integer Philox), so the CPU path the unit tests exercise
is the SAME code that runs on device.

This is a DIFFERENT sampler from Gumbel-max — it does NOT reproduce the
Gumbel path's token bit-for-bit (a different but equally-valid draw from
the SAME filtered categorical). It is gated behind ``ARBI_SAMPLER`` and
is not the default; the Gumbel path is canonical.

32x32 -> (hi32, lo32) unsigned multiply, overflow-safe in int64.

``a`` is a uint32 python constant (a Philox multiplier); ``b`` is an
int64 tensor holding uint32 values. A direct ``a * b`` overflows
signed int64 when both operands approach ``2**32``; splitting into
16-bit halves keeps every partial ``< 2**34``, well inside int64, and
reproduces the full 64-bit product's high and low words. Bit-identical
on CPU and CUDA (pure integer arithmetic), which is what makes the
counter stream rank-symmetric and device-independent.

``(B, n_draws)`` fp32 24-bit uniforms in the open interval ``(0, 1)``.

Each entry ``[b, j]`` is one Philox4x32-10 draw keyed on
``row_seeds[b]`` with the counter set to ``(draw_index=j, 0,
draw_offset, 0)`` — so the uniform is a pure function of
``(row_seed, draw_index)`` and byte-identical on every rank / device.
``draw_offset`` namespaces this draw class on counter lane 2 away from
any other Philox consumer sharing the seed. Ten Feistel-style rounds
of uint32 arithmetic (wraps), matching the Triton Gumbel kernel's
:func:`_philox4x32_lane0`.

The uniform is built from the TOP 24 bits of the lane-0 output as
``(bits + 0.5) * 2^-24`` in fp32 — deliberately NOT the full 32-bit
fp64 form, because ``bits < 2^24`` makes this arithmetic reproducible
BIT-FOR-BIT inside a Triton fp32 kernel (no fp64 path). That is the
determinism anchor the fused kernel
(:mod:`arbi_serve.sampler._rejection_kernel`) matches exactly.

Counter-based rejection draw from the filtered categorical -> ``(B,)`` int64.

Proposal is ``softmax(proposal_logits)`` over the full vocab; the
target is that categorical restricted to ``keep_mask`` and
renormalized. Draws ``max_draws`` candidate tokens by inverse-CDF of
the proposal at the counter-based uniforms
:func:`philox_uniforms` produces, and returns the FIRST candidate that
lands in ``keep_mask`` (all candidates accepted when ``keep_mask`` is
``None`` — the no-filter default). A row that accepts none within the
budget commits ``fallback`` (or, when ``fallback`` is ``None``, the
highest-proposal survivor of ``keep_mask``) — a deterministic
in-support token.

Exactness: on the kept set the proposal and target agree up to the
per-row constant ``1 / Q_S``, so accepting the first in-support
candidate is an exact draw from the filtered categorical (see the
module docstring). Deterministic and rank-symmetric: uniforms depend
only on ``(row_seed, draw_index)``; the CDF and mask only on the
rank-identical logits and knobs.

Per-row survivor LOGIT threshold ``tau`` for the fused kernel.

The top_k / top_p / min_p kept set is a logit-threshold set (softmax is
monotone in the logit, so the nucleus / top-k / min-p survivors are
exactly the highest-logit tokens): ``keep = {c : logit_c >= tau}`` with
``tau`` the smallest surviving logit. Returns ``-inf`` per row when
``keep_mask`` is ``None`` (no filter — every token accepted).

Fused single-pass rejection draw on device -> ``(B,)`` int64.

Launches :func:`._rejection_kernel._rejection_sample_kernel` (one
program per row): online-softmax normalizer pass, then the in-kernel
rejection loop (counter-Philox candidate -> inverse-CDF -> threshold
accept, bounded, deterministic fallback). Byte-matches
:func:`rejection_sample_from_seeds` on the drawn token (the Philox
uniform is the same 24-bit fp32 value; the inverse-CDF crossing agrees
whenever the cumulative gap exceeds fp32 rounding — always, bar an
adversarial near-tie). ``proposal_logits`` is read-only.

Per-request sampling on ``(K, num_seqs, vocab_size)`` logits.

The pipeline operates on a 3-D logits tensor; ``K=1`` in the non-MTP
path. The leading dim costs zero at runtime and keeps the surface
MTP-ready (no ``if K > 1:`` branch anywhere).

Per-request sampling chain:
  1. Repetition / frequency / presence penalties (over the output-id
     history).
  2. ``logit_bias`` — raw additive shifts on specific token IDs.
  3. Structured-decoding processors (xgrammar) — mask disallowed
     tokens to ``-inf``.
  4. Greedy short-circuit (T=0).
  5. Temperature scale.
  6. top_k → top_p → min_p mask.
  7. Gumbel-max sample (categorical via the ``-log(-log(u))`` trick).

Vectorization. The chain runs as a small set of batched tensor ops over
the full ``(K, B, V)`` (or flattened ``(K*B, V)``) tensor — no Python
loop over requests in the steady state. The greedy fast path skips
penalty / bias / temperature / sampling entirely when every row in the
batch is greedy and has identity-valued knobs (``temperature <= 0``,
no penalties, no bias) — in that case sampling is one ``argmax`` over
the whole tensor and a single ``.cpu().tolist()`` host sync.

Module layout. The per-step metadata (:class:`_SamplingMetadata` and its
build/key helpers) lives in :mod:`arbi_serve.sampler.sampling_metadata`;
the batched logit masks and the sampler-private noise generator live in
:mod:`arbi_serve.sampler.sampling_masks`. Both are re-imported below so
every name remains importable as ``arbi_serve.sampler.sampler.<name>``.
``_fused_sample`` and ``_USE_COUNTER_GUMBEL_FILTERED`` deliberately stay
in THIS module: tests replace/mutate them as module attributes and
:class:`Sampler` must observe those mutations at call time.

Counter-based rejection draw over ``(K*B, V)`` logits -> ``(K, B)`` int64.

Proposal is ``softmax(flat)`` over the full vocab; the target is that
categorical restricted to the top_k/top_p/min_p survivors. For a
FILTERED slate the survivor mask is built with the SAME sort-free
masks the Gumbel path uses (Qrita top-k/top-p Triton on CUDA, the
vectorized masks on CPU) applied to a CLONE — ``flat`` stays pristine
as the proposal — and the draw accepts the first candidate that lands
in the mask (see :mod:`arbi_serve.sampler.rejection_sampler`). The
no-filter slate needs no mask: every candidate is accepted, so the
draw is a single inverse-CDF map of one counter-based uniform.

Determinism / rank-symmetry: the candidate uniforms depend only on
``(row_seed, draw_index)`` and the proposal CDF + survivor mask only on
the rank-identical logits and knobs, so every rank draws the
byte-identical token with no noise exchanged — the same contract the
Gumbel path holds.

Single batched stochastic sample over (K*B, V) logits.

Steps:

    1. Top-k + Top-p in **one fused Triton kernel launch** (vLLM's
       Qrita pivot algorithm) — replaces the two-launch
       ``_vectorized_top_k_mask`` + ``_vectorized_top_p_mask``
       chain when fp32 + CUDA.
    2. Apply per-row min-p mask on softmax-probs (logit-space
       threshold; one batched ``masked_fill_``).
    3. Softmax → Gumbel-max via the ``probs / Exp(1)`` trick → argmax.

Returns sampled token ids of shape ``(K, B)`` int64. Greedy rows
are NOT excluded here; the caller merges them in via
``torch.where`` against a parallel argmax pass (vLLM pattern).

Numerical contract. Greedy rows are short-circuited above the
fused path entirely — temperature=0 routes through the
``argmax``-only path that never enters this function. The
resulting token-id distribution is parity-tested via KL-divergence
in ``tests/test_sampler_triton_parity_gpu``.

Cudagraph-safety. The Gumbel-max noise draw uses a sampler-private
``torch.Generator`` (see :func:`_get_noise_generator`) instead of
the default CUDA generator. Default-generator ``exponential_``
raises ``RuntimeError: Offset increment outside graph capture
encountered unexpectedly`` on every stochastic completion when
cudagraphs are ON, because the engine's captured forwards register
the default generator with their CUDAGraphs and the post-capture
main-stream RNG call trips a state-mismatch check in
``CUDAGeneratorState::increase``.

Stateless sampling coordinator.

Operates on ``(K, num_seqs, vocab_size)`` logits. Non-MTP callers
pass ``K=1``; :class:`MtpDriver` passes ``K = draft_len + 1``
without changing any signature.

External logits processors (e.g. :class:`XGrammarLogitsProcessor`)
plug in via :meth:`set_logits_processors`. They run after penalty
/ bias / temperature and before greedy / top-k masking, so the
mask they apply (``-inf`` on disallowed tokens) survives every
downstream masking step.

Return ``_SamplingMetadata`` for the slate, reusing the cached
tensors when the slate's sampling key is unchanged.

Cache HIT requires the slate membership (request ids, in slate
order), the params of every row, the device, AND the vocab size
to all match the previous build. On any mismatch we rebuild and
refresh the cache. Bit-exact with calling
:func:`_build_sampling_metadata` directly: the returned tensors
are identical-valued (a HIT returns the *same* tensor objects the
prior build produced; downstream consumers never mutate the
metadata tensors in-place — they only read them and broadcast /
compare, mutating their own logits copies).

Drop the cached slate metadata. Returns 1 when there was one.

The tensors above are built with the forward arena routing, so the
cache holds a block in a pool whose physical can only be returned by
destroying it (:func:`~arbi_serve.engine.inprocess_capture.
release_idle_forward_arena`). This is the seam that gives that block
up, and it is sound there for a reason the cache key states: the key
carries the slate's request ids, and the release runs only with no
running and no waiting row, so no later slate can hit this entry.

The active chain, in the order :meth:`sample` runs it.

The speculative verify draw does not go through :meth:`sample`, so it
reads the chain here and runs each member's
:meth:`~arbi_serve.sampler.processors.LogitsProcessor.apply_to_verify_slate`
— the same members, in the same order. Every constraint a request is
granted on the ``K=1`` draw therefore reaches its speculative draw
too, or is refused there.

True ⇒ request needs nothing but ``argmax``.

Cheap CPU-side predicate: temperature == 0 (greedy), no
``top_k`` / ``top_p`` / ``min_p`` constraint, no penalty, no
``logit_bias``, no structured-output (``response_format`` /
grammar). When this returns True for the only row in a batch
and the sampler has no external processors attached, the
whole sampling chain collapses to ``logits.argmax(-1)`` —
which lets the B=1 greedy path skip the ``_SamplingMetadata``
build entirely.

Run the full sampling chain and return ``(K, B)`` int64 on device.

``noise`` (SPMD rank-symmetric path): when given, the stochastic
Gumbel draw uses this caller-supplied Exp(1) tensor instead of the
sampler's own generator, so every rank draws identical Gumbels.
Greedy rows ignore it. Passing ``noise`` keeps the SAME chain — it
is not a separate code path.

The pure-GPU kernel chain — no host sync. Keeping the host sync
out of this method lets :meth:`sample_step_keep_gpu` reuse the
chain without any host work, so the engine task can hand the
GPU tensor straight to the next step's persistent ``input_ids``
buffer via D2D.

Bit-identical contract with :meth:`sample`: every fast-path /
branch returns the exact same ``(K, B)`` int64 GPU tensor that
:meth:`sample` produces before its host sync.

Sample tokens per (k, request) row.

Returns a list of length ``num_seqs``; each entry is a list of
length ``K`` of sampled token IDs. In the non-MTP path
(``K=1``), the returned inner lists are length-1.

Hot-path strategy.

Trivial single-request greedy fast-path: when ``B == 1`` and
the only request is greedy with no penalties / bias / advanced
sampling / structured output, sampling collapses to a single
``argmax`` and one host sync — and the ``_SamplingMetadata``
build (five (B,) host-tensor allocations + H2D copies + a
Python loop over requests) is skipped entirely. This is the
dominant decode-step shape in low-QPS production.

Greedy fast path: when every row has ``temperature <= 0`` and
no penalties / bias / external processors, sampling collapses
to a single ``argmax(dim=-1)`` over the whole ``(K, B, V)``
tensor and one ``.cpu().tolist()``. No fp32 cast, no penalty
scatter, no mask passes — the dominant decode-step shape in
production.

Stochastic / mixed path: penalty / bias / temperature / mask
passes are each ONE batched tensor op over the full
``(K, B, V)`` (or flattened ``(K*B, V)``) tensor — no Python
loop over requests. Greedy rows in a mixed batch are merged in
via a ``torch.where`` against a parallel argmax pass (vLLM
pattern); this keeps the stochastic kernels launched once.

Sample one token per request from a 2-D ``(num_seqs, V)`` view.

The engine's non-MTP step calls this — it adds a leading K=1
dim, dispatches through :meth:`sample`, and unwraps. ``noise`` is
the rank-symmetric SPMD Gumbel noise (``(B, V)``, K=1);
``row_seeds`` is the cheaper rank-agreed per-row Philox seed for
the no-filter counter-based Gumbel-max fast path (``(B,)``, K=1).

Sample one token per request and return BOTH GPU + CPU views.

Returns ``(samples_gpu, sampled_ints)`` where ``samples_gpu`` is a
``(B,)`` int64 device tensor (a ``[0]`` slice of the internal
``(K=1, B)`` result) and ``sampled_ints`` is the per-request
Python int list (length B). The GPU tensor lets
``ModelRunner._build_batch`` write the next step's
``input_ids`` slot via D2D, skipping the int-list →
``torch.as_tensor`` → H2D path. The host int list is still
produced because ``Scheduler.commit`` / ``post_token`` /
detokenization need real Python ints.

The single ``.cpu().tolist()`` host sync is the only sync
point — and it is unavoidable as long as the run loop drives
detokenization synchronously after each step. The win here is
the elimination of the H2D round-trip ON THE NEXT STEP's
batch build.

No silent fallback: every sampling mode (trivial greedy,
all-greedy fast path, full stochastic chain) keeps the merged
result on GPU through :meth:`_sample_to_device` and produces
a normal ``(K, B)`` int64 tensor here. Logprobs aren't
produced by this path — callers requesting logprobs go
through :meth:`sample` which has its own host-copy budget.

Sample one token per request and return ONLY the GPU tensor.

Async-output path. Identical sampling to
:meth:`sample_step_keep_gpu` but issues NO ``.cpu().tolist()`` —
the host materialization is deferred one engine tick on the
run-loop side via the dedicated copy stream + event (see
:mod:`arbi_serve.runtime.async_output`). Returns the ``(B,)``
int64 device tensor; the caller stamps ``last_sampled_gpu`` from
it (so the next forward's input is a D2D copy) and issues the
non-blocking D2H itself.

Batched logit masks + the sampler-private noise generator.

Holds the vectorized top-k / top-p / min-p masks, the fused-Triton mask
helpers, and the per-device cudagraph-safe noise generator +
counter-Gumbel seed resolver. All names are re-exported from
:mod:`arbi_serve.sampler.sampler` so the public import surface (and the
``smod._vectorized_*`` / ``smod._get_noise_generator`` test hooks) are
unchanged.

Mask all but the top-k entries per row to ``-inf``.

Single batched ``topk`` over the whole (N, V) tensor at the max
K active in the batch; rows whose k=0 (or k>=V) are unmasked. The
mask is applied in-place.

Fused top-k + top-p mask via the vendored vLLM Triton kernel.

Replaces the two Python-batched ops
(``_vectorized_top_k_mask`` + ``_vectorized_top_p_mask``) with
ONE kernel launch, avoiding the global ``sort`` + ``softmax`` +
``cumsum`` + ``scatter`` kernel sequence the Python top-p path
triggers.

Tensors must be CUDA fp32. The caller (``_fused_sample``) is
responsible for that contract — see the cast at the call site.

Args:
    flat: ``(N=K*B, V)`` fp32 logits tensor; modified in-place.
    sm: per-step sampling metadata.
    K, B: original logits dimensions.

Build the per-row (k, p) tensors for the fused Triton kernel.

Broadcasts the (B,) metadata to (N=K*B,) and maps our "disabled"
sentinels to the kernel's convention (top_k==0 → vocab_size so the
kernel's ``k >= VOCAB_SIZE`` early-exit fires; top_p>=1 is already
disabled per row). Returns ``(None, None)`` for disabled axes.

Return the sampler-private noise generator for ``device``.

Lazily allocated on first request per device; seeded from the
default generator's current seed so determinism contracts the
engine sets via ``torch.manual_seed`` are preserved (the user-
visible "set seed → reproduce sample" guarantee). All subsequent
calls reuse the same generator instance — its Philox offset
advances normally on each ``exponential_(generator=g)`` call.

Threading. The engine's run loop is single-threaded per device,
so no lock is needed. The dict-set on first miss is safe under
the GIL.

Resolve the ``(n,)`` int64 per-row Philox seeds for the counter kernel.

``row_seeds`` given: validate the shape and use it verbatim — the
caller derived the seeds from each row's ``(seed, position)``
(:func:`~arbi_serve.sampler.rank_symmetric.request_row_seeds`), the
same keys on every rank and topology. ``row_seeds`` None (a caller
with no requests to key on): draw a fresh per-step seed per row from
the sampler-private (cudagraph-safe) generator.

Per-step sampling metadata: the ``(B,)`` knob tensors + fast-path predicates.

The :class:`_SamplingMetadata` dataclass and its build/key helpers marshal
each request's :class:`SamplingParams` into the device-side scalar tensors
the sampling chain consumes. Re-exported from
:mod:`arbi_serve.sampler.sampler` so the public import surface is unchanged.

Per-step pre-built sampling knobs (one allocation per ``sample()``).

Lives only for the duration of one ``Sampler.sample`` call. Holds
GPU-side per-request scalar tensors (temperature, top_k, top_p,
min_p) plus boolean fast-path predicates. The fields here are the
minimum needed to drive the stochastic path with zero Python loops
over requests.

An uninitialised host buffer whose H2D copy to ``device`` is truly async.

``torch.empty(n, dtype=...)`` returns PAGEABLE host memory, and the CUDA
runtime serializes a pageable H2D as ``cudaMemcpyAsync`` +
``cudaStreamSynchronize``. The ``non_blocking=True`` on the matching
``.to(device)`` is therefore silently ignored and the host blocks until
the copy stream drains. Allocating the staging buffer in PINNED host
memory makes that flag effective: byte-identical device values, no host
stall. The pinned block is owned by the CUDA caching host allocator,
which defers its reuse until the recorded copy event completes, so it is
safe even though the host tensor is a temporary.

Falls back to a pageable buffer for a non-CUDA device, where
``pin_memory`` is unavailable and the copy is a plain host memcpy.

Build an invalidation key for the per-step sampling metadata.

The key is a pure function of *exactly* the inputs
:func:`_build_sampling_metadata` reads — and nothing else — so two
slates with the same key are guaranteed to produce byte-identical
``_SamplingMetadata`` tensors. It captures, in slate order:

  * ``device`` and ``vocab_size`` — both feed the build (the
    ``top_k >= vocab_size`` clamp, the H2D target device);
  * per row: ``request_id`` AND every :class:`SamplingParams` field
    the build touches (temperature, top_k, top_p, min_p, the three
    penalties, and ``bool(logit_bias)``).

Why ``request_id`` *in slate order* matters. ``_SamplingMetadata``
tensors are indexed by **slate position**, not request id. An
admission, an eviction, or any reorder changes the position→request
mapping even if the multiset of params is unchanged. Folding
``request_id`` into the key at its slate index makes any such churn
flip the key, so the cache can never serve a stale row-to-position
mapping. This is deliberately conservative: identical params on a
*different* request still invalidate.

Note we intentionally do NOT consult ``logit_bias`` contents,
penalty histories, or grammar state here — those do not affect the
five scalar knob tensors or the six CPU bool predicates this
metadata holds (penalty/bias VALUES are read later, per step, by
the processor chain, not by the metadata build). Only the
``bool(logit_bias)`` / penalty-active flags live in the metadata,
and those are captured.

``True`` iff every request is greedy with default penalties /
bias / xgrammar.

The full sampling chain (``apply_penalties`` ->
``apply_logit_bias`` -> external processors -> ``apply_temperature``)
short-circuits row-by-row when each row's :class:`SamplingParams`
has the default values. But each processor still iterates the
``requests`` sequence, reads attributes, and walks ``logit_bias``
keys. At B=1 single-stream decode (the production hot path) every
one of those calls is a no-op that still costs Python overhead.

This predicate folds the four short-circuits into one flat loop so
:meth:`Sampler.sample` can skip the chain entirely. Mirrors
:attr:`SamplingParams.is_greedy` but additionally checks penalties,
bias, and grammar state (the only fields the processors touch).

Structured-output sampling parameters with mutual-exclusivity validation.

``StructuredOutputsParams`` is one container for every
structured-decoding mode (json_schema, regex, choice, grammar,
json_object, structural_tag) with **exactly one** mode active per
request, validated at construction time. The structured-decoding
backend (xgrammar, today) reads whichever field is set.

This module ships the param dataclass + validation only. Wiring
xgrammar against this dataclass is not implemented; the existing
:class:`arbi_serve.engine.request.ResponseFormat` drives xgrammar
through the existing path.

One field per structured-decoding mode; **exactly one** must be set.

Fields:
  * ``json_schema`` — JSON validating against this JSON Schema dict.
  * ``regex`` — output must match the regex.
  * ``choice`` — output must be one of these literal strings.
  * ``grammar`` — raw EBNF grammar.
  * ``json_object`` — free-form JSON (any well-formed JSON value).
    ``True`` enables, ``False``/``None`` disables. (boolean to keep
    the "exactly one set" semantics — ``True`` means "this mode is
    the active one".)
  * ``structural_tag`` — vendor tag (e.g. ``"tool_call"``) that the
    decoding backend translates into a built-in grammar.

Mutual exclusivity is enforced in ``__post_init__``: zero set
raises ``ValueError`` (caller should pass ``None`` instead);
two-or-more set raises ``ValueError`` with the field names.

Reasoning-block token accounting + the thinking-budget close guard.

A reasoning model spends tokens inside a ``<think>…</think>`` block before
its answer. This module owns BOTH halves of the engine's reasoning-block
bookkeeping, so the boundary is detected exactly once, in one place:

  * the COUNTING half — :func:`count_reasoning_token`, run per committed
    token by :func:`arbi_serve.engine.run_step.post_token`. It advances
    ``Request.reasoning_tokens`` (the boundary index the API surfaces as
    ``usage.completion_tokens_details.reasoning_tokens``, reported through
    :func:`reported_reasoning_tokens`), closes the block when the marker
    lands, and arms the budget when ``SamplingParams.thinking_token_budget``
    is exhausted;
  * the ENFORCEMENT half — :class:`ThinkingBudgetGuard`. For every row whose
    request is in the force-close state it masks all logits except the
    closing marker's next token id to ``-inf``, so the marker is really
    sampled — the token flows to the next step's input through the same path
    as any sampled token, keeping KV / positions consistent (a post-hoc
    overwrite of the sampled list would NOT reach the on-device input in the
    async-output pipeline). Multi-token markers are forced one id per step
    via ``Request.force_close_pos``, which the counting half advances as each
    forced id commits.

The marker is matched on TOKEN IDS, not text: ``Engine.think_end_ids`` is the
reasoning parser's ``think_end`` string encoded once at build
(:func:`resolve_think_end_ids`), so the per-token work stays a fixed handful
of integer compares with no detokenize, no string scan and no allocation on
the engine step. The text-side split
(:mod:`arbi_serve.engine.reasoning_parsing`) is the ROUTE's view of the same
boundary and stays independent — it consumes decoded text the engine loop
never has to wait for.

The guard's registration surface mirrors
:class:`~arbi_serve.sampler.omni_guard.OmniOutputGuard`
(``Sampler.set_logits_processors``): logits arrive ``(K, num_seqs, V)`` and
rows map 1:1 to ``requests``. Zero-cost pass-through for the common case —
no row forcing — mirroring the no-repeat-ngram blocker: one bool read per
row, then ``return logits`` unchanged.

Encode the reasoning parser's closing marker to token ids. Build-time.

The engine-side boundary detector. ``()`` for a model with no reasoning
surface (``parser is None``) — which turns BOTH halves off: nothing is
counted, nothing is forced, and ``reasoning_token_count`` stays ``None``
on every request. Called once per build (and per hotswap), never on the
step path.

``validate_markers`` has already proven this marker tokenizes and
survives a ``skip_special_tokens`` round-trip, so an empty encoding here
can only mean "no parser".

Assumes the model emits the marker as its own canonical encoding — true
for every registered family (Qwen3.x ``</think>`` is one dedicated id,
which is also what the chat template feeds back in). A checkpoint that
emitted some other tokenization of the same string would never close the
block, so the request would report ALL its tokens as reasoning — the
conservative direction (it never claims answer tokens were reasoning),
and visible in the ``done#`` log line's ``content=0``.

Does this request have a reasoning block to account for at all?

The ONE predicate behind both "count this token" and "report a count":
the model has a closing marker AND the RENDERED prompt left generation
inside an unclosed block (``SamplingParams.opens_in_reasoning``, read off
the prompt by ``ReasoningParser.prompt_opens_reasoning``). A raw
completion, a thinking-OFF chat turn, or a model with no reasoning
markers is False here and reports ``None`` — never a misleading ``0``.

Advance the reasoning state machine for one COMMITTED token.

``end_ids`` is the ``</think>`` marker used to DETECT a natural close;
``force_close_ids`` (defaults to ``end_ids``) is the sequence the budget
guard FORCES — a short transition + the marker — advanced one id per step.

Called from :func:`arbi_serve.engine.run_step.post_token` once per
committed token — including the token that FINISHES the request, so a
request truncated at ``max_tokens`` while still inside the block reports
``reasoning_token_count == output_token_count`` (every token was
reasoning) rather than a count that stops short of the truth.

Counts the closing marker itself, so ``reasoning_tokens`` is the boundary
INDEX: ``output_token_count − reasoning_tokens`` is exactly the
visible-answer length. Once the block closes, answer tokens are not
counted (the ``thinking_closed`` gate at the top makes the post-answer
steps a single bool read).

Cost: O(len(end_ids)) integer compares — one for the usual single-id
``</think>`` — plus one int increment. No allocation, no decode, no
device work.

The running request ``request_id`` names, by either of its two ids.

THE CALLER'S OWN ID WINS. ``client_request_id`` is the ``X-Request-ID`` the
caller supplied and can correlate; the engine's ``request_id`` is a
per-process counter that means nothing outside this process and is only
discoverable by reading the timeline. So a caller acting on their own
request is matched first, and the engine counter is the fallback for an
operator working from the swimlane -- where the bar under the cursor IS
the engine id.

That precedence also settles the only ambiguity: a client id that happens
to be all digits resolves to the client id, never to whichever engine
request currently holds that number.

Make one in-flight request stop reasoning and start answering.

The same transition the thinking budget performs, decided by an operator
instead of by a token count: arm ``force_close_thinking`` and the
:class:`ThinkingBudgetGuard` — already attached and already consulted for
every row on every step — masks the row to the closing marker on the next
forward pass. The model therefore TRANSITIONS to answering rather than
being truncated: the response stays well-formed, ``thinking_closed`` flips,
and ``stamp_first_answer`` fires, so the reasoning→answer boundary is
recorded exactly as a natural close would record it.

Refuses rather than no-ops. Three ways this cannot act, and each returns
the reason instead of a success a caller would believe:

  * the model has no reasoning surface (``think_end_ids`` empty) — the
    guard degrades to pass-through, so arming the flag would change
    nothing and report that it had;
  * the request is not running (matched by the caller's own
    ``X-Request-ID`` first, then by the engine's per-process counter --
    see :func:`_find_request`);
  * it is not inside an open thinking block, so there is nothing to close.

FIELD ORDER IS LOAD-BEARING. This runs on the caller's thread while the
engine loop reads both fields; writing ``force_close_pos`` first means the
guard can never observe ``force_close_thinking`` set against a stale
cursor, which would emit the tail of the marker rather than its head.

``FinishOut.reasoning_token_count`` for a finished request.

An int — the boundary index — for any request that entered a reasoning
block, INCLUDING one truncated at ``max_tokens`` before ``</think>`` (the
count then equals the output length: it was all reasoning). ``None``, and
therefore ``usage.completion_tokens_details`` omitted entirely, for a
request that never entered one.

The token sequence the guard FORCES at the budget: a short transition +
the ``</think>`` marker, so the hand-off to the answer reads smoothly.

Distinct from ``think_end_ids`` (which stays the DETECTION marker for a
natural close): forcing emits this whole sequence one id per step. Falls
back to the bare marker when there is no reasoning surface or the transition
is disabled. Resolved once at build, never on the step path.

The tool-call START marker (e.g. ``<tool_call>``) as token ids — treated
as an IMPLICIT reasoning close.

For MOST models a tool call is the end of the turn, so a call emitted while
still inside ``<think>`` means the reasoning is done (and the budget guard
must never force ``</think>`` into the middle of a call — the open vLLM
bug). Model-specific: a few checkpoints are trained to call tools MID-
reasoning; set ``ARBI_TOOL_CALL_ENDS_REASONING=0`` for those and this
returns ``()`` (disabled). ``()`` too when the model has no tool parser.
Resolved once at build, never on the step path.

The close-marker id to FORCE for this row THIS step, or ``None``.

The ONE force-close decision, shared by :class:`ThinkingBudgetGuard`'s
main-path ``__call__`` and its verify-slate form, so the ``K=1`` draw and
the speculative draw can never disagree about which rows are forcing.
``None`` unless the row is in the force-close state AND the marker is not
already fully emitted — then the id at ``force_close_pos`` (one id per
step for a multi-token marker). Duck-typed on ``force_close_thinking`` / ``force_close_pos`` so it
reads a real ``Request``, the SPMD mirror, or the verify shim identically.

Mask verify position ``p == 0`` to the forced close marker.

The verify-slate force-close body behind
:meth:`ThinkingBudgetGuard.apply_to_verify_slate`, kept a free function so
the mask can be driven against a slate directly in a test. No-op when the
model has no marker or no row is forcing (one bool read per column).

Under a vocab-parallel ``window`` every rank still bans its whole slice —
the ban is the whole vocabulary minus one token — and only the rank that
OWNS the forced id reopens it. The finite value lands exactly once across
the group, which is what makes the restriction exact rather than merely
consistent.

Logits processor forcing ``</think>`` for over-budget reasoning rows.

``think_end_ids`` is the model's closing marker as token ids (resolved
once at build from the reasoning parser's ``think_end`` string). Empty
means no reasoning surface — the guard is not attached in that case, but
an empty tuple degrades to an unconditional pass-through rather than
forcing a bogus id.

Force the next close-marker id at the IMMEDIATE verify position.

Only ``p == 0`` is constrained: it is the one position whose context
is the committed history the budget was measured against, so it is
where the marker must land. Positions ``p > 0`` are the
post-``</think>`` answer and stay free — the row's drafts at those
positions are then rejected against the forced marker and exactly
the marker is emitted.

Build an xgrammar *structural tag* that constrains tool-call generation.

This is the enforcement half of tool calling (the parsing half lives in
:mod:`arbi_serve.engine.tool_parsers`). Given the request's ``tools`` and
``tool_choice`` it builds an ``xgrammar.structural_tag.StructuralTag`` that
forces the model's output into the *exact* surface its
:class:`~arbi_serve.engine.tool_parsers.base.ToolParser` reads back — so
enforce and parse never drift.

The per-family wire format (the ``<function=>`` Qwen XML vs Hermes JSON body,
the trigger, the scaffolding) is owned by the active ``ToolParser``, resolved
from the model's chat template at engine build and registered here via
:func:`set_active_tool_parser`. This module only owns the family-agnostic
``tool_choice`` scaffolding:

  * ``"auto"`` / ``None`` → ``TriggeredTagsFormat`` on the family trigger: the
    model may emit plain text OR open a tool call; once opened it is forced
    valid + parseable.
  * ``"required"`` → ``TriggeredTagsFormat(at_least_one=True)``: same shape but
    at least one valid call; parallel calls + inter-call text allowed.
  * named → ``TagsWithSeparatorFormat`` restricted to that function with
    ``stop_after_first``: exactly one call to the named function.
  * ``"none"`` → no grammar (handled by the caller; never reaches here).

xgrammar is imported lazily so importing this module never needs it present.

Build the ``StructuralTag`` for ``(tools, tool_choice)``, or ``None``.

The tag set + trigger come from ``parser`` (defaulting to the active
model's parser), so the enforced surface matches that model's own tool-call
format. Returns ``None`` when there is nothing to enforce (no tools, or no
tool matched a named choice). Per-tool schema problems degrade to permissive.

The tag set is built in ONE call
(:meth:`~arbi_serve.engine.tool_parsers.base.ToolParser.build_tool_tags`)
rather than per tool, because how many tags a request needs is a
family property: Hermes / Qwen3-Coder put one call per marker span, so
each tool is its own alternative tag, while NemotronVoiceChat wraps a
JSON ARRAY of calls in one ``<TOOLCALL>`` span and therefore needs a
single tag covering every offered tool. The default implementation is
still one ``build_tool_body_tag`` per tool, so the two existing
families' tag sets are unchanged.

Fanned-out top_k/top_p/min_p mask for the small-batch decode row.

WHY THIS EXISTS. :func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
(the vendored Qrita pivot kernel) launches ``NUM_PROGRAMS = min(num_sm,
batch_size)``. Its cost is therefore set by the SLATE HEIGHT, not by the
work: at the ``B = 1`` the speculative drafter serves it is ONE CTA of 128
threads running a multi-pass ternary pivot search over the whole vocabulary
while 127 SMs idle. Measured live (RTX 4090, Qwen3.8-27B, ``V = 248320``,
``K = 4``): **276 us** per drafter slot (x4 = 1.11 ms/step) and **271 us**
for the ``(K+1) * B = 5``-row verify mask. The row it reads is 993 KB; at
~1 TB/s that is ~1 us.

The verify mask is on the BASELINE path — every request pays it, greedy or
stochastic, no flag — but it is the GATHERED verify tail, so a TP>1 config
running ``sharded_verify`` (default ON) does not reach it: that path never
builds a dense ``(rows, V)`` row to mask, it compacts to per-shard candidates
first. The 271 us is the TP1 cost, which is the configuration it was profiled
on. The drafter mask needs ``ARBI_TRUE_STOCHASTIC_DRAFT != 0``.

WHAT REPLACES IT. Two launches, both gridded ``(rows x vocab-chunk)`` so one
row occupies every SM:

  1. :func:`~arbi_serve.sampler._topc_kernel._topc_partials_kernel` — each
     program ``tl.topk``-selects the top-``capacity`` of its own contiguous
     vocab chunk. This is the only pass that reads all of ``V``.
  2. :func:`~arbi_serve.sampler._topc_kernel._topc_mask_kernel` — each
     program re-derives the row's global top-``capacity`` from the level-1
     partials (their union provably contains it), turns the filter chain
     into ONE scalar threshold, and applies ``keep iff logit >= T`` to its
     own chunk, writing only the dropped lanes.

THE NUCLEUS, AND THE THIRD LAUNCH. One row shape has no bound a candidate set
can supply on its own: ``top_p`` active with ``top_k`` DISABLED — the standard
OpenAI request shape, and every row of a checkpoint whose
``generation_config.json`` ships no ``top_k``. The nucleus renormalises over
the WHOLE vocabulary and its depth is set by the row's flatness, so it can run
past any capacity.

Two things make it representable anyway. The denominator is real: the level-1
pass also emits each chunk's softmax mass in its own max's units, which the
level-2 reduction rescales onto the row max, so the nucleus is taken against
the row's true normaliser and not against the candidates — for the price of
the ``exp`` on bytes the selection was already reading. And the row is served
only when the candidates PROVABLY contain the boundary, which the threshold
kernel decides per row. A row it cannot claim keeps its ``top_p`` in a
device-written residual and reaches
:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
UNMASKED; a served row's residual reads ``1.0``, the DISABLED encoding, so
that launch walks past it. The escalation is unconditional in SHAPE and
selective in EFFECT, which is what keeps the chain capturable: nothing about
any launch depends on a host-visible value.

That escalation also takes over ``min_p`` and the fused LSE whenever it is
built, because both must observe the FINAL kept set and an escalated row is
not masked until it runs.

The nucleus path runs at its own capacity (:data:`_NUCLEUS_CAPACITY`), which a
``top_k``-capped slate never pays: holding a served ``top_k`` and holding a
whole nucleus are different requirements, and the mask's cost is linear in the
capacity.

The threshold is recomputed redundantly by every masking program on purpose:
a grid-``(rows,)`` kernel to compute it once would reintroduce exactly the
single-CTA step this design exists to delete, and the programs run
concurrently so the redundancy is free in wall time.

WHY A THRESHOLD, NOT A KEPT-SET. top_k, top_p and min_p are all
"keep the high tail" rules, so each is a cut at a VALUE:

  * top_k → the ``k``-th largest logit;
  * top_p → the smallest value whose descending prefix reaches ``p``, over
    the top_k survivors, renormalised;
  * min_p → ``row_max + log(min_p)`` (the logit-space rewrite the eager
    chain and the pivot kernel both use).

``T`` is the largest of the three and the kept set is ``{logit >= T}``.
Placing each cut AT a real logit value rather than strictly between two is
what makes this agree with the pivot kernel on an exact tie: both keep the
WHOLE tie group at the cut (the Qrita mask pass admits the full duplicate
group, ``_topk_topp_kernel``'s sixth pass).

WHAT IS AND IS NOT BIT-IDENTICAL. The masked ``(rows, V)`` tensor is
bit-identical to the pivot kernel's — same kept set, and kept lanes are never
written so they carry their input bits. ``lse_out`` is NOT, and cannot be: the
pivot kernel accumulates ``sum(exp(logit - max))`` serially inside one
program, and a mask that fans the row across programs necessarily groups that
sum differently. Measured within a few fp32 ULPs (pinned at 8 in
``tests/_topc_mask_interpret_harness.py``; 0 across the whole served grid at
``V = 248320`` on a 4090). That is the SAME drift class the fused-LSE epilogue
already ships with against a separate ``torch.logsumexp``, and its only
consumer effect is that a verify accept can flip at an exact-tie boundary —
the output distribution is unchanged (Leviathan-2023 Thm 1 holds for the
target law; the accept test is comparing the same two real numbers).

EXACTNESS OF THE NUCLEUS BOUNDARY. A nucleus boundary is where a cumulative
mass crosses ``p``, and two exact statements of that rule still disagree about
the crossing when they sum the same terms in different orders and the boundary
token's own mass is comparable to the summation error. A flipped boundary
token is a different sampled distribution, not a rounding detail, so the row
is served only when the crossing clears :data:`_NUCLEUS_MARGIN` of relative
slack on BOTH sides — no reordering smaller than that can move the cut — and
handed to the pivot kernel otherwise. Serving fewer rows is the safe
direction; serving one whose boundary is in doubt is not.

EXACTNESS OF THE SELECTION. It is exact, not heuristic: a chunk can
contribute at most ``capacity`` members to the global top-``capacity``, so
the top-``capacity`` of the union of the per-chunk top-``capacity`` sets IS
the global one. A chunk too wide for one tile is folded by a running merge
inside the program, which is the same argument applied again.

The one boundary this cannot see is a tie group at the ``k``-th logit long
enough to run PAST ``capacity`` entries: top_p's renormaliser would then sum
a truncated top_k set and close the nucleus early. It needs
``capacity - k + 1`` bit-identical fp32 logits at the cut, which an fp32
lm_head over a 248k vocabulary does not produce — but "does not happen" is a
claim, so the behaviour is recorded rather than assumed away
(``test_a_tie_group_wider_than_the_capacity_is_a_known_boundary``). The top_k
cut itself is unaffected: it is a value, so the dense pass keeps every tied
entry wherever it sits in the vocabulary.

CAPTURE SAFETY. ``capacity``, ``split``, ``chunk`` and every block size are
launch constants and the two grids are static in ``(rows, V)``. There is no
host readback, no data-dependent shape and no branch on a device value, so
the pair records into a cudagraph and replays against whatever ``copy_()``
last wrote into the parameter buffers. Because ``capacity`` is BAKED at
record time, a slate it cannot represent must be kept off the captured
chain by the routing layer — :func:`topc_mask_eligible` is the single
predicate both the router and the live dispatch ask. Whether the escalation
was RECORDED is the same class of fact and rides the same predicate
(``nucleus=``): a launch cannot be added by ``copy_()`` either.

Largest power of two ``<= value`` (``0`` for anything non-positive).

The level-2 selection is a ``tl.topk``, which takes only power-of-two
widths. Rounding DOWN rather than up matters: rounding up would quietly
serve slates the operator's number excludes, so the flag would be lying
about what it baked.

``None`` if the slate is representable, else WHY it is not.

The single decision :func:`topc_mask_eligible` is defined as, split out
so a refusal carries a machine-readable reason to the flag-truth counter
instead of collapsing to a bare ``False``. A bare ``False`` makes a
dormant fast path indistinguishable from a fast path the traffic keeps
stepping around, which is the exact question a fallback count exists to
answer.

``top_ks`` / ``top_ps`` are the encoded per-row values
(:meth:`~arbi_serve.spec_decode.rejection_sampler_ops.DrafterSamplingTensors.encode_params_host`
conventions: ``top_k == vocab_size`` means DISABLED, ``top_p >= 1.0``
means disabled).

A row is representable when its top_k caps the kept set inside the
capacity, or when nothing needs a rank/nucleus search at all:

  * ``0 < top_k <= capacity`` — the whole chain (top_k, then top_p over
    its survivors, then min_p) lives inside the top-``capacity`` set;
  * ``top_k`` disabled AND ``top_p`` disabled — nothing but min_p can
    cut, and min_p needs only the row max, which the partials carry.

``top_k`` disabled with ``top_p`` ACTIVE needs the third launch
(``nucleus``): the nucleus renormalises over the whole vocabulary and can
extend past the capacity, so the mask can only serve the rows whose
boundary it proves lies inside the candidates and must hand the rest back.
A caller that can run that escalation passes ``nucleus=True``; a caller
replaying a chain that did not RECORD it passes ``False``, because a
launch cannot be added by ``copy_()``.

``capacity <= 0`` disables the path outright (the rollback).

True iff every row of a slate can be masked from a top-``capacity`` set.

The one predicate the routing layer and the live dispatch both ask; see
:func:`topc_mask_refusal` for the rule and the per-shape reasons. Kept as
a thin wrapper rather than a second copy of the rule so the two can never
disagree about a slate.

Capacity the nucleus path runs at for a row of ``vocab_size``.

A nucleus needs a candidate set deep enough to CONTAIN it, which is a
different requirement from holding a served ``top_k`` — so the path widens
past the operator's capacity, and only for the slates that present the
shape. The widening stops below the vocabulary: a candidate set that IS
the row selects nothing and leaves no tail to cut, so a narrow vocabulary
keeps the operator's capacity and escalates more rows instead.

``(split, chunk, block_size, part_block, n_tiles)`` for a row width.

``split`` — how far one row fans — is derived from the level-2 budget
rather than chosen independently, because those two are the same number:
every masking program re-selects over a ``split * capacity``-wide tile, so
letting ``split`` grow to fill the device is exactly how the level-2 cost
runs away. Bounding the product bounds both, at every vocabulary size.

``block_size`` is then capped separately and a wider chunk is folded by a
running merge, so neither tile scales with ``V``. Both are powers of two
(``tl.topk`` takes no other width) and both are at least ``capacity``.

Both bounds are arguments rather than constants so the geometry can be
TESTED and MEASURED across them. In particular the running-merge branch is
unreachable at any served width — it needs a chunk wider than ``max_tile``
— so without a way to shrink the cap it would ship exercised only by
arithmetic.

Mask ``logits`` to the top_k → top_p → min_p kept set; returns it.

Drop-in for :func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
on a slate :func:`topc_mask_eligible` accepts (the caller MUST have
checked). Masked positions get ``-inf``; kept positions are untouched,
so they keep the exact fp32 bits they came in with.

``lse_out`` reproduces the pivot kernel's fused LSE epilogue — the
masked row's ``logsumexp``, which the fused MTP verify tail consumes in
place of a separate ``torch.logsumexp`` launch. It is finished on the
host from per-chunk partial sums, so it differs from the pivot kernel's
single-program accumulation by fp32 reduction order only — the same
drift class that path already ships with.

``nucleus`` adds the third launch that makes a ``top_k``-disabled row with
an active ``top_p`` representable: the mask kernel serves the rows whose
nucleus it can bound exactly and hands the rest to the pivot kernel
through a device-written residual ``top_p``, which reads ``1.0`` —
DISABLED — for every row already served. The escalation is therefore
unconditional in SHAPE and selective in EFFECT, which is what a captured
chain needs: nothing about the launch depends on a host-visible value.
It also takes over ``min_p`` and the LSE, because both must observe the
FINAL kept set and an escalated row is not masked until it runs.

Pure-torch twin of :func:`apply_topc_mask_`'s threshold — ``(rows,)``.

The device-free statement of the rule the kernel implements, so the CPU
suite can pin the SEMANTICS (which value each filter cuts at, and that
the cuts compose by ``max``) without a GPU, leaving the GPU suite to pin
the one thing only it can: agreement with the kernel this replaces.
Deliberately written against ``torch.topk`` rather than by reusing the
kernel's host helpers — a reference that shares the implementation
proves nothing.

Under ``nucleus`` a ``top_k``-disabled row renormalises over the whole
row, and a row whose boundary the candidates cannot pin down with
``margin`` to spare reads ``-inf``: not "keep everything", but "this row
is the escalation's, and the mask must leave it alone".

Combined Top-K + Top-P Triton kernel — vendored from vLLM.

Source: ``vllm/v1/sample/ops/topk_topp_triton.py`` (vLLM project, Apache 2.0).

Vendoring rationale. The Python-batched
sampling chain in ``arbi_serve.sampler.sampler`` runs one PyTorch op
per stage: ``torch.topk`` for top-k threshold, ``sort`` + ``cumsum`` +
``scatter`` for top-p, ``logits.max(...)`` for min-p, ``softmax`` +
``argmax`` for the Gumbel-max sample — several kernel launches per
step of pure launch overhead. vLLM fused
top-k + top-p into one Triton kernel using the **Qrita pivot
algorithm** (Park et al., arXiv:2602.01518): one kernel launch +
one buffer scratch, no global sort. This module brings that kernel
into ``arbi-serve`` so the Python-batched chain collapses to a
single launch for the top-k / top-p mask. Min-p stays in PyTorch
(one batched ``masked_fill_`` using the
``logit_i >= log(mp) + max_logit`` rewrite).

Local adaptations vs. the vLLM source:

  * Imports rewritten to use top-level ``triton`` / ``triton.language``
    instead of ``vllm.triton_utils`` (we don't have that wrapper).
  * ``next_power_of_2`` lifted from ``vllm.utils.math_utils`` (3 LOC).
  * ``num_compute_units`` replaced with
    ``torch.cuda.get_device_properties(...).multi_processor_count``
    (one call site).
  * ``SAMPLE`` epilogue (arbi-serve): optional fused softmax → Gumbel-max
    → argmax tail so the masked logits never round-trip to Python
    (:func:`sample_top_k_top_p_triton`).
  * ``MINP_ENABLED`` / ``WRITE_LSE`` epilogue (arbi-serve): optional
    per-row min-p mask (the ``logit >= max_logit + log(min_p)`` rewrite —
    the same form :func:`arbi_serve.sampler.sampler._vectorized_min_p_mask`
    uses) and/or the masked-row ``logsumexp`` written to ``OUT_LSE`` in
    the SAME launch. This is what lets the fused MTP rejection sampler
    (``ARBI_MTP_FUSED_REJECTION``) serve a filtered slate without adding
    a launch: the mask kernel REPLACES the ``torch.logsumexp`` launch the
    no-filter fused path issues, keeping the launch count identical.

DIVERGED FROM This file
used to promise that "the pivot-search body is byte-for-byte identical to the
upstream version", so the mask math stayed cherry-pick-able. That is no
longer true and must not be assumed by a future sync.

The upstream ``top_p`` search is a ternary search over the CONTINUOUS
probability range, capped at 18 iterations with a ``1e-9`` convergence floor
and an interval-MIDPOINT fallback when neither probe satisfies its acceptance
test in budget. Ternary search narrows by ~1/3 a step, so 18 steps from
``[0, 1]`` bottoms out at ~``1.4e-9`` — the floor itself. Two adjacent
probabilities closer than that cannot be separated, and the midpoint lands on
whichever side of the true boundary the interval happened to straddle. The
served nucleus was therefore not the one ``top_p`` asks for: measured against
a ``torch.sort`` reference it comes out too small OR too large, both observed,
on rows whose boundary sits inside a plateau of near-identical logits — which
is what a real ``lm_head`` produces over a 248k vocabulary. On well-separated
rows it is exact, which is why this survived.

It is replaced by :func:`~arbi_serve.sampler._topp_exact._nucleus_pivot`,
which bisects the IEEE754 ORDERED-INTEGER KEY in LOGIT space: representable
floats are a contiguous integer range, so bisecting them terminates on an
exact value with no tolerance constant anywhere, and staying in logit space
avoids a ``log(p * Z) + max`` round-trip that would reintroduce a ULP of slop
at exactly the boundary being fixed. It also runs FEWER passes than what it
replaces — one per iteration instead of two, and no probability buffer at
all.

The top_k pivot search is UNCHANGED and still tracks upstream: it was
measured exact on every case tried, including plateaus down to zero logit
spacing (``tests/test_topp_nucleus_exact.py``, and the ``top_k`` sweep in its
history). The arbi-serve epilogues remain appended, constexpr-gated blocks
whose disabled specialization compiles out entirely.

A future vLLM sync must therefore MERGE rather than overwrite: taking the
upstream search body wholesale reintroduces the defect silently, and the
regression test that catches it
(``tests/test_topp_nucleus_exact.py``) needs rows with a plateau at the
boundary to fail at all.

Engine boot hook: route sampler scratch/table allocations into a named pool.

Drops the existing caches so the next sampler call rebuilds them INSIDE
the pool; without that the process would keep serving the unpooled tensors
it had already cached and the ledger row would read zero while the bytes
stayed unbudgeted.

Passing ``None`` uninstalls (used by tests and by member teardown).

Materialise every mask-kernel buffer the serving path will need.

Two things make this a boot step rather than a lazy one.

The pivot scratch is cached per ``(device, row width)`` and reallocated
only when a WIDER row count arrives, so the widest slate must run first or
the widening lands later. And the cache is keyed per LOGITS width: the
full-vocab mask and the sharded candidate-width mask
(``n_cand * tp_size``) are separate entries, so covering one leaves the
other to allocate on first use.

Both matter because the pool is capped at the bytes it holds when the boot
freeze snapshots it. A later entry needs a fresh segment mapped into that
pool, and the cap denies the map however much device memory is free — the
itself wants.

Returns the bytes allocated.

Drop the process-global sampler scratch/table device tensors.

These caches (:data:`_TRITON_BUFFER_CACHE` — the per-(device, vocab) fp32
pivot-search scratch; :data:`_TRITON_TABLE_CACHE` — the CDF→sigma lookup
tables) hold DEVICE tensors first allocated by whichever member happened to
run the sampler first. They are process-global, NOT per-member, so a
stable-VA member swap does NOT rebuild them — the incoming member reuses
the OUTGOING member's tensors. When the outgoing member is parked its pool
physical is unmapped/evicted; the retained cache tensors then reference
freed/unmapped VRAM and the next ``_topk_topp_kernel`` launch reads
unbacked memory → ``cudaErrorIllegalAddress``. Clearing on member reset
forces the incoming member to re-allocate them in ITS OWN live context
on the next sampler call.

Same class as :func:`~arbi_serve.engine.engine_model_state._EngineModelStateMixin.
_reset_workspace_pool_state` (the FLA / cuBLAS workspace globals). Returns
the number of cache entries dropped (for the reset log). Safe to call any
time — the caches lazy-rebuild on the next sampler call.

Launch the fused top-k/top-p (+ optional epilogue) kernel once.

Shared by :func:`apply_top_k_top_p_triton` (mask, optionally + min-p
mask and/or masked-row logsumexp) and :func:`sample_top_k_top_p_triton`
(mask + in-kernel Gumbel-max). ``min_p`` / ``lse_out`` and the SAMPLE
tail are mutually exclusive (no caller combines them; the assert keeps
an accidental combination loud instead of silently mis-normalized).

Apply combined top-k and top-p masking with one Triton kernel launch.

Top-k is applied first (by logit value), then top-p is applied to
the surviving logits' softmax probabilities. ``logits`` is modified
in-place; the same tensor is returned for chaining.

Args:
    logits: ``(batch_size, vocab_size)`` fp32 OR bf16/fp16 tensor;
        modified in place. Must be CUDA-resident and contiguous.
        bf16/fp16 logits are upcast to fp32 IN-REGISTER per tile (the
        softmax max/sum accumulators stay fp32); the in-place mask
        write-back downcasts the kept (already-bf16) logits losslessly,
        so the masked output is bit-identical to the fp32-input path
        while halving the HBM read bytes.
    k: ``(batch_size,)`` int tensor of per-row top-k caps, or
        ``None`` to disable top-k. Rows whose ``k >= vocab_size``
        are no-ops in the kernel.
    p: ``(batch_size,)`` fp32 tensor of per-row top-p (nucleus)
        probabilities in ``[0, 1]``, or ``None`` to disable top-p.
        Rows whose ``p >= 1.0`` are no-ops in the kernel.
    mask_value: Value to write at masked-out positions. Default
        ``-inf`` (so a downstream softmax zeroes them). Must stay
        ``-inf`` when ``min_p`` / ``lse_out`` are supplied.
    min_p: ``(batch_size,)`` fp32 tensor of per-row min-p thresholds
        (``0.0`` = disabled per row), or ``None`` to disable. Applied
        AFTER top-k/top-p in the same launch via the logit-space
        rewrite ``keep iff logit >= max_logit + log(min_p)`` — the
        identical form :func:`~arbi_serve.sampler.sampler._vectorized_min_p_mask`
        uses, so the kept set matches the eager chain's.
    lse_out: optional ``(batch_size,)`` fp32 tensor; when supplied the
        kernel ALSO writes each row's masked-logits ``logsumexp``
        (fp32 ``log(sum(exp(l - max))) + max`` over the post-mask
        row) — the softmax normalizer, produced in the same launch so
        a fused consumer (the MTP rejection sampler) needs no separate
        ``torch.logsumexp`` launch.

Returns:
    The input ``logits`` tensor (modified in-place).

Fused top-k/top-p mask + Gumbel-max sample in ONE kernel launch.

Masks ``logits`` exactly as :func:`apply_top_k_top_p_triton` (in
place), then folds the stochastic tail —
``softmax`` -> ``probs / noise`` -> ``argmax`` — into the same
kernel so the masked logits never round-trip to Python. Returns the
sampled token index per row.

The result is op-for-op identical to the unfused chain
``argmax(softmax(masked_logits) / noise)`` (same fp32 softmax
denominator, same two divides, same lowest-index argmax tie-break).

Args:
    logits: ``(batch_size, vocab_size)`` fp32 OR bf16/fp16 CUDA
        tensor; masked in place (so it is NOT reusable as raw logits
        afterwards). bf16/fp16 is read directly and upcast to fp32
        in-register (half the HBM read); the result is bit-identical
        to passing the same logits pre-cast to fp32, because a bf16
        value widens to fp32 exactly and the masked write-back
        round-trips losslessly.
    k: per-row top-k caps (``None`` disables); see
        :func:`apply_top_k_top_p_triton`.
    p: per-row top-p (``None`` disables).
    noise: ``(batch_size, vocab_size)`` fp32 Exp(1) Gumbel-max draw
        — supplied by the caller so every TP rank draws the same
        token (rank-symmetric SPMD). Consumed read-only.

Returns:
    ``(batch_size,)`` int64 sampled token indices.

The global-vocab range a rank-local logits slice covers.

A vocab-parallel rank holds a contiguous slice of the vocabulary, so every
per-token logits transform the K=1 sampler applies to a full row — a
``logit_bias`` offset, the three penalties, a grammar mask — has an exact
restriction to that slice: keep the entries whose GLOBAL id lands inside
the window, drop the rest, and index by ``global_id - start``.

This exists so those transforms take the window as a PARAMETER rather than
being reimplemented shard-side. A second copy of penalty or bias arithmetic
would be a second thing to keep in step with the K=1 chain, and the two
would drift silently — the sharded draw would honour a knob the gathered
draw had since changed the meaning of, with nothing comparing them.

``None`` everywhere a window is accepted means the full vocabulary, and the
full-vocab case must stay byte-identical to the un-windowed code it
replaced: the shard-resident path is an optimisation, never a change to
what the K=1 draw would have produced.

Global ids ``[start, start + width)``, out of a ``vocab_size`` vocab.

``width`` is the number of REAL vocab columns the local logits tensor
carries, which is not always its last dimension: a short final shard is
zero-padded up to the uniform per-rank width, and those pad columns are
not tokens. The caller passes the tensor already trimmed to ``width``.

Restrict a full-vocab ``(..., vocab_size)`` tensor to the window.

A VIEW, not a copy: the shard-resident path's whole point is not to
materialise anything vocab-shaped per step, and a windowed view of
a resident buffer (the penalty accumulator's counts, its additive
entry) costs nothing. Row stride survives, so a Triton kernel that
takes ``stride(0)`` reads it directly.

Keyed-noise text watermark for the counter-based Gumbel-max sampler.

The Aaronson-family watermark (the scheme SynthID-Text descends from)
changes ONLY the source of the sampler's randomness: instead of a fresh
per-step Philox seed, the stochastic draw for each row is keyed on
``PRF(watermark_key, last-H sampled tokens)``. The drawn token is still
``argmax_v(logit_v + Gumbel(seed, v))`` — an exact categorical draw from
``softmax(logits)`` for any fixed seed source — so the per-token sampling
distribution, support, temperature/top-k/top-p semantics, and greedy
behaviour are all unchanged. What changes is that the choice becomes a
deterministic function of ``(key, context)``, which a key-holder can
verify after the fact (see :func:`detect_tokens`).

Context state lives ON DEVICE. Under async output the host does not yet
know the previous sampled token when the next step's seeds are needed, so
the last-``H``-token window per request is kept in a pooled ``(rows, H)``
device ring advanced from the sampler's own output tensor each step (a
D2D shift — no host sync). A row is (re)built from the host-side
``prompt_token_ids + output_token_ids`` tail only when the host state is
provably current: at a request's first sample (nothing sampled yet) and
after a preemption/recompute prompt rewrite (detected by a changed prompt
length, exactly like :mod:`arbi_serve.sampler.penalty_history`).

Repeated-context masking (the SynthID rule): a context window that has
already keyed a draw for this request must not key another one — the same
seed would force the same token and can lock a generation into a cycle.
Each row carries a device-side Bloom filter of the context hashes it has
used; a hit falls that step back to the normal (non-watermark) seed. A
Bloom false positive merely skips the watermark on one extra step (the
statistic just loses one token of signal); a false negative is impossible,
so a genuinely repeated context can never re-key a draw.

Detection needs no model and no GPU: recompute each position's seed from
the preceding ``H`` tokens with the same hash, evaluate the SAME
Philox4x32-10 uniform the sampling kernel generated at the chosen token
(:func:`philox_uniform_host` mirrors ``_philox_uniform`` bit-for-bit),
and test the uniforms for the upward bias the argmax leaves. Under the
null (text not produced with this key) each scored uniform is i.i.d.
``U(0,1)`` and the summed ``-log(1-u)`` statistic is ``Gamma(n, 1)``.

Scope: the watermark keys the K=1 counter-Gumbel decode paths (the
canonical ``ARBI_SAMPLER=gumbel`` backend) AND the MTP spec-decode path,
where keyed determinism collapses verification to EXACT-MATCH against
the target's keyed draw (see
:meth:`~arbi_serve.spec_decode.rejection_sampler.GraphSafeRejectionSampler.
_sample_batched_keyed_exact`) — full-density watermark, emitted stream
bit-identical to non-speculative keyed sampling. The rejection sampler
backend (``ARBI_SAMPLER=rejection``, inverse-CDF) and the TP>1
sharded-verify kernel consume randomness differently and are refused at
boot rather than silently left unwatermarked.

Unsigned 64-bit per-request salt for the unkeyed fallback draws.

The server's watermark key is folded in. Without it the salt would be a
function of the request id and a published constant alone, so anyone
could reproduce the fallback noise for a position, compare it against
the emitted token, and learn which positions were drawn UNKEYED — a
map of the gaps in the watermark, which is what stripping or forging
one needs. The key makes the fallback unpredictable to everyone who
does not already hold it, and the drafter and the verify read the same
cache row, so their coupling is unaffected.

Positive 63-bit seed for an UNKEYED draw at absolute token ``position``.

Used wherever a draw must not be keyed on its context: a Bloom-masked
depth (context window already keyed once) and every depth of an
opted-out row. It depends only on the row's salt and the position of
the token being drawn, both of which the drafter that proposes the
token and the verify that judges it read from the same cache state, so
the two draw with the same noise and their argmaxes agree whenever
their distributions do. It never repeats within a request (positions
are unique) and never keys on the context.

True iff ``raw`` is a watermark key an operator actually set.

``ARBI_WATERMARK_KEY=`` (empty) is how compose spells "leave this off",
but the generic ``str | None`` parser deliberately PRESERVES the
unset-vs-empty distinction — ``ARBI_FOO=`` is a real way to override an
earlier export back to empty — so the flag lands as ``""``, not ``None``.
Every ``is not None`` test on it then answers "configured" for a server
that has no key, and the report lies in the direction of MORE configured
than it is: the boot line masks a nonexistent secret as ``<set, masked>``
and per-request traces mark unwatermarked requests as watermarked.

One predicate, so the two readers cannot disagree about what "set" means.

True iff the configured key is the publicly disclosed demo key.

Both write-doors already refuse it — the boot guard
(``engine.build_phases_load``) and the live override coercion
(``config_overrides._as_watermark_key``) — so a running server should
never answer True. That is precisely why a status surface ASKS rather
than assumes: a key that reached the flag by any other route would
otherwise be presented as a private mark, and a mark anyone can
reproduce or strip proves nothing about who wrote the text.

A fresh 64-bit key from the system CSPRNG, as a parseable hex string.

``secrets`` and not ``random``: this is the secret a detector's whole
claim rests on, so it may not come from a PRNG seeded by anything
reproducible. Two values are redrawn — the publicly disclosed demo key
(which both write-doors already refuse) and zero (a key an operator
could reach by typo, and one that reads as "unset" to every truthiness
test in the tree).

Whether this server is watermarking — the ENABLE question, not the secret.

The two have been one thing in practice: ``watermark_enabled`` defaults
True and does nothing without a key, so what actually turned the feature
on was the presence of a secret. That is why a boot could sit in the
state the directive forbids — notionally on, no key — and why an
unrelated commit could re-parent two warm blocks under
``if runtime_flags().watermark_key:`` with nothing noticing (#2117).

Separated here into three independent readings:

* ``watermark_enabled`` — the opt-out. False means this server samples
  unkeyed whatever else is configured, and no key is needed or made.
* ``watermark_required`` — the fleet policy "every generated token is
  watermarked". It is an ENABLE, and it does not carry a secret.
* a configured key — an operator handing over a secret is also a request
  to use it, which is the behaviour every existing deployment has.

A server that is ON always has a key, because
:func:`effective_watermark_key` makes one when it is on and none was
supplied. There is no fourth state.

The key in force and where it came from — never ``None`` when ON.

Returns ``(key_string_or_None, source)``. ``source`` is one of
:data:`KEY_SOURCE_ENVIRONMENT`, :data:`KEY_SOURCE_GENERATED`,
:data:`KEY_SOURCE_OFF`. Every reader that decides behaviour off "is
there a key" goes through here rather than reading the flag, so the
generated key is not a second truth the tree has to be kept in step
with.

A generated key is a real key with one property an operator has to be
told about: it exists only for this process, so text marked with it
cannot be scored by any other process, and the mark does not survive a
restart. :func:`report_watermark_resolution` is where a boot says which
of the two it has.

Whether keyed watermarking is ACTUALLY applied, and why not when it is not.

The three conditions :func:`get_watermark_cache` gates on, answered off
the flags alone so a status reader needs no device and no engine. The
gate and the key are INDEPENDENT (``watermark_enabled`` says whether to
apply the mark, ``watermark_key`` says whether there is one to apply), so
neither flag read alone is the outcome: the default boot has the gate on
and no key, and reporting that gate as the state tells an operator text
is marked while every generation leaves unmarked.

Returns ``(active, inactive_reason)`` — the reason is a plain sentence
naming what to change, and is ``None`` exactly when ``active``. It never
contains the key.

:func:`watermark_state` against the key actually in force.

A status surface that reads the raw flag reports "no key" for a server
that is marking every token with a generated one, which is the exact
inversion this module's own docstrings warn about in the other
direction.

Resolve the key ONCE at boot, name its source, and count it.

Fires ``watermark_key_from_environment`` when the operator supplied the
secret, and refuses it — with the source as the reason — when the key
was generated here or when the server is not watermarking at all. The
counter therefore separates three states that a boolean could not: an
operator-supplied key, a key this process made, and a resolver that
never ran (fired 0, refused 0), which is a different and earlier
failure than either.

Returns ``(key, source)``. The key is never logged: the status line
carries :func:`mask_watermark_key`'s two-hex-digit form, which is what
every other watermark status surface in the tree already prints.

Pooled device state: ``(rows, H)`` context ring + per-row Bloom.

One instance per device (see :func:`get_watermark_cache`). Row
lifecycle mirrors :class:`~arbi_serve.sampler.penalty_history.
PenaltyHistoryCache`: keyed by ``request_id``, LRU eviction of rows
not in the batch being built, process-wide caches dropped in
``Sampler.__init__`` (request ids are unique only within one engine
build).

The device's watermark cache, or ``None`` when the flag is unset.

Raises at first use (fail-loud) when the watermark is combined with
the rejection sampler backend, whose draws the detector cannot
reproduce.

Per-row keep-mask for the keyed watermark (opt-out-only contract).

``keep[i]`` is False only when request ``i`` explicitly opted out
(``sampling.watermark is False``); ``None`` / ``True`` — and any request
with no ``watermark`` field — keep the row watermarked. Consumed by the
K=1 sampler and threaded into :meth:`WatermarkContextCache.mtp_seed_grid`
/ :meth:`~WatermarkContextCache.mtp_draft_seed_row` for the MTP paths.

The sampling kernel's Philox4x32-10 uniform at ``(seed, index)``.

Bit-for-bit mirror of ``_gumbel_argmax_kernel._philox_uniform``:
counter lane 0 = ``index`` (uint32 wrap), lanes 1-3 zero, key = the
64-bit seed split lo/hi, ten rounds, lane-0 mapped to (0, 1) as
``(x + 0.5) * 2^-32``.

Score a token sequence for the ``key`` watermark. Model-free.

``token_ids`` is the full available sequence (prompt + completion when
the prompt is known — pass ``prompt_len`` so prompt tokens provide
context but are not themselves scored). Positions whose ``h``-window
context repeats an earlier scored context are skipped, mirroring the
sampler's repeated-context masking (the sampler's Bloom false
positives are not mirrored — those tokens score as null noise and
only dilute the statistic).

Returns ``n`` (scored tokens), ``mean_u`` (null: 0.5), ``stat``
(``Σ -log(1-u)``; null: ``Gamma(n, 1)``), ``p_value`` (survival
probability of the null — small ⇒ watermarked), and ``z`` (normal
approximation of the same test).

Size the device pools for ``rows`` concurrent requests, then FREEZE
their addresses.

Call this with the engine's ``cfg.batch.max_batch`` — the scheduler's
admission width, and so the widest slate the pools can ever be asked
for — before recording a cudagraph over any keyed path. Afterwards
``_ring``, ``_bloom``, ``_slate_rows`` and the keep mask keep the
addresses the graph records, and a slate that would need more rows
raises instead of quietly reallocating underneath it.

The refusal is the point. A reallocation under a live graph does not
fail — the graph replays against freed storage and the draw is keyed
on whatever now occupies it, so the watermark is WRONG rather than
slow and nothing raises to say so.

Reallocate every device pool at ``capacity`` rows — the ONLY event
that moves their ``data_ptr``s, and refused once reserved.

``torch.inference_mode(False)`` is load-bearing for the same reason
as the xgrammar bitmask buffers: these pools are cross-step state
written from BOTH the K=1 sampler (outside the forward's
``inference_mode``) and the MTP seed/verify helpers (inside it).
Allocated under ``inference_mode`` they become inference tensors and
the next outside-the-window in-place write raises.

Make ``rows`` pool rows claimable.

Widening the pools' CAPACITY is a separate, checked event
(:meth:`_grow_pool`); claiming more rows inside an existing capacity
only extends the host-side ownership bookkeeping, so the device
addresses a graph recorded stay put.

Watermark seeds for the slate → ``(B,)`` int64 on device.

``base_seeds`` is the ``(B,)`` seed tensor the sampler would have
used without the watermark (the SPMD rank-agreed seeds, or the
private-generator draw); a row falls back to it when its context
window has already keyed a draw (repeated-context masking).
Everything after the host bookkeeping is device tensor ops — no
host sync.

Shift each slate row's ring left and append this step's token.

``sampled`` is the sampler's merged ``(B,)`` int64 output tensor —
a pure D2D update, valid under async output where the host has not
yet materialized the token.

Per-(depth, row) watermark seeds for one MTP verify step.

Returns ``(kp1 * B,)`` int64 (flat row ``d * B + b`` — the verify
kernel's row layout). Depth ``d``'s context is the last ``H`` of
(committed ring ++ ``draft_tokens[:d]``): the residual draw at the
first rejected depth sees exactly this context because the
accepted prefix IS the drafted prefix, and the bonus draw sees all
drafted tokens. Bloom is CHECK-ONLY (a masked depth falls back to
the row's unkeyed seed at that absolute position,
:func:`unkeyed_fallback_seed`); inserts happen in
:meth:`mtp_advance` for the depths actually consumed.

``keep`` is the per-row opt-out mask (``keep[b] is False`` ⇒ request
``b`` opted out of the watermark). An opted-out row gets the UNKEYED
per-(depth, row) fallback at EVERY depth — the same
fallback a Bloom-masked depth uses — so its verify draws the target
distribution from the rank-agreed engine seed (a valid unkeyed
sample) and, coupled with the identical draft seeds
(:meth:`mtp_draft_seed_row` receives the same ``keep``), its
draft↔verify accept behaviour is preserved. ``None`` keeps every
row watermarked (the historical behaviour).

Seeds for depths ``d_start..d_end-1``, flat row ``(d - d_start) * B + b``.

The one definition of the MTP keyed-seed computation, sliced two
ways: the verify tail asks for every depth, a drafter asks for the
single depth it is about to draw. A depth's hash is a function of
the context TOKENS alone — ring ++ ``draft_tokens[:d]``, never
another depth's hash — and the Bloom probe here is check-only, so
a sliced range is bit-identical to the same rows of a full sweep.

DIAGNOSTIC: per row, whether each depth ``0..kp1-1`` is Bloom-masked.

Host-side re-derivation of the seed kernel's Bloom probe over the same
(ring ++ drafts) windows; a per-probe device read, so never on a perf
path. Rows must already be prepared for the slate.

Overwrite opted-out rows' seeds with the UNKEYED per-(depth, row)
fallback (:func:`unkeyed_fallback_seed`) — the SAME fallback a
Bloom-masked depth uses, so an opted-out row is unkeyed at every
depth. The fallback keys on the ABSOLUTE depth ``d_start + i``, so a
sliced range reproduces the full sweep's rows — the seeds differ at
every depth by construction and are NOT hoistable across a walk.
``keep is None`` (or all True) is the no-op keyed path — it allocates
and launches nothing.

The mask lands in the persistent ``_keep_pool`` rather than a fresh
upload per call: a pageable host-list H2D cannot be captured at all,
and a re-uploaded one would move under a graph that had recorded it.

Per-row keyed seeds for draft slot ``step`` — depth ``step`` of
:meth:`mtp_seed_grid` (context = last ``H`` of ring ++
``draft_prefix``). A stochastic drafter drawing its slot-``step``
token with these seeds (zero-lane Philox layout) shares its Gumbel
noise with the verify pass's keyed choice at the same depth, so the
draft argmax and the target's keyed argmax agree whenever their
distributions do — the accept-restoring coupling for
``ARBI_TRUE_STOCHASTIC_DRAFT`` under the watermark.

``keep`` (same per-row opt-out mask :meth:`mtp_seed_grid` takes) makes
an opted-out row draw from the UNKEYED per-position fallback, matching
the seed its verify pass uses so accept behaviour is preserved.

Returns ``(B,)`` int64.

Commit one MTP verify step: Bloom-insert the consumed contexts
(depths ``0..n``) and advance each row's ring by its ``n+1``
committed tokens. Pure D2D on CUDA — host-sync-free under async
output. ``insert_blooms=False`` is the ring-only variant for
slates that committed WITHOUT keyed draws (greedy-routed steps):
the ring must track the context, but no keyed context was
consumed, so nothing is marked seen.

Row claim / prompt-rewrite rebuild / LRU touch for a slate —
the bookkeeping prefix of :meth:`row_seeds`, shared with the MTP
path (which computes its seeds in its own kernel).

The slate's row indices are WRITTEN INTO the persistent pool rather
than uploaded as a fresh tensor: ``_slate_rows`` is always a prefix
view at the pool's base address, so a graph that recorded it keeps
reading the current slate. The

Return every row to the free pool and clear its context.

The device pools are re-filled in place, not dropped: their addresses
outlive the reset, so a graph recorded against them stays valid. The
reservation is released — a rebuild may serve a different
``max_batch``, and re-reserving is how the next capture states it.

Constrained decoding via :pypi:`xgrammar`.

Surface
-------

The processor is constructed once per engine and shared across every
request. Each request owns its own :class:`xgrammar.GrammarMatcher`
(per-request automaton state); the *compiled* grammar
(:class:`xgrammar.CompiledGrammar`) is shared across requests via a
content-hash-keyed in-process cache, since compilation costs seconds
for large JSON Schemas while matcher construction is microseconds.

Per-step contract
-----------------

For each ``(k, request)`` row whose request carries a
``response_format``:

  1. Look up the request's matcher (``request.grammar_state``); if
     missing, lazily build it from the cached
     :class:`CompiledGrammar`.
  2. Fill that row's slot in a *shared* pinned host bitmask via
     ``matcher.fill_next_token_bitmask(host_bitmask, row_index)`` —
     all CPU work, no device traffic.
  3. After every constrained row is filled, perform ONE H2D copy
     of the populated host bitmask region into a persistent device
     bitmask, then ONE fused
     ``xgrammar.apply_token_bitmask_inplace(logits_2d, device_bitmask,
     indices=constrained_row_idx)`` kernel launch — disallowed tokens
     are set to ``-inf`` for every constrained row in a single pass.

After sampling, the engine calls :meth:`accept_token` so the matcher
advances its per-request automaton state.

Batching rationale
------------------

xgrammar's kernel takes an ``indices=`` argument so a single dense
bitmask covering every row in the flat ``(K*B, V)`` logits view can be
filled CPU-side, copied once, and applied to every constrained row in
one launch, instead of one per-request bitmask copy + kernel launch
per ``(k, request)``.

Cache safety / thread safety
----------------------------

The cache is a plain dict guarded by an :class:`asyncio.Lock`; entries
build via ``asyncio.to_thread`` (compilation is CPU-bound).
``GrammarMatcher.fork()`` is the documented way to spawn an
independent matcher for the same grammar; matchers do not share
mutable state with the ``CompiledGrammar``.

The shared host / device bitmask buffers are written to at *distinct
row indices* per request per step (the row index is the request's
flat (k, i) coordinate in the current logits batch), so concurrent
``fill_next_token_bitmask`` calls touch non-overlapping memory. The
processor's :meth:`__call__` runs synchronously in a single thread
(the engine's run loop), so no lock is needed around buffer access.

Return the on-disk grammar cache root, or None if disabled.

The Dockerfile sets ``ARBI_SERVE_GRAMMAR_CACHE_DIR=/cache/xgrammar``.
Outside the container we fall back to ``~/.cache/arbi-serve/xgrammar``.
A literal empty string disables disk persistence (in-memory cache only).

The single grammar source for a request — used by BOTH the single-rank
attach path and the SPMD admit path so every rank agrees.

Precedence: an explicit non-text ``response_format`` wins; otherwise tools
+ an enforcing ``tool_choice`` synthesize a structural-tag grammar. Returns
``None`` when nothing constrains decoding.

Pickle-friendly grammar spec for a request, or ``None``.

Carried once in the SPMD ``AdmitRow`` so every rank reconstructs the
identical grammar (and thus an identical matcher) locally. Reuses the
on-disk ResponseFormat payload shape.

One in-flight request's structured-decoding state.

Lives on ``Request.grammar_state`` so concurrent requests share no
mutable matcher state at the processor level. Bitmask buffers are
shared engine-wide and addressed by row index per step (see the
processor docstring); this struct only carries the per-request
matcher + diagnostics.

Build from our :class:`arbi_serve.tokenizer.Tokenizer`.

Constructs an :class:`xg.TokenizerInfo` directly from the
decoded vocab list and known stop tokens (no HF
``PreTrainedTokenizerBase`` is required).

``vocab_size`` should match the model's lm_head width (often
larger than the tokenizer's vocab — Qwen3-4B for example has a
tokenizer of 151,669 tokens but a 151,936-wide lm_head with
padding for fast matmul). When the bitmask is sized to the
narrower tokenizer vocab the model can sample logits at the
padding indices, which xgrammar then rejects with an
out-of-range warning.

Allocate a pinned (rows, ceil(V/32)) int32 host bitmask.

``pin_memory=True`` is best-effort: on CPU-only hosts (CI test
runners with no CUDA) it raises ``RuntimeError``; we fall back
to an unpinned tensor so the CPU code path still works.
``inference_mode(False)`` per :meth:`_ensure_buffers`.

Grow / re-home the persistent bitmask + indices buffers.

Buffers grow if ``rows`` exceeds the current cap; they are
also reallocated when the device changes (hot-swap). Steady
state never reallocates — the common path is a no-op except
on the first call where the device buffer must be created.

``torch.inference_mode(False)`` is LOAD-BEARING, not defensive.
These buffers are persistent cross-step scratch, but the two
writers reach them from opposite autograd contexts: the main
sampler chain (:meth:`__call__`) runs OUTSIDE the forward's
``inference_mode`` window, while the speculative verify masker
(:meth:`mask_verify_logits`) runs INSIDE it. A tensor allocated
under ``inference_mode`` is an inference tensor and can never be
updated in place again once execution leaves that region, so
letting the verify path win the lazy-allocation race would make
every later ``__call__`` raise "Inplace update to inference
tensor outside InferenceMode" — killing the whole co-batched
step, not just the constrained row. Allocating with inference
mode explicitly OFF yields normal tensors, which both writers may
update in place from either context.

LRU insert + evict down to the cap. The ONE writer of
``_compiled_cache``.

Idempotent by key: the compile is a deterministic function of the
grammar source and the key is its content hash, so a re-insert (async
attach and the SPMD sync build racing on the same key) replaces an
entry with an equivalent automaton. Callers hold a strong ref to the
object they got, so an eviction under them is safe.

Return a cached :class:`CompiledGrammar`, compiling on miss.

Compilation is CPU-bound and runs in a worker thread.
Concurrent calls for the same grammar key share one compile
job via a per-key future, so N identical-schema requests
spawn one ``to_thread`` worker rather than racing on
xgrammar's internal compile lock.

On in-memory miss the on-disk cache is probed: a hit still
recompiles (the binary form isn't stable across xgrammar
versions) but skips the network / file fetch.

Build the per-request matcher. The ONE builder of
``Request.grammar_state``.

Called by the engine before the request is admitted, on every routing.
Idempotent: a second call is a no-op once ``grammar_state`` is
populated.

Every exit either installs a matcher or explicitly clears
``grammar_required`` — a request must never leave here requiring a
grammar it does not have.

Build a mirror's per-rank matcher from a broadcast grammar payload.
The ONE builder of ``_MirroredRequest.grammar_state``.

Returns ``None`` on a compile failure rather than diverging — every
rank degrades identically since the payload is the same everywhere.
The caller must then clear the mirror's ``grammar_required``.

Whether any row carries a live matcher.

The SAME condition :meth:`mask_verify_logits` short-circuits on, so
the answer cannot disagree with what the mask would do. A row that
merely REQUIRES a grammar whose matcher is not built yet writes
nothing this step.

The chain's verify-slate form of the grammar mask.

Per-position by construction: :meth:`mask_verify_logits` advances
each constrained row's matcher through its own drafts, so position
``p`` is masked by the state the model actually conditioned on. The
matcher walk is CPU work every rank already does identically, so a
``window`` narrows only the device-side mask application.

Token-budget continuous-batching scheduler.

Two surfaces:

  - :class:`Scheduler` — single-FIFO continuous batcher for the active
    resident.
  - :class:`MultiGroupQueue` — the graceful cross-resident park queue
    with swap-cost-aware selection, drained by the run loop's stable-VA
    residency advance.

Admission-time accounting for one step's forward-activation bytes.

Admission commits KV PAGES per row and refuses a row it cannot page
(``Scheduler._can_allocate_reserved``). It commits nothing for the bytes that
row's tokens will allocate in the forward, so a slate whose pages fit but whose
GEMM / gather working set does not is admitted anyway and the forward OOMs —
after the Phase-2 freeze, with KV already locked, i.e. as a client-visible 500.
The serving floor
(:func:`~arbi_serve.engine.inprocess_capture.serving_floor_for_grow`) reserves
free VRAM for that working set, but a reserve is only a promise about a state
nothing prevents; this module is the matching enforcement, so the state the
floor fears is unreachable rather than merely unlikely.

The model
--------
One step's forward activation bytes, as a function of the three quantities
admission already has in hand while it builds the slate::

    demand = per_prefill_token x prefill_tokens
           + per_row           x rows
           + per_gather_kv_token x gather_kv_tokens
           + per_query_token   x query_tokens

``prefill_tokens`` and ``rows`` are the axes the boot activation profile
measures directly (:mod:`arbi_serve.runtime.activation_profile` runs a pure
decode shape at ``max_batch`` rows and a pure prefill shape at
``max_batched_tokens``), so the two slopes are a two-point calibration on
measured data rather than an analytic model of the layer stack.

``query_tokens`` is the DECODE-ROW query width past one token each. Every
profiled shape above runs its decode rows at exactly one token, but a served
MTP verify step runs ``K + 1`` flat tokens per opted row
(:func:`~arbi_serve.spec_decode.mtp_verify_plan.build_verify_plan`). On a
capture hit those bytes are the graph pool's and are budgeted there; on a MISS
the step falls to a full eager per-layer forward at that width
(``EagerModelRunner``'s ``eager_verify_offladder`` rung) whose activations land
in ``scratch.forward_arena``. The axis is therefore reachable, unbudgeted
without this term, and measured by its own profiled shape
(:data:`~arbi_serve.runtime.activation_profile.SHAPE_VERIFY_FORWARD`) rather
than derived from the decode shape. The PAGE side of the same step already
reserves it — :meth:`~arbi_serve.scheduler.scheduler_allocation.
_AllocationMixin._decode_page_commit` commits the ``1 + K`` extension's
pages — so this is the byte half of a reservation the scheduler already makes.

``gather_kv_tokens`` is the axis the profile CANNOT see. The tkv Turbo prefill paged
prefill takes a padded short-circuit whenever the step carries exactly one
prefill row and that row's accumulated KV is at or below
``PREFILL_PAGED_GATHER_MAX_TOKENS``; the short-circuit materializes the whole
accumulated context per (chunk, layer). The profile's prefill probe is a FRESH
sequence, so it measures that term at chunk width and never at the context
width a long chat reaches. It is the one term that scales with something other
than the step's own token count — and therefore the only one that can carry a
slate past what the floor reserved.

Charging the gather
-------------------
The short-circuit is mutually exclusive with co-scheduling: it needs ``B == 1``,
so a step carrying two prefill rows takes the zero-copy varlen route and stages
nothing context-scaled. A slate is therefore charged the LARGEST single prefill
row's gather, never a sum — an upper bound that only falls as more prefill rows
join the step.

Enforcement contract
--------------------
The budget is the bytes the serving floor actually left for the step, so the
reserve and the gate are one quantity rather than two that can drift. Both are
calibrated on the same profile, so every slate whose cost the profile SAW fits
the budget and the gate does not fire — the whole of ordinary decode, mixed and
chunk-prefill traffic. What can exceed it is the term the profile did not see:
one prefill row's context-scaled gather. Whether the gate fires at all on a
given boot is therefore an observable property of that boot's numbers, not an
assumption — :mod:`arbi_serve.engine.activation_admission` computes the widest
admissible slate against the reserve at arm time and says which way it came out.

A PREFILL row is narrowed rather than refused: :meth:`StepActivationModel.
max_prefill_tokens` inverts the model in its token argument, so the scheduler
caps the chunk at what the budget covers and the row still advances every step.
A row that cannot fit the budget even at its minimum width is admitted anyway
with a bounded warning — one row is the minimum unit of progress, and starving
it silently is worse than running it. So admission bounds CO-RESIDENT demand
exactly, narrows the one term it can narrow, and names a row that exceeds the
budget on its own.

Pure arithmetic — no CUDA, no engine state, no allocation.

tkv's live gather ceiling, falling back to the mirrored constant.

Reading tkv's own value keeps the budget correct across a tkv bump that
moves the ceiling — including the direction that matters most, a tkv that
caps the gather at one chunk, which collapses the context-scaled term and
lets the floor follow the real peak down without a second hand-tuned
number here.

Calibrated per-step activation cost and the budget admission enforces.

Every field is bytes except ``gather_max_kv_tokens`` (KV tokens) — the
ceiling past which tkv stops materializing the gather, so charging beyond
it would price a state the kernel cannot enter.

``budget_bytes`` is ``None`` when no budget is armed — no reading was taken,
the gate is inert, and every slate is admissible. An integer is a READING,
and ``0`` is a reading like any other: the card saying a step may spend
nothing. Those are opposite instructions to admission, so they cannot share
a value. A single sentinel for both makes the inert branch the one an
exhausted card always reaches, which turns the gate off at the moment it is
the only thing left.

One profiled shape, as the calibration reads it.

``gather_kv_tokens`` is the accumulated KV the probe's own attention
gathered — nonzero only for the single-row prefill probe, since the padded
gather needs ``B == 1`` and the multi-row decode / mixed probes take the
zero-copy route.

``query_tokens`` is the EXTRA query tokens the probe's rows carried beyond
one each — nonzero only for the verify-width probe, whose rows each run
``K + 1`` tokens through the model. It is what separates a shape that
calibrates the query axis from one that calibrates the prefill-token axis:
the two are different forwards at the same token count (a decode-class
batch with a wide ``max_query_len`` against a prefill batch), and fitting
one slope over both would price neither.

Fit the token and row slopes to the profiled shapes.

The boot profile measures the two axes admission controls at their
extremes: a pure-decode shape (many rows, one token each) and a pure-prefill
shape (one row, ``max_batched_tokens`` tokens). The token slope is the rise
between them once each shape's own gather is netted out; the row slope is
what the decode shape costs beyond its handful of tokens.

A CONSTANT is then lifted until the model reproduces every profiled shape at
or above its measurement. That is what keeps a linear fit honest against a
peak that is not linear: the mixed shape sits above the line through the two
extremes (it runs a prefill chunk and a decode batch in one forward), and a
model that under-predicts a shape the engine demonstrably runs would hand
admission a budget the forward exceeds.

Every coefficient is clamped at zero, so a profile with a missing or
zero-recorded shape (an OOM'd probe records 0 bytes by design) yields a
smaller model rather than a negative slope that would make wider slates
look cheaper. Returns ``None`` when fewer than two distinct token widths
were profiled — there is then no slope to fit, and an unarmed gate is the
correct outcome.

THE QUERY SLOPE IS FIT LAST, AND ONLY AGAINST ITS OWN RESIDUAL. An
observation with ``query_tokens > 0`` is a decode-class batch whose rows
each carry ``K + 1`` tokens; it is held out of the three fits above and then
charged only what the finished model does not already predict for the same
rows at one token each. Two reasons it cannot be one of the token-slope
points. Its tokens are not prefill tokens — admission calls
:meth:`~StepActivationModel.demand_bytes` for a pure-MTP decode slate with
``prefill_tokens=0``, so a slope fit through it would price bytes no term
ever charges. And folding it into the constant instead would charge its
whole excess to every slate, including the ``mtp_k == 0`` traffic that never
runs the shape, which narrows prefill chunks for a step that cannot occur.
Clamped at zero like every other coefficient: a verify probe that came in at
or below what the row slope already predicts leaves the model exactly what
it was, which is the honest reading of "this axis costs nothing extra here".

Modelled forward-activation bytes for a slate of this shape.

``query_tokens`` is the slate's EXTRA decode-row query tokens — the
drafts an MTP row's verify step carries beside its committed tail, i.e.
``K`` per opted row, not ``K + 1``. The tail token is the one every
decode row already pays through ``per_row_bytes``, so charging ``K + 1``
here would price it twice.

Prefill tokens one more row may take and keep the slate in budget.

``None`` means "unbounded by this model" — no budget armed, or a token
slope of zero, either of which makes a token cap meaningless. A budget
of zero is NOT unbounded: it caps the row at zero tokens, and the
scheduler's empty-slate floor then advances it one token at a time,
which is the narrowest progress admission can make.

This is the inverse of :meth:`demand_bytes` in its one continuous
argument, and it is what lets a long prompt stay servable on a card
whose free VRAM is below a full-width chunk. Refusing such a row would
starve it and admitting it whole would OOM the forward; narrowing the
chunk does neither — the row advances every step, just by fewer tokens.

Demand of the widest slate the admission limits allow.

The number the budget must be at or above for the gate to be
TTFT-neutral: every slate admission can build costs at most this.

``max_query_tokens`` is ``max_batch * mtp_k`` — every row opted into
speculation at the deepest draft the boot resolved. It is a separate
argument rather than a product taken here because the model does not
know ``mtp_k``: the draft depth is the engine's, and a widest-slate
number that guessed it would report a bound admission does not enforce.

Per-attention-DP-rank admission: one scheduler per request shard.

At ``attn_dp_size > 1`` an attention-DP set owns whole REQUESTS — their
tokens and their whole KV — so admission is a PER-SET decision. One
scheduler over one page table cannot make it: the sets' pools fill
independently, and a slate every set forwards would put the same requests
(and the same KV) on every set, which is the replication the axis exists
to remove.

So rank 0 runs ``attn_dp_size`` schedulers, one per set, each over its own
page table, and this group is the single object the engine holds in
``eng.scheduler``. Each member is an ORDINARY
:class:`~arbi_serve.scheduler.scheduler.Scheduler` with no attention-DP
awareness at all: per-set page accounting, preemption and backpressure
fall out of the fact that a member can only see its own table.

Rank 0 backs member 0 with the REAL pool (its own set's KV) and every
other member with an :class:`~arbi_serve.cache.attn_dp_shadow_pool.AttnDpShadowPagedPool`
— see that class for why a page-count mirror is sufficient and what would
break if the ids had to match.

OWNERSHIP IS PERMANENT. :meth:`add` stamps ``Request.attn_dp_rank`` once
and every later call routes on it; a request's pages live in exactly one
set's pool, so re-homing it would strand them. New requests go to the
member with the fewest resident requests (lowest index breaks the tie),
which is a rank-0-local decision and never crosses the wire — the
per-member slates it produces do.

THE STEP BUDGET IS SPLIT, NOT REPLICATED. Every rank's FFN runs over the
GATHERED rows of all ``attn_dp_size`` sets, so the whole step's token
count — not one set's — is what the activation arena and
``max_batched_tokens`` bound. Each member therefore gets ``1 /
attn_dp_size`` of the batch and token budget, so the gathered width stays
inside the profiled envelope.

One attention-DP member's share of the step budget.

``max_batch`` and ``max_batched_tokens`` bound the GATHERED step (see
the module docstring), so each member may contribute at most its
share. Both floor at 1 — a member that could admit no row would be a
set that never serves.

Read-only concatenation of the members' ``waiting`` / ``running``.

Admission and the health surfaces read these as a sequence
(``len`` / iteration / truthiness) and never mutate them, so a view
beats rebuilding a list on every arriving request.

The members' page tables as one read-only accounting surface.

:func:`arbi_serve.engine.admission.check_admission` measures KV
occupancy against ONE rank's ``pool.num_pages``, so the availability
it must see is the FULLEST set's — the set a new request could land
on. ``min`` is therefore the aggregate, not ``sum``: summing would
report ``attn_dp_size`` pools' worth of headroom against one pool's
denominator and never watermark at all.

``attn_dp_size`` schedulers presented to the engine as one.

Exposes the exact surface the engine, the admission gate and the SPMD
driver call on a :class:`~arbi_serve.scheduler.scheduler.Scheduler`.
Anything not listed here is deliberately absent rather than forwarded
to member 0: a call that silently hit one set's queues while the
others held requests would be wrong in a way no test would show.

The least-loaded set, lowest index breaking the tie.

Load is resident request COUNT, not tokens: the pages a request
will take are unknown at admission (chunked prefill grows them
over many ticks), while the count is what bounds each member's
share of the split step budget.

One slate per attention-DP set, in set order.

The SPMD driver's entry point: each set forwards ITS list and
nothing else, so this is the per-rank ``StepPlan`` source. Members
are scheduled in index order on rank 0, which makes the split a
deterministic function of the arrival order — the same property
the single-scheduler path relies on.

Make admission policy compression-aware.

The KV pool's page CAPACITY already scales with the codec: num_pages =
kv_bytes // (backend.bytes_per_token The
scheduler's admission GATE (FlatPageTable.can_allocate) correctly consumes
that larger free list.

But the admission POLICY caps — ``cfg.batch.max_batch`` (default 32) and
``cfg.batch.max_batched_tokens`` (default 8192) — are static constants that do
NOT scale with the codec. So the recovered capacity sits idle: the engine
admits the same 32 sequences whether

This derives a compression-aware ``max_batch`` from the realized page capacity
and a target context, so the policy fills the capacity the codec bought.

Concurrent sequences of ``target_context`` tokens the pool can hold.

``num_pages`` reflects the active backend's compressed per-token cost
(per-layer under smart-mix), so this rises with the compression ratio
— unlike a static ``max_batch``. ``hard_cap`` clamps to an operator
ceiling (e.g. a scheduler/runtime limit).

Admission-time MTP K-bucketing policy.

The verify pass requires a SINGLE uniform K every step (cudagraph capture
invariant — baked-in ``cu_seqlens_q`` / ``tokens_per_seq``). Rather than
coerce all rows down to ``min(K)`` (which would throw away every K=7
row's draft chain that step), the scheduler picks ONE K per step and
drops rows at other K values from the slate; the dropped rows stay in
``Scheduler.running`` and run next step.

``MtpBucketer`` owns the per-K consecutive-skip bookkeeping
(``skip_counts``) so the starvation guard
(``BatchConfig.mtp_admission_starvation_steps``) can force a K-bucket
pick after the bucket has been skipped that many steps in a row.

It also applies the speculation load valve, which decides whether a step
drafts at all. The valve's tri-state and the measurement behind ``auto``
live in :mod:`arbi_serve.scheduler.spec_valve`; this module owns only the
mutation (transiently zeroing ``sampling.mtp_k``) and its observability.

Seed ``auto``'s starting arm from the architecture fallback.

The seed decides only which arm a width runs BEFORE it has been
measured; the first phase pair at that width replaces it.

Coerce the whole step to mtp_k=0 when speculation is not paying.

FIRST restore any rows zeroed last step, THEN — if the resolved
valve says stop — transiently zero ``sampling.mtp_k`` on every
opted decode row so the step routes through the proven MTP-OFF
decode path (no seed-drafter, no verify, no worker-bridge
drafter/seed ops).

Which rows are zeroed is the same in every mode; the modes differ
only in what decides. ``off`` never zeroes, ``on`` zeroes at a
pinned decode width, and ``auto`` asks
:class:`~arbi_serve.scheduler.spec_valve.SpecValveController`,
which compares the two arms' measured committed tokens/sec at the
width being served.

Rank-symmetric by construction: only rank 0 runs the scheduler,
and zeroing the K makes rank 0 emit the SAME op stream a
genuinely non-opted slate emits — the MTP-fill gate
(``mtp_fill_enabled=False``) + ``mtp_per_req_k=None`` ride the
per-step broadcast, so the worker rebuilds an identical batch.

Mutation is scoped to a single engine step: the restore at the
top of the NEXT call runs before any reader sees the request.

Record the load valve's state, logging only on a transition edge.

Off the per-step cost path by construction: the common case is a
bool compare against the state this method last stored. Only an
edge reaches the clock or the logger, and consecutive edges inside
``_TRANSITION_LOG_MIN_INTERVAL_S`` are folded into the next line's
suppressed count so an oscillating batch cannot flood the log.

Drop MTP-opted decode rows whose K differs from the chosen step K.

Policy (strict per-step single-K slates):

  1. Find the MTP-opted decode rows in the slate
     (``mtp_k > 0`` and ``not is_prefill``). Prefill rows and
     ``mtp_k == 0`` rows pass through untouched — they don't
     enter the verify pass and never need K-uniformity.
  2. If 0 or 1 distinct K values are present among MTP rows,
     return the slate unchanged. Common case at low traffic.
  3. Otherwise, group MTP rows by K. Pick ``step_k`` as the K
     with the most MTP rows. Tie-break: the K of the
     slate-earliest MTP row at any tied count (FCFS-like —
     oldest-waiting request wins).
  4. Starvation guard. If any OTHER K bucket's consecutive-
     skip count has reached ``batch_cfg.mtp_admission_starvation_steps``,
     override step 3 and pick THAT K instead (oldest-skipped K
     wins on ties). Prevents one K bucket from being
     permanently starved by a heavy stream of another K.
  5. Drop every MTP row whose K != step_k from the slate. The
     dropped rows stay in ``Scheduler.running`` (we never popped
     them) so they're retried next step. Update skip counts:
     ``+1`` for each K bucket we just skipped; reset to ``0``
     for the K bucket we picked.

Idempotent: calling on an already-uniform slate is O(slate)
and a no-op past step 2.

Returns the filtered slate (same rows for prefill / non-opted
decode; only MTP-opted rows at ``step_k`` survive). The
scheduler's ``commit()`` only operates on the returned rows;
the engine's downstream verify pass sees a uniform-K slate by
construction.

Graceful cross-resident park queue for the stable-VA residency swap.

Each routing key is one **group** FIFO — a stable-VA resident's key is
``(resident, "")``. The :class:`MultiGroupQueue` holds one FIFO per
key; the run loop drains the currently-active resident's scheduler,
then asks :meth:`pick_next_group` for the cheapest non-empty next
target. ``cheapest`` is decided by the swap class
:func:`arbi_serve.engine.params_hash.cheap_param_diff` reports between
``current_group``'s params and the candidate's params, with an
age-based promotion to prevent indefinite starvation of an
expensive-swap target.

Key design points:

  - The queue carries **opaque** request handles. It never touches the
    model / pool / page table — those live in the engine. This keeps
    the queue trivially testable and lets the engine evolve its own
    request shape without the queue moving in lockstep.
  - Per-key FIFO is preserved (no priority inversion within a
    resident). Cross-resident selection is the only place a "newer"
    request can be served before an "older" one — and only when the
    older request's resident is the most expensive swap target AND its
    age is under the anti-starvation threshold.
  - Time injection (``_now``) is plumbed through every method that
    reads the wall clock, so unit tests are deterministic.

The single-FIFO :class:`Scheduler` serves the active resident; this
queue holds the parked non-active residents.

Per-group FIFO + swap-cost-aware cross-group selector.

The selector is independent of the swap orchestrator: it scores
candidates against ``current_group`` using
:func:`cheap_param_diff` on the cached engine-params dicts the
engine registers via :meth:`register_group`. When no params dict
is registered for a candidate yet (which happens when the engine
has never seen this group), we fall back to "model swap" cost.

Total requests parked across all groups (FIFO not yet admitted).

These requests are NOT on the GPU — they sit pre-admission waiting
for their group to become active. The critical-section drain
(:func:`arbi_serve.engine.critical.has_inflight`) subtracts this from
the live-request count so a run-loop-driven group swap does not wait
on the very request that triggered it (the self-deadlock the multi-
group cycling path otherwise hits).

Pick the cheapest non-empty group to drain next.

Algorithm:

  1. If no groups have pending requests → ``None``.
  2. Compute an age for each candidate — seconds since its
     oldest pending request was enqueued.
  3. Any candidate whose age exceeds ``max_age_s`` is "starving"
     and gets promoted to the front, regardless of swap cost.
     Tie-break starvation by oldest-age (most starved wins).
  4. Otherwise, score each candidate by swap cost vs
     ``current_group``'s params (via
     :func:`cheap_param_diff`). Lowest cost wins; ties broken
     by larger queue depth (drain the bigger backlog first),
     then by oldest age, then by group_key for determinism.

``current_group`` may be ``None`` (engine hasn't picked one
yet — every candidate is "model swap" cost).

``sleep_threshold_s`` is accepted but not used by the selector.

Build the params dict the selector scores against.

When the engine has registered explicit engine params for
``group_key`` (via :meth:`register_group`), use those — they
carry K_BITS / V_BITS / recipe / etc, so the
:func:`cheap_param_diff` classifier can distinguish recipe vs
backend vs model swaps with full fidelity.

When no explicit params are registered, fall back to a
synthetic ``{"model": <model_path>, "_params_hash": <hash>}``
dict — cheap_param_diff will at least correctly tag "different
model" as the model-tier swap, and "different params_hash on
the same model" as a backend swap (because _params_hash isn't
a known engine key, but it's stable per group). To make the
backend-tier classification work without a known key, we copy
the params_hash into the ``attention_backend`` slot — a known
backend-scope key.

Suspended-job KV offload (batch-backlog parking).

When a SUSPENDED batch decode job's paged-KV pages (and recurrent state if
hybrid) are snapshotted GPU→host→(disk) and its GPU pages freed, a huge
pooled backlog can wait with KV parked on disk instead of pinning VRAM.
On resume the pages are faulted back into a fresh allocation BEFORE the
job rejoins a slate.

``OffloadManager`` owns the per-step suspend / resume policy and the
pool-pressure watermarks; the parked-job deques (``Scheduler.suspended``,
``Scheduler._dropped_suspended``) and the snapshot store
(``Scheduler._suspended_store``) stay on the :class:`Scheduler` instance
the manager is constructed with, so the ``_add_inner`` self-heal and
code that reaches into ``Scheduler`` attributes directly keep working.
The manager reads/mutates those via the ``sched`` reference.

Per-step suspend / resume policy. Interactive is the priority.

Suspend (offload GPU→host→disk, free pages):
  * ANY interactive request present → suspend EVERY resident batch
    decode job (interactive runs alone at idle; batch parks).
  * else, pool pressure → suspend oldest-admitted resident batch
    decode jobs until free pages recover above the high watermark.
Resume (fault back into fresh pages):
  * only when NO interactive present AND the pool has headroom;
    bounded to ``_offload_resume_per_step`` per step so the
    fault-back work stays small.

Return + clear jobs whose resume failed irrecoverably.

The engine run loop calls this each step and runs the terminal
path (``finished`` + on_finished) for each — they already carry a
``finish_reason``. Empty list in the common case.

Offload ``req``'s KV off the GPU and park it; free its pages.

Snapshots the request's paged-KV page content (D2H) — and its
recurrent state if hybrid — into the :class:`SuspendedJobStore`
(host RAM, spilling to disk under host-budget pressure), then
frees the GPU pages so the pool reclaims them. The request leaves
``Scheduler.running`` and enters ``Scheduler.suspended`` in
SUSPENDED state.

Only DECODING batch jobs are offloadable (a mid-prefill job has
no stable resume contract here — we leave it resident and let the
prefill-cap path bound its step instead). Returns True on offload.

Runs on the engine loop thread (the snapshot D2H + free are CUDA
ops). The D2H gather is BLOCKING (``async_copy=False``, so the
staged entry carries no ``pending_event``): the forward runs on a
separate worker-thread stream, which an async copy on the loop
stream is not ordered against, and suspend frees the pages the
moment the gather returns. See the call site for the ordering
argument.

Fault ``req``'s KV back into fresh GPU pages and rejoin running.

Inverse of :meth:`suspend_to_offload`. Looks up the snapshot
(host fast-path or disk read), allocates ``num_pages`` fresh
pages, scatters the page content back (H2D), restores recurrent
state, re-registers the request with the page table, and moves it
from ``Scheduler.suspended`` → ``Scheduler.running`` (DECODING).
Returns True on success.

Fails (returns False, leaves the job suspended) when the pool
can't allocate the pages right now — the caller retries on a
later step once pages free up. A fault MISS (snapshot evicted /
TTL / corrupt) drops the job (it can't be resumed correctly);
surfaced via the metric.

Preemption victim picker (priority + fairness aware).

When a running request needs a new KV page and the primary paged pool is
full, :meth:`Preemptor.preempt_for_space` bumps a running request back to
WAITING to free its pages. Batch-priority running requests are evicted
before any interactive request; within a priority class the
lowest-fairness-weight tenant is preempted first, then the
most-recently-admitted row.

The picker operates over the live :class:`Scheduler` it is constructed
with — it reads/mutates ``sched.running`` / ``sched.waiting`` and reads
``sched.page_table`` / ``sched.pool`` / ``sched.primary_paged_kind`` /
``sched._metrics`` and the tenant fairness helper.

Tiebreak key used by the preemption picker.

Higher ``fairness_weight`` = more important; we preempt the
LOW-weight victim first when the pool is exhausted. Anon /
no-tenant requests are weight 1.0 by default.

The running set in the order :meth:`preempt_for_space` evicts it.

Primary key: priority class (batch = 0 sorts before interactive =
1, so batch is evicted first). Secondary: fairness weight (lowest
first). Tertiary: recency (most-recently-admitted / rightmost
first).

STABLE UNDER REMOVAL, which is what lets a caller that needs
several victims compute this once and consume it. Removing an
element shifts the deque index of everything to its right down by
one, uniformly, so the ``-index`` tiebreak preserves the relative
order of every survivor; the priority and weight keys do not
change at all. The all-interactive/unweighted fast path is the
same order the sort produces for those keys, so a mid-sequence
transition into it (the last weighted or batch row was the victim)
agrees with the precomputed remainder as well.

Bump a running request back to WAITING to free KV pages.

Duplex-lane requests (``Request.is_duplex_frame``) are

Priority bias (opportunistic batch). The victim picker is
priority-aware: low-priority ``"batch"`` running requests are
evicted BEFORE any interactive request. Within a priority class
the legacy ordering holds — lowest fairness-weight tenant first,
then most-recently-admitted (rightmost of the deque). So an
interactive request never loses its KV pages while any batch
request is still resident, which is what keeps interactive TTFT
flat under a batch backlog.

``admitting_is_batch`` — when the request DRIVING this
preemption is itself batch-priority, interactive requests are
excluded from the victim pool entirely: we never preempt
interactive to make room for batch. If no batch victim is
available the preempt fails and the batch admission is deferred
to a later step (it only ever runs on genuinely spare capacity).

Fairness tiebreak. Within a priority class we PREEMPT THE
LOWEST-WEIGHT TENANT FIRST. Equal-weight tenants fall back to
most-recent-first. Without any priority or quota config every
request is interactive / weight 1.0, so the ``ordering`` below
reduces to a most-recent-first deque walk.

``slate_req_ids`` is the set of request ids already attached
to the in-progress slate this :meth:`Scheduler.schedule` call.
We refuse to preempt those: the slate row holds a Python
reference to the live ``Request`` object, and preempting it
in-place would rebuild ``prompt_token_ids`` / reset
``prompt_consumed`` / clear ``output_token_ids`` while the slate
row still points at the same object. Downstream
:meth:`_build_batch` then either slices a stale prompt at line
945 or — when both the rebuilt prompt and the cleared output are
empty (the ``chunk_prefill=256`` second-admission case) — raises
``IndexError: list index out of range`` reading
``prompt_token_ids[-1]``. Excluding slate reqs from the victim
pool keeps preempt scoped to genuinely-idle running rows; if
no eligible victim remains, the caller's
``can_allocate / preempt`` gate fails and admission breaks
out of its loop (the request retries next step).

What the victim keeps. Its pages are COMMITTED into the radix
tree before they are released (:meth:`_commit_written_pages`), so
the restarted prefill can alias back what it already computed
instead of recomputing it. That claim is made at ADMISSION, not
here — see ``Scheduler._rebind_preempted_prefix`` — because a
matched page is ref-ed back up and stops being evictable, and a
request sitting in ``waiting`` holding its whole prefix takes the
pool away from every request that could actually run. Here the
pages go to ref 0: still in the tree, matchable, and reclaimable
by whoever needs them first.

``ordering`` — a victim order this call should consume instead of
building its own, for a caller taking several victims in a row
(:meth:`~arbi_serve.scheduler.scheduler_allocation._AllocationMixin._preempt_until_allocatable`).
Entries already gone from ``running`` are skipped by the
``_running_remove`` guard below, and :meth:`victim_ordering`
documents why a precomputed order picks the same victims in the
same sequence as rebuilding it per victim does.

Promote the victim's completed full pages into the radix tree.

``commit_full_pages`` runs at end-of-prefill and again at
:meth:`Scheduler.finished`; a preempt is neither, so a victim that
was DECODING has its generated tokens' pages sitting private and
unmatchable at the moment they would be most useful — the victim
is about to re-prefill exactly those tokens.

Truncated to the tokens whose KV is actually WRITTEN, which is the
page table's own ``length``. Two positions in the full sequence do
not have KV and the table is the only thing that knows it: a decode
row's most recently sampled token is appended to
``output_token_ids`` at sample time but is not fed until the next
step's forward, and a chunked-prefill row has consumed only part of
its prompt. Committing past either would publish a page with an
uncomputed slot into the shared tree, where any later prefix match
would read it as valid. The truncation costs at most one page of
recompute.

Prefix-serial admission for concurrent requests sharing a prompt.

The scheduler lets one active request publish a shared prefix before related
followers enter their first slate. Followers are rebound to the radix cache
and skip that work; unrelated requests may fill every spare batch slot.

The common no-sharing path uses a tenant-scoped index of the first full page.
It returns the FCFS head after one dictionary miss. Full page-aligned LCP work
and bounded queue scanning happen only after a first-page collision.

``ARBI_PREFIX_GROUPING=off`` disables the policy and restores strict FCFS.

Common page-aligned prefix length, in tokens, between two requests.

The metric the engine actually cares about: ``score_pair == K``
means batching A and B together aliases ``K`` tokens of prefill
into one physical computation instead of two.

Per-request memo of :func:`_page_chunks`.

``group_for_batch`` runs on EVERY ``schedule()`` step while the
waiting queue is non-empty, and rebuilt each candidate's page-chunk
tuple from scratch — at window=128 and 16k-token prompts that is
~2M tuple elements per engine step, a measured double-digit share
of the engine thread under a c64 long-context backlog. The prompt
is immutable once queued, so the chunk tuple is computed once per
(request, block_size) and cached on the request. Guarded by
identity on ``prompt_token_ids`` so any exotic prompt swap
invalidates naturally. Decisions are bit-identical to the uncached
path.

Stateless-ish grouper — maintains no cross-step state.

The current implementation is pure-function over the pending queue.

Args:
    block_size: page size; the grouper works in page-chunk units.
    window: cap on how many leading pending-queue entries are
        considered as candidates. Bounds the grouper's per-step
        cost at O(window²).

Resolve the effective grouper window from ``ARBI_PREFIX_GROUPING_WINDOW``.

Factored out of :func:`build_grouper_from_env` so the scheduler can
re-resolve the window on every step (``prefix_grouping_window`` is one
of the 3 fresh-read runtime flags — a live change must apply without a
member rebuild, same as ``prefix_grouping`` itself). Both call sites
share this one implementation so the truth table stays single-sourced.

Re-resolve ``grouper.window`` from the CURRENT env, in place.

Called every scheduling step while grouping is enabled (see
``Scheduler.schedule``) so a live ``ARBI_PREFIX_GROUPING_WINDOW`` flip
takes effect on the next step without discarding the grouper object —
preserving its cross-step ``blocked_possible`` hint, which a full
rebuild would otherwise reset.

Total shared-prefix tokens within a candidate batch.

Computed as the sum of (matching-page count × ``block_size``)
savings against the leader. The leader pays its own prefill;
each follower saves up to its LCP with the leader. This is the
exact same accounting the radix cache will produce when the
batch hits.

Pick a batch from ``pending`` maximising shared-prefix density.

Algorithm: take ``pending[0]`` (FCFS leader — preserves
head-of-line fairness; we never starve the oldest waiting
request), then greedily admit up to ``max_batch - 1`` more
from ``pending[1 : 1+window]`` in DESCENDING order of LCP with
the leader. Ties broken FCFS.

Returns the chosen requests in admission order (leader first,
then highest-LCP siblings). The scheduler's main loop then
admits them into the batch as if they were the head of the
FCFS queue — which they effectively are.

Returning a SUBSEQUENCE of the original pending order is
important: the scheduler will pop these from the actual deque
in the order returned, and the leader-first invariant
guarantees no head-of-line starvation.

Per-tenant rate-limit + concurrency-cap slate filter.

This is the LAST admission gate before a step's slate goes out: it
drops slate rows whose tenant has exhausted its quota, leaving those
requests queued (``Scheduler.running`` / ``Scheduler.waiting``) for the
next step. Zero cost when every request is anon / UNLIMITED.

``QuotaFilter`` owns the :class:`QuotaTracker` and the per-reason
deferral counters (the test / metrics surface). The
:meth:`Scheduler._tenant_for` resolver is passed in so tenant resolution
stays single-sourced on the scheduler.

Drop slate rows whose tenant has exhausted its quota.

Hot-path order: skipped rows leave the requesting request in
``Scheduler.running`` / ``Scheduler.waiting`` for the next step's
slate. Pre-existing FCFS / chunk-prefill / page-boundary
admission gates have already passed; this filter is the LAST gate
before the slate goes out, so pages aren't allocated only to
be rolled back.

``cost_tokens`` is the new tokens the request will produce
this step (``n_tok`` from the slate row). For prefill rows
``n_tok`` is the chunk size; for decode rows it's 1
(or ``mtp_k`` for a verify pass — the verify slate emits the
post-MTP-bucketing K rows separately).

Empty slate, no-tenant slate (anon / UNLIMITED everywhere),
and no-tracker → fast no-op.

Side effect: increments per-reason deferral counters readable
via :meth:`deferred_by_quota` for tests / metrics.

Recurrent-state savepoint snapshot worker (device-side).

The :class:`Scheduler.commit` per-token accounting loop stays in
``scheduler.py`` (the off-by-one class hides in that loop and can't be
GPU-validated here); only the device-side snapshot worker — the async
D2H clone + ``put`` into the host-RAM savepoint store — lives here.

The admission-side resume-key resolution stays inline in
:meth:`Scheduler._add_inner`: it is woven through the radix
``add_request`` / ``prompt_consumed`` accounting and the KV-coverage
cross-check, where lifting it would risk the documented token-accounting
off-by-one. ``SavepointSnapshotter`` reads ``sched._savepoint_store`` /
``sched.pool`` via the ``sched`` reference.

True when snapshots land in a bounded pinned-host ring.

The ring changes what "the entry is gone" MEANS: a wrap is routine, an LRU
eviction between match and get is a budget fault. The two must not share a
failure path.

Resolve an admission-time savepoint resume for ``req``.

Returns ``(savepoint_resume_key, savepoint_entry)`` — both
``None`` when no host-RAM checkpoint covers a page-aligned
prefix of the request's prompt. The ``prompt_consumed`` / radix
``match_len`` accounting that consumes this result stays inline
in ``_add_inner`` (the documented off-by-one site).

Raises ``RuntimeError`` on a best_match→get eviction race (the
"no silent fallbacks" hard-fail).

Run the radix ``add_request`` + savepoint reconciliation.

Performs the radix-cache ``add_request`` (returning the
page-aligned ``match_len``), the KV-coverage cross-check that
releases a savepoint whose coverage outran the
radix match, and the hybrid-model force-disable fallback that
re-admits with ``cache_enabled=False`` when a hybrid model
matched pages but has no savepoint to restore the recurrent
state at the boundary.

``readmit`` routes the registration through
:meth:`RadixPageTable.readmit_preempted` instead of ``add_request``
— the preempt path, where the request is ALREADY registered and the
seam has to free its pages and carry its cache identity across the
drop. Everything downstream (the coverage cross-check, the hybrid
force-disable) is identical, which is the point of routing it here
rather than restating the reconciliation at the preempt site.

``page_budget_cap`` is a second, independent cap on the match, in
tokens, MIN-ed with the recurrent-coverage cap. Preemption sets it
to the part of the victim it must genuinely give back (a re-matched
page is ref-ed up and stops being evictable, so an uncapped
re-match frees nothing for the request the preemption was for);
``None`` on the ordinary admission path, which has no such budget.

Returns the FINAL ``(match_len, savepoint_entry, savepoint_resume_key)``.
The caller (:meth:`Scheduler._add_inner`) consumes ``match_len``
to credit ``prompt_consumed`` inline — that ``prompt_consumed``
write (the documented off-by-one site) stays in the scheduler.

Queue the resolved savepoint restore + bump ``prompt_consumed``.

Order-critical: the caller invokes this AFTER
``alloc_recurrent_state`` (the slab row must exist before we can
write into it) and BEFORE the request enters ``Scheduler.waiting``
(so the next start-of-step flush sees the queue entry). The only
``prompt_consumed`` mutation here is the savepoint-coverage bump
(savepoint-specific, not the radix off-by-one).

Snapshot the recurrent state at ``pc_after`` IFF this step
crossed a ``chunk_size`` grid line and ``pc_after`` is page-aligned.

Called from the prefill branch of :meth:`Scheduler.commit` AFTER
the ``prompt_consumed`` accounting (the off-by-one site, which
stays in the scheduler) has produced ``pc_before`` / ``pc_after``.
A snapshot failure is opportunistic and MUST NOT break the prefill
commit — the store's ``put`` already returned False and bumped the
metric.

THE ONLY STATE THAT EXISTS IS THE STATE AT ``pc_after``
=========================================================
The recurrent slab row holds exactly ONE state: the state after the
chunk that just ran, i.e. after ``pc_after`` tokens. The recurrence
for the whole chunk is swept inside a single fused kernel — no
intermediate token position is ever materialized in the row. So a
snapshot is always labelled ``pc_after``, never a grid point the
step swept through.

WHICH ``pc_after`` IS WORTH A WRITE
===================================
``chunk_size`` is the write-rate governor, not an address: the
step that carries ``prompt_consumed`` across a multiple of it
writes one snapshot, wherever it lands — and so does the step that
completes the prompt (``completes_prompt``, the one write the
ration must not skip: see
:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_fold_split`).
Under concurrency a
chunk is routinely narrower than ``chunk_size`` — the step budget
is what is left after the rows ahead of it in ``running`` (one
token per decode row, and the prompt-completing tail of the
prefill before it), see ``Scheduler.schedule`` — so the request's
chunk grid is offset from the ``chunk_size`` grid from its first
chunk on and never realigns. Keying the snapshot by its true
token count makes every crossing addressable, and snapshotting
only on crossings keeps the rate at one write per ``chunk_size``
tokens of prompt regardless of how the steps were cut.

WHICH ``pc_after`` CAN BE RESTORED
==================================
Only a page-aligned one. A resume credits ``prompt_consumed`` for
exactly the tokens the slab state absorbed, against KV the radix
tree aliases in whole pages; the page holding an unaligned
position is private to this request and can never be aliased.
``Scheduler._page_align_prefill_end`` trims a truncated hybrid
chunk to end on the page grid for that reason; a step that still
lands unaligned (the prompt-completing chunk, or one shorter than
a page) is not snapshotted, and a later admission resumes from an
earlier honest one.

EVERY DROP IS COUNTED AND SIZED
===============================
``arbi_serve.savepoint.snapshot_skipped_total`` records why a
step wrote nothing, and for a step that crossed the grid but
landed off the page grid, the chunk's width and where it landed
(bucketed, so the label cardinality stays bounded). The drop
used to be silent, which is how the trigger rate was measured
three times and explained zero times: the rate is a one-bit
shadow of ``pc_after``. ``width=full`` with a nonzero residue
names an offset grid rather than a narrowed chunk; a short
``width`` names the step that narrowed it.

Store the MID-step boundary snapshot, if the forward staged one.

Returns True when a snapshot was stored — the caller then counts no
drop. False means the write did not happen, and every path to False
names itself in
``arbi_serve.savepoint.snapshot_skipped_total{reason=...}``: this is
the write path of a 147 MiB-per-write ring, and a silent decline here
is exactly the shape of defect #2130 exists to prevent.

THE LABEL AND THE STATE COME FROM THE SAME DECISION
===================================================
The boundary is recomputed here with
:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_fold_split`
from THIS step's own ``pc_before``/``pc_after`` — the same call the
scheduler made when it armed the split — and the result is then
checked against the boundary the snapshot itself carries. Recomputing
and then agreeing is not redundant: the snapshot's buffers were filled
by a forward that ran against an armed plan, and the only thing that
can prove the plan the forward executed is the plan this commit is
labelling is the two numbers matching. They disagree if a step was
skipped, if the pool handed back a stale snapshot, or if the request
this commit is for is not the request the forward split. Any of those
would put a coverage label on a state that is not the state at that
token, which the store's own guard
(``recurrent_savepoint.py`` ``covered == key.num_tokens``) cannot
catch, because both halves would be consistent and both wrong.

A snapshot that is not COMPLETE is refused with the attributes nobody
staged: a partial restore keeps a stale tensor for the layers that
were missed, which is silently wrong output rather than a miss.

Take a recurrent-state snapshot of the state after ``boundary`` tokens.

Called from :meth:`snapshot_crossed_boundaries` for the step that
just ended at ``boundary``. Looks up the request's slab row,
clones each per-(layer_idx, attribute) tensor slice into host
RAM, digests the ``boundary``-token prefix, and ``put``s the
resulting :class:`SavepointEntry` into the store.

ASYNC D2H.
==========
The snapshot is issued via ``snapshot_recurrent_row(async_copy=True)``
+ ``snapshot_recurrent_event()``. This issues one
``non_blocking=True`` D2H per (layer × attr) into a pinned
host buffer and records a single CUDA event on the current
stream. The commit critical path returns BEFORE the D2H is
complete on the host. The event is attached to the
:class:`SavepointEntry`; the store's ``get`` / eviction paths
call ``entry.synchronize()`` to barrier on it before reading
or freeing the host buffers.

Why this is correct without a chunk-N stream sync:

  * The engine's slab writes (recurrent kernels in chunk N's
    forward) and our D2H copy are on the SAME stream
    (``torch.cuda.current_stream(pool.device)``). CUDA
    stream-ordering guarantees the writes happen-before the
    D2H — we don't need a host-side fence between them.
  * Chunk N+1's forward is also on the same stream and
    happens-after the D2H in stream order, so the slab is
    free to be overwritten without affecting the D2H source
    (the GPU will not start chunk N+1's writes until the
    D2H read of chunk N's content has been queued).
  * The host bytes are guaranteed to be present only after
    ``cudaEventSynchronize`` — which the consumer side
    (admission-time restore, eviction) performs.

Falls back to synchronous on CPU-only pools (test fixtures);
``snapshot_recurrent_event`` returns ``None`` and the entry's
``pending_event`` stays ``None``.

Returns silently on a no-recurrent-pool model (``snapshot_fn``
returning an empty dict — we don't ``put`` empty entries).

Continuous-batching scheduler with chunked prefill.

Talks to :class:`MultiStatePool` exclusively — never reaches into a
specific :class:`StatePoolView` impl. The scheduler uses the
multi-pool ``free_pages(StateKind)`` and
``bytes_per_token_total()`` API; the page table handles per-pool
alloc / free internally.

Per step:
  1. Drain a batch from running (decoding) requests, capped at
     ``max_batch`` and ``max_batched_tokens``.
  2. Top up from waiting (prefilling) requests, contributing
     ``min(remaining_prompt, chunk_prefill, budget_left)`` tokens
     each.
  3. Hand the engine a list of ``(request, num_tokens)`` tuples; the
     engine materializes the :class:`ScheduledBatch` and runs
     forward.
  4. After forward, the engine calls :meth:`commit` to advance per-
     request state.

Preemption: when a decode-phase request needs a new page and the
PAGED_KV pool is full, the most recently-admitted decode request is
bumped back to WAITING (its pages free, its prompt + so-far-decoded
tokens become the new prompt).

FCFS continuous-batch scheduler with chunked prefill.

Pool-aware via the :class:`MultiStatePool` interface only.

The scheduler owns the page-table for exactly one paged state kind —
the model's "primary paged" kind. For Qwen3-style dense / hybrid
models that's :attr:`StateKind.PAGED_KV`; for DeepSeek MLA models
it's :attr:`StateKind.MLA_SHARED`. The kind is wired in at
construction (``primary_paged_kind=pool.primary_paged_kind``); the
scheduler never branches on which kind it has.

The armed activation gate, or ``None`` while it is inert.

Readable so the engine can re-arm the SAME model against a re-measured
budget when the card contradicts the one it was armed with
(:func:`~arbi_serve.engine.activation_admission.rearm_step_budget`),
rather than a second copy of the calibration being kept somewhere else
to read the first one back from.

Arm (or clear, with ``None``) the admission-time activation gate.

Called from the boot seam that finalizes the serving floor, so the
budget admission enforces and the bytes the floor reserved are the same
number — and again from the memory-pressure seam, which re-arms the same
calibration against a freshly measured budget after an OOM proved the
boot's reading stale. Until it is called the gate is inert: the page
tally alone governs admission, which is the behaviour every engine had
before.

KV tokens tkv's padded gather can materialize for this prefill row.

``req.total_length`` is the row's accumulated KV before this chunk and
``want`` its query tokens, so their sum is the ``S_kv`` the gather
would span — clamped by the model to the ceiling past which tkv takes
the zero-copy varlen route and stages nothing context-scaled. Two int
reads and a min; zero when no model is armed or the engine has no
paged tkv backend.

Flat tokens this decode row takes from the step's token allowance.

``1`` — its committed tail — on every boot that serves a mixed step as
two forwards, where the verify rows' ``K`` drafts run in a forward of
their own and never share the chunk's width. With
``ARBI_MTP_FUSED_MIXED`` armed the verify rows and the chunk are ONE
forward, whose width must stay inside the step width the activation
profile measured (:mod:`arbi_serve.runtime.fused_mixed_step`), so the
row debits its whole ``1 + K`` and the co-admitted chunk is sized
against what is left. The same read of ``K`` the page side makes
(:meth:`~arbi_serve.scheduler.scheduler_allocation._AllocationMixin.row_draft_depth`)
— an upper bound the step can collapse below, never exceed.

EXTRA query tokens this decode row's step will run beyond its tail.

``K`` for an MTP-opted row, 0 otherwise. A served verify step stages
``K + 1`` flat tokens for the row (:func:`~arbi_serve.spec_decode.
mtp_verify_plan.build_verify_plan`); the ``+1`` is the committed tail
every decode row already pays for through the model's per-row term, so
the charge is the ``K`` drafts beside it.

THE SAME READ THE PAGE SIDE MAKES, literally — both halves of one step's
reservation call :meth:`~arbi_serve.scheduler.scheduler_allocation.
_AllocationMixin.row_draft_depth`, so they cannot come to disagree about
which rows speculate or how deep. The only thing this adds is the gate's
own precondition: an unarmed model charges nothing, while the pages are
reserved either way.

A ZERO SLOPE SKIPS THE READ, and that is an identity rather than a
shortcut: ``demand_bytes`` multiplies this by ``per_query_token_bytes``
and ``max_prefill_tokens`` subtracts the same product, so at a zero
slope every caller's answer is unchanged whatever this returns. The
slope is zero for every boot that does not speculate — the profile emits
no verify shape and the calibration has no residual to fit — so a
non-MTP decode step pays one attribute compare per row and not a
``getattr`` chain, and only the traffic the term exists for pays for it.

Trim a truncated hybrid prefill chunk so it ends on the fold grid.

THE GRID IS THE KERNEL'S, NOT THE PAGE TABLE'S
==============================================
A recurrent savepoint is restorable only where the state it names is
the state an uncut prefill would hold at that token. The recurrent
kernel reduces over its own fixed chunk width (``GDN_CHUNK_SIZE``),
so any decomposition whose every cut is a multiple of that width
yields the same state, and one whose cuts are not yields a state no
uncut prefill produces — the precondition
:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_resume_grid`
already states, applied to the chunk cuts that create the boundaries
rather than to the boundary alone. Because each chunk starts where
the last one stopped, ONE off-grid cut puts every later position of
that prefill off the grid for good.

This used to trim onto the KV PAGE grid, because the snapshot could
only be taken where the step ended. It no longer has to be: the
forward splits the fold at the page boundary INSIDE the step and the
snapshot is taken there (see
:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_fold_split`).
What is left for the chunk's end to satisfy is only the bit-identity
precondition above — the fold grid, which the page grid is a
multiple of. The cut therefore gives up at most ``fold_grid - 1``
tokens instead of at most ``page_bs - 1``.

GOING BACK TO THE PAGE GRID IS NOT THE ANSWER, AND IT WAS MEASURED
=================================================================
While the mid-step emit did not run on the COMPILED prefill path
(arbicity/arbi-serve#2238, before the ops read the plan host-side —
:mod:`arbi_serve.cache._fold_emit_staging`), "cut onto the page grid
so the end-of-step snapshot can take the crossing" looked like the
obvious workaround. It is worse, and stays worse with the emit
working: measured on the served 27B, 6 sessions x 4 turns, same boot
shape both arms, mean turn-2+ cached fell from ~65% (2816-8960
tokens) to ~41% (2048-8192), ``off_page`` rose 35 -> 50, and
``savepoint_fold_split_armed{this request's fold history left the
grid}`` rose 2 -> 21. Giving up to ``page_bs - 1`` tokens narrows the
step below the ``chunk_size`` write governor often enough that
crossings are LOST (``no_crossing`` 30 -> 35), and the page trim is
refused as costing a step far more often than the fold trim is, which
forfeits the grid for the whole request.

Trims only when the trimmed step still covers at least one fold
chunk: a step too small to reach the next grid line keeps its full
width, because progress outranks a snapshot the next wider step will
take anyway. Never trims the prompt-completing chunk — its end is
the prompt length, which is not a choice.

AND ONLY WHEN THE TRIM COSTS NO STEP
====================================
The trim gives up ``want % grid`` tokens, and those tokens are free
exactly when the prefill still finishes in the number of steps it
would have taken uncut. The predicate is that, evaluated on this
step's own numbers::

    1 + ceil((remaining - aligned) / want) <= ceil(remaining / want)

— take the narrower step now, then full ones, and see whether the
total is still the uncut total. When it is, the trim is free and
better than free (a narrower step is a shorter step, and it occupies
the same number of GEMM tiles because the GEMM pads M to whole ones).
When it is not, the trim buys a whole extra step, which is hundreds
of milliseconds against the few the narrower step saves, so it is
refused.

WHY THE PREDICATE LOOKS AHEAD RATHER THAN AT THE WHOLE PROMPT
=============================================================
The obvious form — cut only when ``ceil(remaining / aligned) ==
ceil(remaining / want)``, i.e. only when cutting EVERY remaining step
is free — declines on the FIRST step of any prefill whose length does
not permit it, and a decline is not local: it sets
:attr:`Request.fold_history_off_grid`, which forfeits every later
savepoint of the request. Declining early therefore throws away the
savepoints of a whole prefill to save a step at its end. The
look-ahead form declines as LATE as possible instead, which costs the
same number of steps and keeps substantially more savepoints
(simulated over the served widths; figures on
arbicity/arbi-serve#2189).

Nothing here is a constant: ``want`` is what the step's budget left
after the rows ahead of it, and ``grid`` is the kernel's own reduction
width read through ``recurrent_fold_grid``.

The caller charges the step the UNTRIMMED width. Handing the trimmed
tokens to the next prefill row would turn the depth-first FCFS slate
(one prefill row advancing at full width, see ``schedule``) into two
rows sharing the step: measured on the served 27B at c=8 that moved
TTFT p50 up by more than half while p95 fell — a redistribution, not
a saving. Charging the full width keeps the slate's shape.

There is no flag. The cut was a flag while it was a TRADE — TTFT for
the savepoint hit rate — and the predicate above is what ended the
trade, so a switch would only offer an operator the losing side of a
decision already made from the step's own numbers. What replaces it is
the invariant it was standing in for, asserted directly:
``test_the_cut_never_costs_a_step_at_any_length``.

0 (no store, or a pure-attention model) leaves ``want`` untouched. It
costs nothing at c=1 in any case: a prefill running alone takes the
whole ``max_batched_tokens`` chunk, which is a multiple of both grids
already, so ``end % grid`` is 0 and nothing is trimmed.

Arm the forward's fold split for the one prefill row that needs it.

Called ONCE per ``schedule()``, over the finished slate, because the
numbers it needs are the step's final ones: a row's chunk start
(``prompt_consumed``) and the width the budget and the fold-grid cut
left it. Deciding earlier would key on a width three later caps can
still narrow.

:func:`~arbi_serve.cache.recurrent_savepoint.savepoint_fold_split`
answers WHERE, and it is the same call
:meth:`~arbi_serve.scheduler.savepoint_admission.SavepointSnapshotter.
snapshot_crossed_boundaries` makes at commit — the forward and the
label must never come from two derivations. Arming opens the host
buffers the forward stages into
(:meth:`~arbi_serve.cache._recurrent_lifecycle.RecurrentLifecycleMixin.
open_boundary_snapshot`) and records the offset the split lands at,
for the metadata builder to hand the GDN layers.

Every refusal is named and counted. A zero refusal count on
``savepoint_fold_split_armed`` means this method never ran, which is
a different and earlier failure than "the step had no boundary" —
the counter's window is SERVE, and it moves on every prefill step.

``want`` narrowed to the prefill tokens the step's budget covers.

The token count is the one term of the step's cost admission can move
without refusing the row: a prefill row advances at whatever width fits,
so a card with less free VRAM than a full-width chunk serves the prompt
in narrower chunks instead of OOMing the forward on it. Every other term
the model prices — the rows, the padded gather — is fixed by the slate
the caller is building.

The gather is charged at the UNNARROWED width, so the cap it yields is
the conservative one: narrowing can only lower the row's own gather.

Returns at least 1 on an empty slate. One row is the minimum unit of
progress, and a row too expensive even at one token is the case
:meth:`_activation_admits` names and admits with a warning — the same
boundary, reached the same way, rather than a second policy here.

Whether this slate shape fits the step's activation budget.

Mirrors :meth:`_can_allocate_reserved`, one dimension over: the caller
passes the tally INCLUDING the row it is about to append, and a False
return means the same thing a page-gate refusal means — leave the row
for a later step.

The one asymmetry is the solo case. A row that does not fit the budget
even alone can never fit, so refusing it would starve it forever; it is
admitted with a bounded warning instead. A prefill row reaches that case
only after :meth:`_activation_chunk_cap` has already narrowed its token
count as far as the budget allows, so what remains is a row whose cost
does not live in its tokens. That is the honest boundary of what
admission can enforce: it bounds CO-RESIDENT demand exactly, narrows the
one term it can narrow, and names a single row that exceeds the budget
by itself rather than letting the forward discover it as an OOM.

Return ``[(req, num_tokens_this_step), ...]`` for one step.

``_now`` — explicit clock injection for the duplex-lane deadline
check below (mirrors the ``_now`` convention already used by
:mod:`arbi_serve.scheduler.multi_group`). Two real callers now:
a test injecting a synthetic clock, and ``run_forever``'s
duplex pre-schedule hook, which must evaluate
:meth:`duplex_tick_due` and then call this method against the SAME
clock reading — otherwise the deadline could cross between the two
calls and the hook would prepare a frame for a tick that does not
fire (or, worse, not prepare one for a tick that does). Ordinary
callers omit it.

The default clock is :func:`time.monotonic`, NOT :func:`time.time`:
the duplex frame grid is a real-time media cadence, and the wall
clock is steerable. ``Request.next_duplex_deadline`` is an absolute
instant on this clock, so an NTP step backwards parks every live
duplex session for the size of the step (every deadline is suddenly
in the future, and only a timer can fire the lane), while a step
forwards resyncs the grid. Neither is hypothetical on a session
whose configured ceiling is ``duplex_max_session_s`` and on hosts
where chrony slews continuously. Every reader of that field must
use the same clock; ``run_forever``'s idle-park cap is the only
other one. Design doc §7.18.

Pool-aware admission mixin (add / remove / recurrent-row self-heal).

Owns the request-lifecycle entry / exit points: :meth:`add` (total,
never raises), the recurrent-row deferral self-heal, and
:meth:`remove`.

Shared attributes are declared as class-level annotations (mirrors
:mod:`arbi_serve.cache._pool_base`) — the real values are set in
:meth:`Scheduler.__init__`.

Internal signal: re-admit this request at the next ``schedule``.

Raised by :meth:`Scheduler._add_inner` when the recurrent-row
self-heal would need to park a resident batch job while a model
forward is in flight (unsafe CUDA mutation). Caught by
:meth:`Scheduler.add`, which re-queues the request for between-
forwards re-admission. Never escapes the scheduler.

The furthest context ``req`` can ever reach: prompt + output budget.

``None`` when the request carries no ``max_tokens`` — it then runs until
a stop matches or the engine's max context does, so the pool resolves the
target against its own reservation. A duplex-lane request is ``None`` for
the same reason from the other side: its context also grows by the tokens
the lane injects (``pending_context_token_ids``), which no admission-time
sum can bound. This is the growth target for a state pool that backs its
stores on demand: growing to the request's CURRENT length instead would
re-map on every token.

True iff the bound pool exposes a recurrent / short-conv view.

Cached lazily on first call (admission is hot); defaults to False
on mock pools without the predicate. Includes ShortConv (LFM2) as
recurrent for the radix-cache gate — it carries the same
per-request slab-row state GDN/Mamba do, which the token-id-only
radix key can't capture.

Enqueue a freshly-arrived request.

Wrapper over :meth:`_add_inner` that is **total** — it NEVER lets
an exception escape. Two failure classes are handled:

* :class:`_DeferAdmission` — the recurrent-row pool is full and a
  row cannot be reclaimed THIS instant (no parkable batch job, or a
  forward is in flight so the self-heal's CUDA suspend is unsafe).
  The rolled-back request (holds no GPU state) is re-queued on
  ``_deferred_admits`` and re-attempted at the top of a later
  :meth:`schedule` — strictly between forwards — once a running
  request finishes and frees a row. This is the
  queue-fairly-and-resume path for a burst whose concurrency
  exceeds the recurrent-row budget (``max_num_seqs`` on a hybrid
  GDN / Mamba model). The deferred admit is counted by the
  admission depth gate, so a SUSTAINED overload still backs off
  with a clean 429 at ``queue_depth_max`` instead of growing the
  deferred queue without bound.

* Any OTHER exception (e.g. a savepoint best_match→get eviction
  race, an unexpected pool error) — finished cleanly via the
  dropped-job terminal path the run loop drains
  (``run_step.drain_dropped_suspended`` → ``on_finished``). Without
  this, the exception escapes into the cross-loop intake
  trampoline (``engine.loop_bridge.invoke_isolated``) which SWALLOWS
  it: the request would never finish, never stream a token, and the
  client would hang until its own (~300 s) timeout. Admission must
  never hang a request.

True iff any interactive request is queued, running, or deferred.

Mirrors the slate-admission gate's notion of interactive presence
so the row layer and the scheduling layer agree on when batch may
make progress. Deferred admits count: an interactive request that
rolled back for want of a row is still interactive demand.

Undo the page-table + per-request side effects of a partial
:meth:`_add_inner` so the request can be cleanly re-admitted.

Mirrors :meth:`remove` minus the queue removals (the request was
never queued yet): releases radix refs + private pages, clears the
state-handle pointer, resets prefill credit, restores ``WAITING``.

Rebind a never-scheduled request to newly published prefix pages.

Prefix-serial followers were registered when they arrived, before an
active leader published its chunks. Re-running the ordinary admission
transaction here refreshes both radix pages and recurrent savepoints.
This happens between forwards and before the request's first slate, so
SPMD's one-time AdmitRow observes only the refreshed state.

Give a WAITING request's matched prefix pages back to the pool.

A prefix match ref-s the pages it aliases, so they leave
``available_pages()`` and stop being LRU-evictable. ``_add_inner``
takes that match the moment a request ARRIVES, while it is still
queued — so a queued request with a long cached prefix holds KV it
cannot use against every request that could actually run, and the
scheduler has no lever against it: preemption's victims are running
rows, and this one is not running. When the held prefix is a large
fraction of the pool, the row that IS running cannot reach the end
of its own prompt, is dropped from the slate, is preempted for the
next queued request, and restarts — the same prefix chunks forwarded
forever with nothing finishing.

:meth:`Preemptor.preempt_for_space` already states the invariant this
restores: only rows that are running hold pages. It keeps it for a
preempt victim by going page-less in ``waiting`` and claiming the
match again at admission; this does the same for a queued request
whose hold admission cannot honour — the pages drop to ref 0 (still
in the tree, still matchable, now reclaimable by whoever needs them
first) and the request re-claims what survives when it is admitted,
through the same re-bind seam and the same per-step page budget.

Returns the tokens released (0 when the request holds no match).

Re-query the prefix cache for a preempted request being admitted.

``Preemptor.preempt_for_space`` commits the victim's completed
pages into the radix tree and then releases them: they stay in the
tree at ref 0, so at this instant every page the victim is about to
recompute is still there. Nothing in the waiting-admit path looks —
the only ``match_prefix`` site is :meth:`_add_inner`, which a
preempted request never re-enters — so the victim re-prefills from
token 0 and then dedups its recomputed pages onto the ones it just
recomputed. Full cost paid, nothing saved.

Called at the point of ADMISSION, deliberately, not at the preempt:
the match ref-s the pages back up and they stop being evictable, so
a victim that took them while still in ``waiting`` would hold the
pool against every request that could actually run — a state the
scheduler cannot escape, since preemption is the only lever it has
and its victims are the ones holding the pages. Here the request
goes onto this step's slate, so it holds pages exactly like any
other running row. ``test_a_tight_pool_still_makes_progress``
(``tests/test_preempt_prefix_reuse.py``) is the regression.

Runs the SAME transaction a fresh admission does —
:meth:`SavepointSnapshotter.resolve_resume` /
:meth:`~SavepointSnapshotter.resolve_match` /
:meth:`~SavepointSnapshotter.queue_resume` — rather than restating
any of it, so the hybrid rules hold identically: the match is
capped at the recurrent savepoint's coverage, a coverage shortfall
drops the savepoint, and a hybrid match with no savepoint
force-disables the cache instead of decoding against zero GDN
state. ``readmit=True`` is the only difference: the request is
already registered (page-less), so the registration must go through
the free-and-re-register seam that carries its cache identity.

Two caps, both measured, neither a constant:

* one token below the rebuilt prompt, so the row this admission is
  building always has something to feed — a FULL match belongs to
  the ``num_remaining_prompt == 0`` branch, which this call site is
  already past;
* ``max_new_pages``, the pages this step can still spare. A matched
  page moves from evictable to ref-ed, so it leaves
  ``available_pages()`` — the same quantity every slate row's
  allocation gate reads. The caller computes it as availability
  minus what the rows already on the slate reserved minus what this
  row still needs, so the re-bind can never take a page a slate-mate
  was already promised — without it an earlier decode row reaches
  ``allocate_slots`` to find the pool gone. ``None`` leaves only the
  first cap, for page tables that model no availability.

Returns the credited ``match_len`` (0 on a miss, on the flat page
table, or when the request holds no rebuilt prompt).

Re-admit requests deferred by the offload self-heal.

Called at the top of :meth:`schedule` (loop thread, strictly
between forwards) so the recurrent-row self-heal's suspend (CUDA
D2H + page free) runs with NO forward in flight. A request that
STILL can't get a row is re-deferred; bounded by the queue snapshot
so a request deferred this tick isn't retried in the same drain.

Enqueue a freshly-arrived request.

With :class:`RadixPageTable`, ``add_request`` returns the
cached-prefix length matched at admission; we credit those tokens
by advancing ``prompt_consumed`` so prefill skips them
(:class:`FlatPageTable` always returns 0 — the credit is a no-op).

Hybrid models (GDN, Mamba, LFM2 ShortConv) also get a per-request
recurrent slab row reserved here via
:meth:`MultiStatePool.alloc_recurrent_state`; without it the
metadata builder would alias request N's row to request N+1 and a
radix-cache hit would decode against stale recurrent state.
Pure-attention models see a no-op.

Page-allocation / preemption reservation-accounting mixin.

Owns the per-step page-commit accounting helpers
(``_decode_page_commit`` / ``_can_allocate_reserved`` /
``_preempt_until_allocatable``) plus the thin ``_preempt_for_space``
delegator to the :class:`~arbi_serve.scheduler.preemption.Preemptor`.

Shared attributes are declared as class-level annotations (mirrors
:mod:`arbi_serve.cache._pool_base`).

This row's resolved MTP draft depth ``K``. THE read of it.

One step's reservation for a speculating row has two halves — the PAGES
its ``1 + K`` extension will take (:meth:`_decode_page_commit`) and the
activation BYTES its ``K`` extra query tokens will cost
(:meth:`~arbi_serve.scheduler.scheduler.Scheduler._query_charge`) — and
they are reservations for the same tokens in the same forward. Two
spellings of the same read is how the halves come to disagree about
which rows speculate, and a slate whose pages are reserved for a verify
step and whose bytes are not is exactly the over-admission each half
exists to prevent.

An UPPER bound by construction: ``_bucket_mtp_k_uniform`` runs after the
slate is built and can only lower a row's ``K`` (it buckets the step to
one uniform depth, and the spec valve can zero it outright), so a row
reserved here at its own ``K`` never runs deeper than it was reserved
for.

New pages appending this row's ``n_tokens`` will grab.

``n_tokens`` defaults to ``1`` — the ordinary decode row, which
grabs 0 or 1 page. A duplex context-injection row passes its full
``injection_row_width`` so the reservation covers every token the
row's ``allocate_slots`` will take; reserving only the first would
under-reserve and let the pool hit true exhaustion mid-step.

A decode row grabs a fresh page iff its incoming token starts a new
page. ``req.total_length % block_size == 0`` sits ONE tick AHEAD of
the page-table length: the just-sampled token is appended to
``output_token_ids`` at sample time, while its page is not allocated
until the NEXT step's ``allocate_slots`` in ``_gather_slate``. So
``schedule()`` sees ``total_length = page_table.length + 1`` for
every decode row, one step ahead of the true page-table crossing (at
``page_table.length % bs == 0``).

Ask the page table directly instead: ``pages_needed`` mirrors
``allocate_slots``'s arithmetic EXACTLY off the page-table length
(and correctly credits any radix-shared prefix). Falls back to the
page-table ``length`` (flat impl / stubs — no radix sharing there,
so ``length % bs`` is the exact predicate) and finally to the
``total_length`` predicate for minimal stubs that expose neither.
KeyError-safe: a not-yet-registered row estimates 1 page.

MTP rows reserve the VERIFY-step peak. An MTP-opted decode row's
step allocates a tail slot PLUS ``mtp_k`` draft slots (``1 + K``
tokens — ``derive_verify`` / ``build_verify_plan``), not the single
decode token; committing only the tail's 0-or-1 page UNDER-reserves
by up to the drafts' page demand per MTP row, so a saturated c8
slate over-admits and the pool hits TRUE exhaustion mid-step.
``draft_pages_needed`` models the ``1 + K``
extension exactly, on the flat page table and the radix one alike.
Rejected drafts are trimmed back at ``finalize_draft_slots``, so
this is a transient peak — reserving it costs at most admission
backpressure, never a crash.

``can_allocate(k)`` minus this step's already-committed pages.

``schedule()`` appends N rows per step but ``allocate_slots`` only
runs later in the forward — so the page table's free count is NOT
decremented as rows are appended. This helper subtracts
``reserved_pages`` (the pages the rows already ON the in-progress
slate will consume) from the availability check, the same
running-tally pattern ``token_budget`` / ``batch_left`` use, so a
row is never admitted against pages a slate-mate already reserved.

``reserved_pages <= 0`` delegates to the page table's own
``can_allocate`` byte-identically (also keeps minimal test stubs
without ``available_pages`` working).

Free pages held by QUEUED rows' prefix matches, until ``want`` fits.

The cheaper half of the two levers admission has for a full pool.
Preemption takes pages away from a row that is making progress;
this takes them from a row that is not running at all — a queued
request whose admission-time prefix match ref-ed pages it cannot
use yet (see
:meth:`~arbi_serve.scheduler.scheduler_admission._AdmissionMixin._release_waiting_prefix_hold`).
Reclaiming those first is what keeps a long queued prefix from
starving the row that could actually finish and free the pool.

Released from the BACK of the queue: the request furthest from its
own admission gives its hold up first. Returns whether ``want``
fits after the reclaim — including the no-op case where it already
did, so a caller can gate on this before reaching for a victim.

Preempt victims until ``can_allocate(want)`` holds (or no victim).

A single ``_preempt_for_space`` attempt is not enough here: one
victim's freed pages can be far short of ``want`` (a 2047-token
prefill continuation needs ~8 pages; a radix victim whose pages are
shared may free none). Loop-and-recheck gives the same guarantee
the waiting branch gets from its ``continue`` retry: either the
demand is genuinely allocatable when the row is appended, or the
caller breaks/defers cleanly. Terminates: every successful preempt
removes one victim from ``running``.

``reserved_pages`` threads the caller's per-step page-commit tally
(see :meth:`_can_allocate_reserved`) so the loop targets enough
free pages for BOTH the already-appended slate rows and this one.

The victim order is built ONCE and consumed across the loop: it is
stable under the removals the loop itself makes, so rebuilding it
per victim re-derives the same sequence at a cost quadratic in the
running set (see :meth:`Preemptor.victim_ordering`).

:meth:`_reclaim_idle_prefix_holds` runs FIRST: pages held by a row
that is not running are cheaper to take back than pages held by a
row that is making progress, and a queued prefix hold large enough
to starve the running row is the case where preemption has no
victim to offer at all.

Post-step commit + terminal-finish mixin.

Owns the prefill→decode state advance (:meth:`commit_state`), the
deferrable output append (:meth:`commit_output`), the synchronous
:meth:`commit`, and the terminal :meth:`finished` hook.

Shared attributes are declared as class-level annotations (mirrors
:mod:`arbi_serve.cache._pool_base`).

Advance per-request state + append sampled tokens after a step.

Synchronous (non-async-output) path: the two halves run back to
back. Under async output the run loop splits them — :meth:`commit_state`
runs SYNCHRONOUSLY when the deferred step is parked (so the next
step schedules the request as decode, not a re-prefill), and
:meth:`commit_output` runs one engine tick later when the deferred
sample is materialized. See :func:`run_step.drain_pending_output`.

Advance the prefill→decode state machine for a step's slate.

Bumps ``prompt_consumed`` for prefill rows, snapshots crossed
recurrent-savepoint boundaries, and on the prompt-completing chunk
promotes the request to ``DECODING`` + promotes its full prompt
pages into the radix cache. Does NOT append output tokens — that's
:meth:`commit_output`.

Also drains ``pending_context_token_ids`` into ``context_consumed``
for any row the batch builder served as an injection row, so the
injected context advances ``total_length`` without ever entering
``output_token_ids``.

This MUST stay synchronous even under async output: ``is_prefill``
keys off ``prompt_consumed`` and the next step's positions key off
``total_length``, so deferring either advance would leave the
request looking mid-prefill (re-processing the prompt, duplicating
the first decoded token) or re-injecting context onto KV slots it
already owns. Idempotent per step: only prefill rows advance, a row
leaves prefill once consumed, and the context buffer is cleared as
it is counted.

Append each row's sampled token to ``output_token_ids``.

The deferrable half of :meth:`commit` — runs at sample-materialize
time (one tick late under async output). A ``None`` token (mid-
prefill chunk, or a row dropped because it finished mid-flight) is
skipped. ``commit_state`` already ran for this slate.

Called by the engine when stop / length / context limit hits.

Commits any decode-completed full pages into the radix cache BEFORE
releasing the request's page refs (the new shared nodes land at
``ref_count==0`` / evictable, so follow-up prompts covering the same
generated tokens hit them). Cheap: one tree walk +
``output_tokens / block_size`` insertions max.

O(1) admission-counter maintenance + engine-wiring hooks mixin.

Provides the four ``waiting`` / ``running`` deque helpers that keep the
priority-split admission counters exact, the debug-only drift
reconciler, the engine-injected predicate setters, and
:meth:`shutdown`.

The shared attributes they read / write are declared as class-level
annotations (mirrors :mod:`arbi_serve.cache._pool_base`) — the real
values are set in :meth:`Scheduler.__init__`.

One logged preemption, in the shape the SPMD delta broadcasts.

Field-for-field the payload of
:class:`arbi_serve.distributed.spmd_delta.PreemptRow`; the driver's
delta build is a straight copy. It lives here rather than being
rebuilt at drain time because every field is a rank-0 DECISION, taken
at a moment a later scan cannot reconstruct: the victim is routinely
re-admitted onto the same tick's slate, after which nothing
distinguishes it from an ordinary prefill row.

Remove ``req`` from ``dq`` by object identity; report whether it was there.

``deque.remove`` compares with ``==``, and :class:`Request` is a
dataclass whose generated ``__eq__`` materializes a tuple of every
field for BOTH operands per candidate — so a miss costs microseconds,
not nanoseconds, and one removal from a deep queue costs more than a
scheduling tick's whole budget. Queue membership is identity (a
request object is in exactly one queue and no two requests carry the
same ``request_id``), so scan for it directly and never call
``__eq__`` at all.

Whether any batch-priority request is queued.

``_n_waiting_interactive`` counts every NON-batch waiting request
(interactive and duplex alike), so the batch population is exactly
the remainder of ``len(waiting)`` — the predicate needs no scan.

Whether any non-batch request is queued or running.

The slate-construction gate behind interactive-always-priority
(:func:`~arbi_serve.scheduler.step_budget.compute_step_budget`),
read off the priority-split counters instead of scanning both
deques. Duplex-lane requests count as non-batch on both sides,
exactly as the scans they replace did.

Reconcile the O(1) admission counters against a brute-force recount.

Cheap O(queue+running) check used ONLY to catch counter drift in
tests / debug builds — a drifted counter would silently break the
admission gate (over- or under-counting interactive demand). Gated
behind ``__debug__`` (off under ``python -O``) and the
``ARBI_ASSERT_ADMISSION_COUNTERS`` env so it never runs on the hot
production path; the load-bearing scheduler test calls it after
every mutation.

Wire the engine's "is a forward in flight?" predicate.

Used by the admission-time recurrent-row self-heal to decide
whether a suspend can run inline or must be deferred to the next
between-forwards :meth:`schedule` tick.

Defer rank-0 finish-side page mutations to the SPMD replay point.

``fn(request_id)`` queues the request on the SPMD driver's
pending-release list; the driver DRAINS it at the top of the NEXT
tick — running the mirror-sourced finish-commit
(:meth:`SpmdRankLoop.commit_full_pages_on_finish`), the
``page_table.remove_request`` free, and the recurrent-slab free at
the byte-exact point a WORKER's ``apply_delta`` evict replay runs
them (start of its next tick: before that tick's preempt replays,
LRU-eviction replays, and per-step allocations). With it set, BOTH
rank-0 finish paths (:meth:`finished` and :meth:`remove`) skip
their immediate page-table / recurrent-pool mutations and queue
instead.

Why DEFERRAL (not just a different commit source) — the c8 radix
divergence is an ORDERING break, two ways:

  1. Rank 0 finish-commits MID-TICK (``post_token`` →
     :meth:`finished` during the verify stream / legacy commit),
     i.e. BEFORE the same tick's later prompt-completion commits,
     while the worker replays the finish at the START of the NEXT
     tick, i.e. AFTER them. Requests race-committing the SAME
     prefix chunks (concurrent same-doc prefill) then dedup in a
     DIFFERENT order per rank → the chunk→page bindings rotate —
     the observed ``add_request_spmd`` broadcast-page mismatch.
  2. Rank 0 FREES the finished request's pages mid-tick: a verify
     tick's legacy sub-pass derive allocates AFTER the verify
     stream, so rank 0's allocation can draw the just-freed pages
     while the worker's byte-identical allocation cannot (its free
     lands one tick later) → pool free-list divergence.

Content also matters: rank 0's Request token list can be SHORTER
than the mirror (the verify stream stops appending at a mid-accept
finish; a FAILED step never commits to the Request at all), and
``remove`` (cancel / timeout) never committed — while the worker's
replay always commits the mirror view. The drain therefore commits
from the SAME replicated mirror the worker reads. ``None`` restores
the single-process behaviour (immediate release, Request-sourced
commit).

Start (or stop) logging preemptions for the SPMD delta builder.

Armed by the SPMD rank-0 driver, which is the only thing that
drains the log. Off elsewhere so a single-process server does not
accumulate an entry per preemption for the life of the process.

Log one preemption for the SPMD delta builder to broadcast.

Called by :meth:`Preemptor.preempt_for_space` the moment a victim's
pages are freed and its decode state reset — the ONE point at which
the decision is unambiguous. A no-op unless
:meth:`enable_spmd_preempt_log` armed the log (i.e. an SPMD rank-0
driver is there to drain it).

``prefix_match_len`` / ``prefix_page_ids`` are the radix re-match
rank 0's re-admit applied, and they ride the wire for exactly the
reason :class:`AdmitRow`'s do: the match is the one non-derivable
decision (its cap depends on rank 0's savepoint store, which a
worker has no copy of), while everything around it is in lockstep.
``cache_enabled`` OVERRIDES the flag the re-admit seam carries off
the victim's own entry; ``None`` (the default) keeps it, which is
right for every preempt. Rank 0 sets it only where its re-bind
resolved a DIFFERENT value — the hybrid force-disable — because a
worker that missed that promotes pages rank 0 did not, splitting
the per-rank free lists. The preempt itself is page-less, so all
three arrive later, through :meth:`amend_spmd_preempt`.

Attach the prefix re-bind's result to this tick's logged preempt.

A preempt is decided in two steps inside ONE ``schedule()`` call:
the victim's pages are freed (:meth:`record_spmd_preempt`), and
then — if the same admission pass puts it back on the slate — it
re-claims what the tree still holds
(``Scheduler._rebind_preempted_prefix``). Both are rank-0 decisions
a worker must replay, and they belong to the same
:class:`PreemptRow`, so the second amends the first rather than
opening a second wire type for it.

Safe because the drain runs after ``schedule()`` returns and the
re-bind can only fire within it (``_preempt_rebind_ids`` is cleared
per call): the entry being amended is always still undrained.
Amends the LAST entry for ``request_id`` — the victim may have been
preempted more than once in a tick, and it is the current residency
being re-bound. A no-op when the log is not armed.

Take (and clear) the preemptions logged since the last drain.

The SPMD rank-0 delta builder drains this once per tick and turns
each entry into a :class:`PreemptRow`, so every worker replays the
SAME page free + re-add + recurrent-slab recycle + mirror reset
that rank 0's :class:`Preemptor` just performed.

This log — not a post-hoc scan of request state — is the source of
truth because a preempted victim is pushed to the HEAD of
``waiting`` (:meth:`_waiting_add` with ``left=True``) and the SAME
``schedule()`` call's waiting-admit pass can re-admit it onto the
very slate the preempt made room for. Such a victim is on the
tick's slate, alive, and mid-prefill again — indistinguishable by
inspection from a request that was simply scheduled, so the scan it
replaced silently dropped its :class:`PreemptRow`.

Sever every reference this scheduler holds to the model's GPU
state, so an in-process model reload can free the KV pool even while
the run loop's suspended frame still pins THIS scheduler object.

The whole-engine reload (:func:`arbi_serve.engine.lifecycle.release_for_reload`)
sets ``eng.scheduler = None``, but the run-loop coroutine is parked
at an ``await`` with a live local ``drain_dropped =
eng.scheduler.drain_dropped_suspended`` bound method — which keeps
alive past ``free_all``. (Full ``shutdown()`` doesn't hit this
because it terminates the run loop, destroying that frame.) Nulling
the pool / page-table / collaborator refs here breaks the heavy path
regardless; the lightweight gutted Scheduler is released on the first
post-rebuild loop iteration, when the loop rebinds its locals to the
new scheduler. Idempotent.

Suspended-job KV offload (batch-backlog parking) delegator mixin.

The offload POLICY lives in :mod:`arbi_serve.scheduler.offload_manager`;
these are thin delegators kept on :class:`Scheduler` so the
``_add_inner`` self-heal and the test reach-ins keep working. The
parked-job deques (:attr:`suspended`, :attr:`_dropped_suspended`) +
snapshot store (:attr:`_suspended_store`) are created in
:meth:`Scheduler.__init__`.

Shared attributes are declared as class-level annotations (mirrors
:mod:`arbi_serve.cache._pool_base`).

Resolve the cache-namespace tenant id for ``req``.

Resolution: ``tenant.tenant_id``, then ``req.tenant_id``, then
``""`` (global namespace — preserves cross-request sharing).

LoRA isolation: a request's KV depends on its active adapter, so
the adapter id is folded into the namespace. Two requests sharing
a prompt but using DIFFERENT adapters (or base vs an adapter)
must NOT share radix / savepoint cache pages — otherwise the
second request decodes against KV computed under the wrong
adapter (silent correctness bug; surfaces as "a zero-delta
adapter changes the output" because it reuses another adapter's
cached prefix). ``lora_id is None`` (no adapter) keeps the
original namespace, so non-LoRA traffic is byte-for-byte
unchanged.

Waiting-queue selection / prefix-serial / MTP-bucketing mixin.

Covers the per-step waiting-candidate picker, the optional
prefix-serial policy, the thin MTP uniform-K bucketing delegator, and
the duplex-lane-only slate builder (design doc
``docs/nemotron-voicechat-duplex-design.md`` §7 step 3).

Shared attributes are declared as class-level annotations (mirrors
:mod:`arbi_serve.cache._pool_base`).

Pick the next waiting request to attempt admitting this step.

Priority-aware, FCFS-within-class selection:

  * Fast path — no batch waiting (``any_batch_waiting=False``,
    all-interactive default) AND the head of the queue isn't a
    duplex-lane request: return ``self.waiting[0]`` verbatim, no
    extra scan.
  * Mixed path — prefer the oldest INTERACTIVE waiting request;
    only when none is waiting consider batch, and then only if the
    interactive backlog sits at/below
    ``batch_admit_waiting_watermark`` (default 0 ⇒ empty queue).

Duplex-lane requests (``is_duplex_frame``) are INVISIBLE to this
picker regardless of path — this method is only ever called from
``Scheduler.schedule()``'s ordinary (non-duplex) tick body, since
a duplex tick short-circuits before reaching it (see the duplex-
lane exclusivity gate at the top of ``schedule()``). A duplex
request is admitted from ``waiting`` exclusively by
``Scheduler._build_duplex_slate`` on a duplex tick, never here.

Returns ``None`` when only watermark-held batch candidates remain
(caller stops admitting; they retry next step). The returned
request is NOT removed from ``self.waiting`` — the caller pops it
by identity once it clears the allocation / budget gates.

Whether the load valve is currently suppressing speculation.

The engine-side truth behind the admin ``server_info`` field and the
``spec-decode DISABLED by load`` log line: ``True`` means MTP-opted
rows are decoding without drafting because the decode batch reached
``mtp_spec_disable_batch``.

The load valve's resolved threshold as of the last admission step
(``0`` = never disable). In ``auto`` this is the smallest width the
valve has MEASURED a stop verdict at, so the number means the same
thing it does under a pinned threshold.

The valve's tri-state as of the last admission step — ``off`` /
``on`` / ``auto``. The number above is not self-describing: an
operator-pinned 2 and an ``auto`` that measured 2 are different
facts about the deployment.

Whether THIS step should serve the duplex lane exclusively.

The single source of truth for the duplex-lane gate's readiness +
deadline question, so the two callers that must agree exactly
cannot drift apart (design doc §7.9):

  * ``Scheduler.schedule()``'s own exclusivity gate, which acts on
    a ``True`` by returning :meth:`_build_duplex_slate`'s slate;
  * ``run_forever``'s duplex pre-schedule hook
    (:mod:`arbi_serve.realtime.duplex_lane`), which acts on a
    ``True`` by setting every due session's
    ``pending_embed_override`` BEFORE ``schedule()`` runs. Both
    must be answered against the SAME ``now`` — the run loop reads
    the clock once and passes it to both — or a deadline crossing
    between the two calls would either feed a frame to a tick that
    never fires (stale override, wrong audio next tick) or fire a
    tick with no frame prepared.

"Ready" is "at least one duplex-lane request exists in ``running``
or ``waiting``" — a duplex session is ALWAYS ready, never skipped
for want of audio: a tick with no client audio yet is fed a
silence frame instead (design doc §2.2's shared 80 ms grid; the
reference's own live-queue fallback), because skipping a tick
would desync the frame-index-based ``agent_idle``/BOS-EOS
bookkeeping from the grid.

Readiness is asked PER SESSION (design doc §7.9.14): the lane is
due when at least one duplex request's OWN
``next_duplex_deadline`` has come due. A request that has never
ticked (``None``) is due immediately, so a freshly admitted
session is not stuck waiting out a deadline nobody armed yet.
Sessions with different cadences — a real-time conversation and a
bounded, unpaced one — therefore coexist on one engine.

Whether any NON-duplex request could be served by an ordinary
slate right now — the fairness valve's second condition (see
:meth:`duplex_lane_should_yield`).

Deliberately "could be served", not "will be": a ``waiting`` row
may still be refused for page room this step. Yielding a step to a
queue that then produces an empty slate costs one 5 ms idle-park
floor (``run_step/loop.py``) and is self-correcting; NOT yielding
because the queue looked unschedulable would be the starvation
this valve exists to bound. Same O(1) ``_duplex_ever_admitted``
fast-out as its siblings, and only ever reached on the cap-th
consecutive duplex step, never on an ordinary one.

Whether this due duplex step should be handed to ordinary
traffic instead — ``batch_cfg.duplex_lane_max_consecutive_ticks``
consecutive duplex steps have already been taken AND there is
ordinary work waiting. See that field's own comment for the
policy and why a paced session can never reach it.

Earliest wall-clock instant any duplex session is next due, or
``None`` when no duplex session exists.

The run loop's idle park needs this. ``duplex_tick_due`` answers
"is a tick due NOW", which is the wrong question for a loop about
to sleep: an empty slate means no session is due *yet*, not that
the engine is idle, and parking past the nearest deadline drops
every frame in between (design doc §7.14). ``0.0`` — a session
that has never ticked, or an UNPACED one — means "due
immediately", so the caller must not sleep at all.

Mirrors :meth:`duplex_tick_due`'s O(1) fast-out and its
``getattr`` guard for the minimal request stand-ins CPU tests use.

``req`` is a duplex session whose own frame deadline is due.

``getattr``-guarded rather than a bare call: minimal request
stand-ins in CPU tests predate ``Request.duplex_tick_due`` and
must keep behaving as always-due duplex rows.

Size this duplex request's next row: ``(n, pages, is_prefill_row)``.

Three row shapes, in the order a session meets them:

  * **Seed prefill** (``req.is_prefill``) — a chunked prefill row
    of ``min(num_remaining_prompt, chunk_prefill, token_budget)``
    tokens, reserving ``ceil(n / block_size)`` pages exactly like
    the ordinary prefill branch of ``Scheduler.schedule``. A
    duplex request's seed prompt is genuinely uncomputed KV (it
    is ``cache_enabled = False`` and has no cached prefix), so it
    gets a * **Context injection** — ``1 + len(pending_context_token_ids)``
    (``injection_row_width``), the tool-response resume row.
  * **Ordinary frame decode** — one token.

``n == 0`` means "no room in this step's token budget"; the
caller skips the row and retries on the next duplex tick.

The one due duplex request whose next row is WIDER than a single
token, if any — that request owns this tick exclusively.

**This is a correctness invariant, not a throughput heuristic**
(design doc §7.9.18). A duplex tick's steady-state rows carry an
embed override, and an override-carrying step is routed onto the
single EAGER whole-batch forward by ``model_runner``'s
``_has_embed_override`` guard — which deliberately skips the
mixed-composition capture rung, the split decode/prefill rung and
every captured-graph replay, because none of them has a seam for a
per-step ``inputs_embeds``. So a duplex step is exactly one forward
over exactly one batch, with no split available to rescue its
composition.

A recurrent (Mamba-2) block dispatches on
``n_tokens == num_rows``: one token per row is the decode path,
anything else is the PREFILL path — for the WHOLE batch. ``1 + B>1 varlen prefill loop.

This isolation is a GUARD, and its cost is unmeasured.
``Mamba2Block``'s varlen loop masks each row's carried state per row
(``models/mamba2_block.py``, the ``row_init`` mask), so a mixed
fresh/continued composition is ordinary there — this guard is not
holding back a live refusal from that block. It is kept because the
blast radius of being wrong is the WHOLE SLATE: a raise inside the
step reaches ``run_forever``'s error branch, which finishes EVERY
row — the total-outage shape §7.9.14 fixed one layer up, reachable
through the same trigger (two duplex connections that did not open
at the same instant, or one tool call resolving while another
session is live). §7.9.14's ``embed_input_ids`` hook stops
``resolve_embed_override`` from raising; it cannot stop a kernel
below it, and no CPU test sees that path because they run a stub
model. Relaxing this therefore needs a GPU run, not a CPU test.

Two compositions are safe and are deliberately still allowed, so
this is the narrowest rule the kernel contract actually implies:

  * **Every row exactly ONE token.** The block takes its decode
    path, which applies a single-token state update per row
    against that row's own slab state — correct for a fresh row
    (zero state) and a continued one alike. So a seed prefill's
    FINAL chunk, when it is one token wide, still shares a tick
    with steady-state rows, and that is still a genuinely MIXED
    override batch served through ``embed_input_ids``.
  * **Every row a FRESH first chunk** (nothing cached for it yet).
    The B>1 varlen prefill loop only refuses ``past > 0``, so two
    connections opening in the SAME tick still prefill together.

Only when a wide row and an already-started row must share a tick
does one of them have to go, and it is the WIDE one that takes the
tick: it is the row that cannot be split, and giving the tick to
the decode rows instead would starve it forever (a session that
gets no row keeps its deadline, so the incumbents would be due
again on every subsequent step too).

Ordering matches :meth:`_build_duplex_slate`'s own two sources
(waiting before running), so a joining session's seed prefill wins
the tick over an already-running session's injection row.

This request's next row starts at sequence position 0, i.e. it
arrives with NOTHING cached (``past == 0`` in ``Mamba2Block``'s own
terms — see :meth:`_duplex_wide_row_owner`).

``getattr``, not bare reads: minimal request stand-ins in CPU tests
predate both fields, and they only ever carry one-token rows, which
never reach this predicate anyway.

Build a duplex-lane-ONLY slate: exactly one row per duplex request.

Row width comes from :meth:`_duplex_row_size` — a chunked seed
prefill row while the request is still consuming its seed prompt,
then ``req.injection_row_width`` (``1`` for an ordinary frame
tick, ``1 + len(pending_context_token_ids)`` when the lane has
queued injected context). Every shape reserves pages for every
token it will allocate, and the lane feeds no audio frame on a
prefill or injection tick.

Called exclusively from ``Scheduler.schedule()``'s duplex-lane
exclusivity gate once a duplex tick is due and ready (design doc
``docs/nemotron-voicechat-duplex-design.md`` §2.2/§7 step 3). No
ordinary (interactive/batch) row is ever considered here — this
mirrors ``resolve_embed_override``'s all-or-nothing-per-tick
invariant (``forward_exec.py:194-204``) as a scheduling policy.

Two sources of duplex-lane rows, both handled here:

  1. Already-``RUNNING`` duplex requests (``PREFILLING`` while
     the seed is still being consumed, then ``DECODING``) — one
     row each, using the exact same page-commit accounting
     (``_decode_page_commit`` / ``_can_allocate_reserved`` /
     ``_preempt_until_allocatable``) the ordinary
     decode-continuation and prefill-continuation loops use.
  2. Still-``WAITING`` duplex requests — this is the ONLY
     admission path they ever take (they are invisible to
     ``_next_waiting_candidate``). A request with seed prompt
     still to consume is admitted into ``PREFILLING``, mirroring
     ``Scheduler.schedule``'s ordinary prefill admission; one
     whose prompt is already consumed goes straight to
     ``DECODING``.

MTP bucketing and batch/interactive priority arbitration do not
apply — duplex requests are outside that priority split
entirely. ``chunk_prefill`` and ``max_batched_tokens`` DO apply,
to the seed prefill rows, so a long system prompt is chunked
rather than served as one unbounded row — but a chunk is NOT
"just another row shape the runtime already handles" once other
sessions are live, which is what the wide-row invariant below
exists for. Bounded by ``batch_cfg.max_batch`` like any other
step (defensive; duplex session counts are expected to be small).

Only sessions whose OWN deadline is due get a row this tick, and
each slated session's own ``next_duplex_deadline`` is advanced by
its own ``duplex_interval_s`` — never another session's, and never
a session's that did not actually get a row (design doc §7.9.14).
A step that merely checked and found nothing due advances nothing.

**A row WIDER than one token is served ALONE** (design doc
§7.9.18) — see :meth:`_duplex_wide_row_owner` for why that is a
correctness invariant of this lane and not a throughput choice.
The other due sessions keep their deadlines (nothing advances a
deadline that did not get a row), so they are still due and are
served on the very NEXT engine step rather than losing a frame;
the lane's own post-step hook returns their unconsumed audio to
the inbound queue through the ordinary prepared-but-not-stepped
path (``DuplexLane._requeue_prepared_frame``).

One pass of :meth:`_build_duplex_slate`'s two-source walk.

``solo`` restricts the pass to a single request; ``allow_wide=False``
skips every row wider than one token. Both exist only for the
wide-row invariant — see :meth:`_duplex_wide_row_owner`.

The speculation load valve: mode resolution + the runtime-adaptive controller.

The valve decides, per decode step, whether the MTP-opted rows draft or
run a plain decode. It is a tri-state knob:

``off``
    Never disable. Every opted row always speculates.
``on``
    Disable at a pinned decode batch the operator chose. A bare ``on``
    with no number takes the architecture fallback
    (:func:`arch_fallback_batch`).
``auto``
    Decided from live measurement by :class:`SpecValveController`.

``auto`` compares the two arms on the quantity the trade is actually
about — committed tokens per second at the current decode width — by
running each arm in phases and keeping a per-width estimate of both. It
needs no accept rate, no model-topology guess and no cached number from a
synthetic sweep: the drafter's payoff is whatever it measures on the
traffic the server is serving.

The counterfactual arm is the signal that disappears once a decision is
made, so the controller re-probes it: on a cadence that stretches as the
verdict becomes more decisive, and sooner when the arm it CAN still see
moves outside its own band — what a workload change looks like from the
surviving arm. A probe runs the arm the valve did not choose, so it has a
price; the cadence floor is what bounds it, at one phase in
``PROBE_PHASES_MIN + 1`` whatever the traffic does.

Stopping is not the mirror of starting. A width that stops speculating
also stops observing the arm it stopped, so a stop taken on one thin
measurement outlives that measurement by however long the probe cadence
takes to contradict it — while the cadence STRETCHES with how decisive the
wrong verdict looks. So a stop needs both arms sampled enough for a median
to mean something and needs consecutive probes to agree, and a resume
needs one probe to disagree.

Why ``auto`` stands down on a dense verifier — the measurement
--------------------------------------------------------------

The stand-down itself is implemented at the gate (see
:meth:`SpecValveController.bootstrap_never_disables` and its caller in
``mtp_bucketing``). What follows is the evidence, kept here so nobody
rebuilds an A/B estimator over a stateful arm.

The estimator is BIASED, not merely noisy, and the bias is one-sided.
Disabling speculation for a step drops that row's cached drafts and
staleness-marks a tap-conditioned drafter's draft-KV slot, so a
speculating phase that opens after a switch is scored while the drafter is
cold. :data:`PHASE_SETTLE_STEPS` discards two steps; the re-warm is longer
than that. The plain arm carries no state a switch can destroy, so the
error only ever runs against the arm the valve is deciding whether to
keep. It is also SELF-CONFIRMING: once a width has stopped, every
speculating phase left is a probe — a short run opening right after a
switch — so every one of them is scored cold and the depressed estimate
re-earns itself. ``test_a_stopped_width_cannot_re_measure_a_drafter_that_
reopens_cold`` pins that.

The counterfactual is also STARVED, which no care taken on the observed
arm repairs. Measured on a live dense 27B + DFlash2 boot driving 4-wide
bursts of 90-token turns for eleven minutes:

  * 57 scored phases at width 4 — 56 speculating, and exactly ONE plain,
    the bootstrap probe. The probe cadence stretches with ``|gain - 1|``,
    and at the measured gain of 1.69 it sat at its
    :data:`PROBE_PHASES_MAX` ceiling, so the arm was never re-measured.
    The whole stop-or-go verdict rested on that single plain sample for
    the entire run.
  * The speculating arm was So a wrong verdict
    here was never a noisy-estimator problem.
  * One draw on either side of that is the entire
    verdict.

A dense verify step runs ``B*(K+1)`` tokens through the same weight
traffic a ``B``-token decode does, so the prior that it cannot lose is
already right. An estimator that is biased against speculation, starved of
its counterfactual, and confirms its own mistakes has nothing to add to
that, which is why on a dense verifier it is not consulted.

``auto`` is a TRACKER, not an optimum. Its ceiling is the better fixed arm
minus what the probes cost, so an operator who has measured their own
deployment beats it by pinning ``off`` or a batch. What it buys is never
being catastrophically wrong on a deployment nobody measured — which is
the failure a single shipped constant produces.

Last-resort disable batch guessed from the verifier's FFN topology.

A GUESS, never a measurement, and labelled as one wherever it is
surfaced. It exists to give ``auto`` a starting arm before it has
measured anything, and to give a bare ``on`` a number. The quantity it
stands in for is the drafter's accepted tokens per step against the
verify step's ``K+1`` token multiplier, which
:class:`SpecValveController` measures directly and supersedes as soon
as it holds both arms at the serving width.

A verify step runs ``B*(K+1)`` tokens where a plain decode runs ``B``.
On a sparse MoE each token routes to its own expert subset, so the
verify pass multiplies the traffic that already dominates; on a dense
model it spreads the same weight traffic over more rows of GEMM.

Returns ``(batch, reason)``; ``0`` means "never disable".

Resolve ``mtp_spec_disable_batch`` into a :class:`ValveSetting`.

Accepted: ``""`` (unset ⇒ ``auto``), ``auto``, ``off`` / ``never`` /
``0``, ``on`` (⇒ the architecture fallback batch), and any positive
integer (⇒ pinned). Anything else resolves to ``off`` — never
disabling is the safe reading of an unreadable flag — and says so.

Round a decode width up to a power of two.

Both terms of the trade move with the width, so estimates are kept per
width; bucketing stops a busy server spreading its samples over every
integer batch it ever ran.

Committed tokens/sec at one decode width, one estimate per arm.

Each arm's estimate is the MEDIAN of its recent phase scores. A live
server's phases are not equally clean — prefill interleaves, requests
arrive and finish, the slate churns — so an estimator a single sample
can move is one the valve will chase, spending its phases re-checking
noise instead of serving on the arm it already knows is faster.

Runtime-adaptive speculation valve.

Runs the two arms in phases and keeps, per decode width, an estimate
of committed tokens/sec for each. Between phases it picks the arm its
estimates favour, holding inside a dead band so noise cannot flip it,
and periodically spends a phase on the arm it did NOT pick so the
comparison stays live.

The decision is necessarily per STEP, not per request: one step runs
one uniform verify width, so the whole slate speculates or none of it
does. Under a mixed workload the estimates converge to the mix average
at each width and the valve picks the better arm FOR THE MIX; the dead
band and the phase cadence are what stop a mix straddling parity from
oscillating.

Whether the CURRENT pair reads as "speculation loses here".

A candidate verdict, not the verdict: it takes
:data:`VERDICT_CONFIRMATIONS` consecutive probes agreeing on this
before :attr:`stopped` follows. Both gates below are what keep a
single unlucky phase from being one — enough samples that the
median is not one draw, and a gap below parity wider than the dead
band the arm choice already respects.

Phase pairs that contradict a stop: speculating beat plain.

The two rings compared as samples rather than as two medians. A
width where speculation genuinely loses separates the rings; a
width sitting near parity interleaves them however the medians
happen to land.

Fold one counterfactual re-read into the stop verdict.

Called at the close of a PROBE and nowhere else: only a probe has
just re-measured the arm the width is not running, so only a probe
is fresh evidence. The streak is a bounded walk rather than a run
counter: an agreeing probe advances it, a disagreeing one walks it
back by :data:`VERDICT_RESUME_STEP`. That is what makes stopping
cost several probes and resuming cost one.

Pin the pair the current verdict rests on.

Called when both arms have just been measured against each other,
so a later drift is measured from a comparison that was true.

Whether the arm in use has moved away from the pinned pair.

This is the only workload-change signal available once one arm is
switched off, and it is what buys the counterfactual back — so it
has to fire on a real change and stay quiet on the arm's own
scatter. Two guards do that: the estimate must be old enough to
have a shape at all, and the move must clear BOTH the fixed floor
and the spread this arm's own recent phases already show. A
speculating arm's throughput moves with the accept rate from one
phase to the next; a valve that read that as a workload change
would spend its phases re-checking noise.

True when the architecture fallback says never disable.

``0`` from :func:`arch_fallback_batch` is the DENSE verdict, and it is
a statement about the verifier's cost structure rather than a guess to
be refined by measurement — see the caller in ``mtp_bucketing``.

Smallest width holding a CONFIRMED "stop speculating" verdict.

``0`` when no width holds one. Reported as the valve's effective
threshold so an ``auto`` deployment surfaces the same quantity a
pinned one does — and it is the same fact the arm choice rests on,
not a second, looser reading of the estimates.

Read off a value maintained at phase boundaries. This is called
once per admission step, so it may not scan the widths or take a
median.

Close out the step that just ran, then answer for the next one.

``opted`` is the MTP-opted decode slate at admission; its rows
carry the token counts the PREVIOUS step committed, which is what
makes the measurement free of engine plumbing. Costs one clock
read and one pass over the slate.

``pure_decode`` is whether the step now opening carries ONLY decode
rows. A step that also prefills spends most of its wall clock on
the prefill, and charging that to whichever arm happens to be
running measures the arrival pattern instead of the drafter. Such
steps are skipped rather than scored; the phase and the arm hold
across them.

Which arm this phase runs, and whether that is a probe.

An arm with no samples must be sampled or the comparison never
comes into existence — so bootstrapping a width is itself a probe,
and it is the ONLY way ``auto`` ever escapes its starting guess.
With both arms present the probe falls back to the cadence.

Whether this phase spends itself on the arm not in use.

Never before both arms exist (there is nothing to hold against),
then on a cadence that stretches with how decisive the verdict is.
A pending drift shortens that cadence to
:data:`PROBE_PHASES_MIN` but never below it: this floor is what
BOUNDS what keeping the counterfactual alive can cost. At most one
phase in ``PROBE_PHASES_MIN + 1`` runs the arm the valve did not
choose, whatever the traffic does.

Three things take the cadence to that floor, all of them cases where
the stored counterfactual is not something to ride: a pending drift,
a width whose arms are not yet both sampled enough to return a
verdict at all, and a width whose live read disagrees with its
standing verdict. The last is what makes confirmation cheap in
TIME — a candidate stop is re-drawn at the fastest rate the cost
bound allows rather than at a cadence that stretches with how
decisive the candidate looks.

The arm this width's standing verdict implies.

Speculating is the resting state, and a width leaves it only on a
CONFIRMED stop — the same fact :meth:`disable_batch` reports, so
the arm the server runs and the threshold it publishes can never
disagree. The dead band and the sample floor live inside
:meth:`WidthEstimate.stop_reads`, and the confirmation across
probes lives in :meth:`WidthEstimate.confirm`.

Per-step slate-budget computation for the continuous-batch scheduler.

This is pure host-side arithmetic computed ONCE at the top of each
step, BEFORE the slate-build loops run — it does NOT touch the
per-token ``commit`` accounting where the documented start_pos /
seq_lens off-by-one class lives. ``schedule`` stays the orchestrator:
it calls :func:`compute_step_budget`, unpacks the result, and runs its
continue-running + admit-waiting loops.

The computed knobs:

  * ``chunk_cap_this_step`` — fair-share chunk_prefill cap: the prefill
    budget is water-filled over the K competing prefills' remaining demand,
    so a row below the equal share takes what it needs and the level for the
    rest rises, instead of the first prefill grabbing the whole budget or a
    short prompt halving a long one's chunk.
  * ``interactive_present`` — True iff ANY interactive request is in the
    system this step (waiting OR running). Drives the interactive-only
    slate construction (batch fully suspended when interactive present).
  * ``batch_rows_left`` / ``batch_tokens_left`` / ``batch_prefill_cap`` —
    batch-only step-size caps bounding the one in-flight batch step an
    arriving interactive request must wait out.

True iff any non-batch request is queued or running.

Delegates to the scheduler's O(1) counter predicate when it has one;
schedulers that fabricate ``waiting`` / ``running`` without the
counter machinery (unit-test doubles) fall back to the exact scan.

Pydantic schemas for the arbi-serve HTTP API.

Layout:
  - :mod:`._shared` — base classes / mixins (:class:`OpenAIRequest`,
    :class:`AdminRequest`, :class:`StrictResponse`, :class:`Usage`,
    :class:`LogProb`, :class:`LogProbContent`, :data:`FinishReason`).
  - :mod:`.openai`  — OpenAI-compatible request / response models for
    ``/v1/chat/completions``, ``/v1/completions``, ``/v1/models``.
  - :mod:`.extensions` — vendor extension fields that ride
    ``extra_body`` (engine params, response_format extensions,
    attention backend override, lora selection).
  - :mod:`.admin`   — admin endpoint request / response models.

The package is import-cheap: no torch, no fastapi, no engine code.
Safe to pull from any tooling (openapi-python-client, codegen,
test fixtures).

Base class for every server-emitted response.

``extra='forbid'`` so accidental field leaks fail validation in CI
rather than slipping into the wire format. We control what we
emit; downstream clients should be able to round-trip strictly.

Base class for every OpenAI-shaped REQUEST model.

``extra='allow'`` so vendor-extension fields ride through
``extra_body`` (the OpenAI Python client surface for "any extra
JSON to send"). The handler is responsible for inspecting
``model_extra`` (or named optional fields) for the bits it
understands.

Tolerating an unknown key is not the same as tolerating an unserved
one: a key in :data:`UNSERVED_REQUEST_FIELDS` names real behaviour
this server does not implement, and is refused so a client cannot set
a knob that does nothing. Everything else still rides through — that
is what keeps a newer OpenAI platform field from breaking a request
it has no effect on.

Base class for every admin-endpoint request model.

Admin is internal-only so we run ``extra='forbid'`` — typos in
operator scripts should surface as a 422 immediately, not silently
succeed with the typo dropped.

OpenAI ``usage.completion_tokens_details`` breakdown.

A DETAIL of ``completion_tokens``, never a subtraction from it: on a
reasoning model ``completion_tokens`` already INCLUDES the reasoning
trace (matches OpenAI), and ``reasoning_tokens`` reports how much of that
total was spent thinking — so ``completion_tokens − reasoning_tokens`` is
the visible-answer length. Shape mirrors OpenAI's four optional fields;
arbi-serve populates ``reasoning_tokens`` and — when MTP speculative decode
was active — the ``accepted``/``rejected_prediction_tokens`` pair from the
speculation accounting. Each field is omitted (not 0) when it has no
producer for the request, so counts are never misleading. ``audio_tokens``
has no producer here.

OpenAI ``usage.prompt_tokens_details`` breakdown.

A DETAIL of ``prompt_tokens``, never a subtraction from it:
``prompt_tokens`` is the full rendered prompt, and ``cached_tokens`` says
how many of those positions were served from a prefix-cache hit instead
of being prefilled.

Unlike :class:`CompletionTokensDetails`, whose fields are omitted when
they have no producer, ``cached_tokens`` is ALWAYS stamped on a server
that has a prefix cache — ``0`` is a real answer meaning "this prompt was
fully prefilled". The distinction matters: an ABSENT field and a zero are
the same value to a client, so omitting it makes a cache MISS
indistinguishable from a server that does not report caching at all. That
ambiguity is not hypothetical — it is what made an HTTP-level
prefix-cache probe unreadable and cost a 27B boot.

OpenAI-shaped token-usage block on chat / completion responses.

``timing`` is the arbi-serve extension carrying detailed
per-request latency breakdown (TTFT / TPOT, cold-tax, warmup
components). It is omitted unless per-step timing is opted into
via the ``return_timing`` body flag, the ``X-Arbi-Return-Timing``
header, or the server-wide ``cfg.timing_debug`` default.

Admin endpoint request AND response schemas.

Every admin request model uses ``ConfigDict(extra='forbid')`` (inherited
from :class:`AdminRequest`) — admin is internal tooling and a typo in an
operator script should fail loudly at the 422 layer rather than silently
dropping the field.

Response models here are declared as each route's ``response_model=``.
They exist to make the OpenAPI spec / generated ``arbi_serve_client``
actually typed for the admin surface. A prior cleanup (86a31f55)
deleted an earlier attempt at this as dead code — those models were
declared but never wired to ``response_model=`` on any route, so they
validated nothing and appeared nowhere in ``openapi.json``; removing
truly-unused code was the right call, but it also means every model
below was written fresh and IS actually attached to a route this time,
which is what makes drift loud (a ``ResponseValidationError``) instead
of silent. ``git grep response_model`` shows only ``tokenize.py`` and
``openai_realtime.py`` used this mechanism before this file did — the
OpenAI-compat completions/chat routes are UNTYPED too (streaming
responses don't fit ``response_model`` the same way in any case; that's
a separate, not-yet-tackled piece of work), so admin was never really
an outlier here, just the first surface actually swept.

Every field's shape was verified against the actual engine-side
``return {...}`` statements it models (see each model's docstring for
its source), not guessed from the route's prose docs. Where the real
runtime response is polymorphic depending on which code path fired
(e.g. a no-op swap vs a real one), the differing fields are typed
``Optional`` with the discriminator field documented — Pydantic fills
in the declared defaults for whichever fields a given call's response
omits, so one model covers every variant without a fragile Union.

Request body for ``POST /v1/admin/model`` — switch-to-or-load a model.

Default ``mode="swap"`` routes through the fast stable-VA park/wake: a model
already resident becomes active by a ~ms park/wake, an absent model is built
into the residency (add-then-wake) and made active — both keep the residency
so a subsequent swap-back is fast. ``mode="reload"`` is the explicit
destructive fallback (drain → free → rebuild via ``areload_model``) for a
topology change a park/wake cannot express.

Response for ``POST /v1/admin/cache/flush``.

Source: ``cache_flush`` -> ``RadixPageTable.flush`` (always this
exact shape; ``AdminMethodUnsupportedError`` -> 501 with no body
is the only other outcome, on a page-table backend with no
cross-request cache).

Response for ``GET /v1/admin/oscar``.

Reports what the RUNTIME resolves, not what was requested: the rotation is
installed by assigning an env var, so a status that echoed the request back
would report "enabled" for a flip that silently failed to take.

Response for ``POST /v1/admin/model`` with ``mode="swap"`` (default).

Source: ``aswitch_or_load_model`` -> ``runtime_resident.aswitch_or_load_model``.
``route`` is ``"switch"`` (already resident, fast park/wake) or ``"loaded"``
(built into the residency and made active). ``active`` is the now-active
resident's routing key.

Response for ``GET /v1/admin/calibration/bundles``.

Source: the SAME directories auto-discovery scans
(``cli.kv_resolve.calibration_search_dirs``), in the same precedence
order. Nothing outside them is read, and the route takes no path
input at all -- there is no directory a caller can point it at.

Response for ``POST /v1/admin/attention_backend``.

Source: ``aswap_attention_backend`` -> ``swap_admin.py``'s
``_swap_status_payload`` (+ extra keys on the real-swap path). The
no-op path (already-active backend) omits ``prev``/``new``/
``kind``/``route``/``kv_capacity`` entirely -- all five are
Optional for exactly that reason, discriminated by ``no_op``.

One swimlane row -- either a finished request (``live`` absent)
or an in-flight one (``live: true``, ``end_ms`` pinned to "now" so
the bar visibly grows across polls). Source:
``request_timeline.py``'s ``RequestTimelineRing.rows`` /
``in_flight_rows``.

Response for ``POST /v1/admin/request/{request_id}/answer_now``.

``closed`` is whether the request was ARMED to stop reasoning, not whether
it has finished doing so: the close lands on the next forward pass, when
the guard feeds the closing marker. ``reason`` is empty on success and
carries why nothing happened otherwise — the operation refuses rather than
no-ops, because a control that silently does nothing is worse than one
that is absent.

Response for ``POST /v1/admin/request/{request_id}/stop_now``.

``stopped`` is whether the request was ARMED to stop, not whether it has
finished: the forced token lands on the next sampling step. ``reason`` is
empty on success and carries why nothing happened otherwise.

Response for ``POST /v1/admin/shutdown``.

Source: ``request_process_exit`` (``engine/shutdown_request.py``). The exit
is ARMED when this returns, not yet taken, so the reply reaches the caller.

``supervised`` is what the DEPLOYMENT declared
(``ARBI_SERVE_SUPERVISED``), because nothing inside the container can see
its own restart policy. ``expect_restart`` follows from it: without a
declared supervisor this call stops the server and nothing brings it back.

Response for ``POST /v1/admin/release_memory``.

Source: ``release_memory_occupation`` (``engine/sleep.py``).
Idempotent no-op when already released: ``{"already_released":
true}`` only -- every other field is Optional for that reason.
409 (not a variant shape) on a non-phase2 pool.

Response for ``POST /v1/admin/resume_memory``.

Source: ``resume_memory_occupation`` -> ``resume_engine_state_sync``
(``engine/sleep.py``). Idempotent no-op when nothing was released:
``{"already_resumed": true}`` only. The real-resume path spreads
``eng.sleep_pool.resume()``'s own return dict at the top level
alongside the fixed fields below -- those extra keys are pool-
implementation-specific, hence ``extra="allow"`` rather than a
guessed, possibly-wrong sub-schema.

One ``POOL_CATEGORIES`` entry: the group headline, its full
lifetime/scaling line (``pool_taxonomy.CATEGORY_LIFETIME``), and the
``tenure`` every row in it carries (``pool_taxonomy.CATEGORY_TENURE``) --
what KIND of number those rows are, which is the axis ``category`` alone
does not carry.

One ``POOL_TAXONOMY`` entry as a UI renders it: category, ``tenure``,
the one-line "what it holds, what sizes it" caption, the row's DEFAULT
provenance, and whether a non-zero value on it is a leak signal rather than
accounting.

``tenure`` follows the category and is never overridden per row: it says
whether the row's memory EXISTS, and a row allowed to disagree with its
category about that would put the two views back on separate schemes.

One :data:`arbi_serve.runtime.pool_taxonomy.PROVENANCE` value: what it
means, and whether a UI must FLAG it rather than render it like the rest
(``SEEDED`` is a guess; ``UNBUDGETED`` is a hole in the plan).

One row of the ``card`` block -- a named slice of the device total.

``bytes`` is what the row holds (signed: the one negative row is
``transient.serving_step.in_use``, reserve a live step has already taken).
``live_bytes`` is its allocated subset. ``provenance`` is how the number was
obtained (:data:`arbi_serve.runtime.pool_taxonomy.PROVENANCE`) --
``SEEDED``/``UNBUDGETED`` must be flagged by a renderer, not shown like a
measurement. ``holds_physical`` is False only for a ghost row.

``card`` block in ``memory_live`` -- the WHOLE device, as named rows.

The closing contract, and the reason a client needs no arithmetic of its
own: ``rows_bytes + free_rows_bytes == total_bytes`` exactly whenever
``closes`` is true. ``rows`` is resident physical (``driver.residual``
included, computed server-side from the card down); ``free_rows`` is the
unallocated remainder named by what will allocate it; ``ghost_rows`` hold no
physical and take part in no sum -- and ``ghost_unverified_bytes`` is the
part of them the allocator does NOT measure as unmapped-and-released, i.e.
resident memory a ghost row would be hiding.

When ``closes`` is false the identity has a hole: ``unresolved`` names every
contribution the snapshot could not establish and ``unaccounted_bytes`` is
the gap, which is NOT booked on any row. A renderer must show that state,
never close over it.

One Level-2 line: an input a group's number is a function of.

``value`` is pre-formatted server-side like every other figure in this
block -- a console that divides one field by another is a console that can
divide by the wrong thing, and the unit is part of the claim. ``detail``
says what the input MEANS, which is the half that makes the number
actionable.

One piece of a cross-cutting view: a row it claims, and where that row
already lives. ``measured_by`` says what took the number, because a part
that is not a row of the card is only as good as its measurement.

A cost spread over more than one group -- a VIEW, never a group.

The groups partition the card, and that property is why any figure here can
be trusted. A cost that straddles the partition (CUDA-graph capture pays
the allocator, the driver's exec heap and the driver's cubin residency)
cannot be folded into a group without booking the driver's share twice. So
it is carried beside them, takes part in NO total, and states in
``overlaps`` / ``overlap_bytes`` exactly which group's bytes it re-reads and
how many. A view that could not produce that number would be a second
partition wearing the first one's clothes.

One headline sentence: the card in exists / idle / promised / unclaimed.

Written server-side because it is a verdict, not a caption -- the numbers in
it are ones this server already computed, and a console that re-derived
them could disagree with the bars beneath.

One weight class inside a group -- what the checkpoint stores for it.

Measured from the safetensors index's own ``data_offsets`` spans, so these
are STORED bytes. What the pool row holds is the same tensors after this
rank's shard and the loader's repack; the console says so rather than
implying the classes add up to the bar.

One named ROW inside an operator group -- the group, opened.

The bar answers "how much"; a bar folding a residency gate, an operator
ceiling and a rounding cushion into one figure cannot answer "what for".
``label`` is the pool's own name with its category prefix dropped, so this
adds no vocabulary: the word under the bar is the word on the row in
Diagnostics.

``value_text`` is THE string a console prints -- formatted server-side so a
reserve and an allocation cannot render as the same kind of number, and so
a measured peak prints BESIDE the reserve rather than instead of it.
``peak_bytes`` is ``None``, never 0, where nothing measured a peak: a 0
beside a reserve claims that nothing ever needed those bytes, which is a
claim no measurement made.

One bar of ``memory_live``'s ``operator`` view: the card's rows folded
into a category a person running the server can act on.

``label``/``blurb`` are operator language for names the card keeps in
engine language, ``growth`` is the one sentence saying what a reader does
when this group grows, and ``rows`` names exactly which card rows were
folded in, so the two views are provably the same bytes.

``tenure`` says what KIND of number the group is (``pool_taxonomy.TENURE``,
plus ``UNKNOWN`` for the unrecognised-pools group); ``resident`` is that
same fact as a boolean, derived from it rather than declared beside it.
``idle_bytes`` is reserved-minus-live INSIDE this group -- the same physical
``bytes`` already counts, carried as a column because a bucket for it would
book those bytes twice. It is 0 on a group that is not resident: a
reservation has no live subset to be idle against.

One section of ``memory_live``'s ``operator`` view: every group whose
bytes are the same KIND of number, and that section's total.

The split the flat group list did not have. ``heading`` states whether the
memory EXISTS; ``meaning`` says what a reader can do about it. Totalled
server-side like everything else here -- a console that adds up the bars it
was shown books whatever it was not shown into the heading above them.
``groups`` names the group keys in the section, in display order.

One entry of ``memory_live``'s ``operator.warnings`` -- a row set that
means something is WRONG, in a sentence rather than a badge.

Every entry restates a verdict the server already reaches (the card not
closing, an unbacked ghost row, a leak-signal row with bytes on it, a
residual past the alarm threshold, a row whose number is a guess). A
healthy boot produces an empty list, which is the property that makes a
non-empty one worth reading.

``operator.capacity`` -- the card's two totals, named apart.

``addressable_bytes`` is what the process can allocate, what the card closes
against, and the denominator of every ``pct`` on this view.
``physical_bytes`` is the whole card and is LARGER; the difference is the
driver's own carve-out, which is on no group because no pool can claim it.
Every figure is pre-formatted here for the same reason the rest of this
block is: a console that picks its own denominator is the defect.

``physical_known`` false means NVML could not be asked -- ``physical_bytes``
and ``not_addressable_bytes`` are then ``None`` and must render as
unavailable, never as zero and never as the addressable total.

``operator`` block in ``memory_live`` -- the card, grouped for a reader.

The same bytes as ``card``, folded into a handful of named groups with the
wrong-looking rows lifted into ``warnings``. Grouped server-side for the
same reason the card is closed server-side: a consumer that sums rows
itself books whatever the payload omitted into a bucket of its own naming.
``summary`` is the one-line statement of whether the card closes, so a
reader is given the verdict instead of arithmetic to check.

One entry in ``memory_live``'s ``unpooled_owners`` -- one allocation
shape holding part of the ``unpooled.torch_default_pool`` row, named by the
boot-time heap walk. Empty when the boot did not reach the freeze.

``bytes``/``count`` are per STORAGE and ``views`` counts the live tensors
over them, so a folded width-max buffer handed out as per-rung views reads
as the bytes it actually holds. ``shape``/``dtype`` say what the allocation
IS, which is what identifies it when ``owner`` names a holder the reader
does not recognise.

``kv_mapped`` block in ``memory_live`` -- the ``state.attn_kv.mapped``
row plus the STATE of the reading, from
:func:`arbi_serve.engine.memory_budget.driver_residency.resolve_mapped_kv_row`.

The growable KV slab drives cuMem directly, so it is outside every
allocator snapshot and is the largest row of a serving card. ``state`` is
``mapped`` (``bytes`` is its mapped physical), ``pooled`` (the slab is not
growable, so its bytes are on the ``state.attn_kv`` pool row), ``absent``
(this engine holds no KV pool) or ``unresolved`` (the read failed and
``bytes`` means nothing).

Response for ``GET /v1/admin/memory/live``. Source:
``memory_live_snapshot`` (admin_registry.py).

Self-sufficient, and CLOSED: ``card`` carries every row of the device
total, resident and free, with the residual already computed and named
server-side. A consumer renders it and computes nothing -- a consumer that
subtracts is computing a quantity only the server can name, and books every
row the payload forgot into it.

``pools`` / ``device`` / ``reconciliation`` / ``driver_measured`` remain the
raw inputs ``card`` is built from, for a caller that wants the measurement
rather than the accounting.

Response for ``GET /v1/admin/config``. Source: ``dynamic_config.py``
-- every field defaults to unset (``None``) until the first
``POST /v1/admin/config`` sets it. ``_keys`` (the leading underscore
is the real wire name) can't be a Pydantic field name directly, so
it's aliased.

One entry in ``config_override``'s ``params`` catalogue --
heterogeneous per registry-entry kind. Source: ``known_params()``
(config_overrides.py). A normal overridable param carries
``target``/``runtime_flag``/``capture_affecting``/``doc``; an alias
entry carries only ``alias_of``; a model-select entry carries
``target`` + ``model_select: true``. All fields Optional for that
reason -- the admin console's own UI already filters to
``doc is not None`` rows (real params only) before rendering.

One param's LIVE value domain. Source: ``engine/param_domains.py``.

Distinct from ``ParamSpecEntry.choices``, which is the static
vocabulary a param has on every boot. A domain is a property of THIS
boot -- the drafter checkpoint loaded, the attention codec selected --
so it cannot live in the registry. A param with no resolvable bound
carries no entry at all, and a client without one keeps free entry.

Response for ``POST /v1/admin/config_override``.

Source: ``aapply_config_overrides`` -> ``config_variant.py`` --
``route`` discriminates which of 4 code paths fired, each spreading
the same :class:`CurrentConfig` base plus its own extra fields:
``noop`` (no extra fields at all for an empty-overrides call, or
``key``+``switch_s`` for an already-active known variant),
``instant_switch`` (``key``+``switch_s``+``swap``),
``live_overlay`` (``key`` only), ``prepared_switch``
(``key``+``prepare_s``). All four extra fields are Optional for
that reason.

Response for ``DELETE /v1/admin/pool``. Source:
``config_variant.py``'s ``adrop_variant()`` / ``adrop_all_variants()``.

``dropped`` lists the variant key(s) whose member was torn down;
``freed_gib`` is the physical VRAM reclaimed; ``active`` is the baseline
member left resident; ``kv_pages`` is the active member's KV capacity after
the reclaimed VRAM was grown back into it (``None`` when nothing was freed or
the pool is not growable). The ephemeral-A/B reset: after this the residency
holds exactly the active member at its full KV capacity — no permanent
KV/memory tax.

Response for ``GET /v1/admin/server_info``. Source:
``server_info.py``'s ``get_admin_server_info()`` — a best-effort
header-strip composite, so every field is independently Optional
(a failure resolving one never fails the whole call).

``memory_released`` mirrors ``Engine._memory_released`` (the
stable-VA sleep flag) — ``None`` on a CPU/no-GPU stub or process
mode where that attribute isn't readable, not just "awake".

One generated word in ``WatermarkDetectResponse.per_token``.

``u`` is a match score from 0 to 1: higher means the key more strongly
favoured this word. ``masked`` words sat on a repeated context and were
skipped (no score). A UI can shade each word by ``u`` to show where the
watermark is coming through.

Request for ``POST /v1/admin/watermark/detect`` (``watermark.py``).

Exactly one of ``text`` / ``token_ids`` must be set. ``key`` overrides
the server's configured watermark key (hex or decimal string) — needed
to score text against a rotated-out historical key; it is never logged
and never echoed back. ``prompt_len`` marks leading tokens as
context-only (scored positions start after them).

Response for ``POST /v1/admin/watermark/detect``. Source:
``watermark.py``'s ``post_admin_watermark_detect()`` wrapping
:func:`arbi_serve.sampler.watermark.detect_tokens`.

``p_value`` is the exact Gamma-tail survival probability and
underflows to ``0.0`` for strong watermarks; ``log10_p_value``
switches to the asymptotic tail there and never underflows.
``key_source`` says which key judged the text (``server`` config or
``request`` override) without revealing either.

Embedding + reranking request/response schemas.

Two wire surfaces:

  * ``POST /v1/embeddings`` — OpenAI's embeddings shape, so any OpenAI
    client (``client.embeddings.create``) points at us unchanged. Served
    by Qwen3-Embedding instances.
  * ``POST /v1/rerank`` — the Jina / Cohere rerank shape (also what
    TEI / Infinity expose), since OpenAI has no rerank surface. Served by
    Qwen3-Reranker instances.

Both endpoints only exist when the instance was booted for the matching
task (``--task embed`` / ``--task rerank`` or auto-detected from the
model config); a generation instance returns 404.

Request body for ``POST /score`` (vLLM score-API compatible).

Unlike ``/v1/rerank`` (which wraps query + documents in the model's
judgement template server-side), ``/score`` treats the texts as
ALREADY templated: each ``text_1 × text_2`` pair is concatenated
verbatim and scored as-is. Clients apply the chat template
themselves before sending ``text_1`` / ``text_2``.

Vendor extension fields that ride OpenAI's ``extra_body``.

The OpenAI Python client lets users pass arbitrary additional JSON via
``client.chat.completions.create(..., extra_body={...})``. We declare
these as named fields on the request models (with
``ConfigDict(extra='allow')`` upstream so that unknown keys are still
tolerated). Each is optional and falls back to safe defaults.

OpenAI-compatible request and response schemas.

The wire shapes match the OpenAI ``/v1`` surface so any OpenAI
Python / JS / Go client points at us with no code changes. Vendor
extensions (tools, response_format extensions, attention_backend,
lora, return_timing) ride on optional fields and are tolerated by
real OpenAI clients via ``extra_body``.

Conventions across this module:

  - Every field carries a ``Field(description=...)`` that surfaces in
    the generated OpenAPI spec and the openapi-python-client docstring.
  - Request models use ``OpenAIRequest`` (extra='allow') so unknown
    keys ride through; response models use ``StrictResponse`` (extra
    forbidden) so we never leak.
  - ``Literal[...]`` is used for finite enums (role, type, finish
    reason).
  - One realistic ``json_schema_extra.examples[0]`` per top-level
    request / response.

Function definition inside a request's ``tools`` entry.

Pydantic counterpart to the engine-side dataclass
:class:`arbi_serve.engine.request.ToolFunction`. Named ``FunctionDef``
to keep symmetry with OpenAI's wire schema (``tools[*].function``)
and to avoid clashing with the engine dataclass when both are
imported in the same module.

OpenAI-compatible ``stream_options``.

``include_usage`` is the OpenAI knob: when True, the streaming
response emits a final SSE data chunk with empty ``choices`` and a
populated ``usage`` block before ``[DONE]``.

``include_token_count`` is an arbi-serve extension. A streaming chunk
carries the text staged since the last engine wakeup: speculative
decoding puts every token a step accepted into one chunk, and an
off-loop detokenizer coalesces several steps into one. A client that
counts chunks is therefore counting the grouping, and any per-token
latency it derives moves with acceptance and with emission timing
rather than with decode speed. Off by default: it costs a field on
every content frame and only a measurement client needs it.

The ``enable_thinking`` / ``chat_template_kwargs`` surface + its merge.

Shared by ``ChatCompletionRequest`` and ``TokenizeChatRequest`` so the
``/tokenize`` fingerprint resolves thinking through the SAME code the
chat route serves with — a parallel re-implementation could drift, and
drift is exactly what this surface exists to detect.

Request body for ``POST /v1/chat/completions``.

Vendor extensions (``attention_backend``, ``lora``,
``response_format.type='regex'|'grammar'``, ``return_timing``,
``enable_thinking``) are accepted as named fields. Real OpenAI
clients pass them via ``extra_body``; the model_config
``extra='allow'`` (inherited from :class:`OpenAIRequest`) makes
that path work without us listing every possible vendor field.

Merge ``enable_thinking`` with ``chat_template_kwargs``.

Returns ``None`` when NEITHER surface specifies thinking — the
MODEL-FAITHFUL case: the renderer leaves ``enable_thinking``
undefined and the model template's own default governs.

FAIL-LOUD contract: unknown ``chat_template_kwargs`` keys and a
conflicting top-level/kwargs pair both raise ``ValueError`` (the
chat route maps it to 400) — a template knob that does not apply
must refuse, not silently serve a different prompt.

The chat-template variables to define, BEYOND ``enable_thinking``.

Only variables the client actually asked for appear — an omitted one
is left UNDEFINED in the Jinja context so the template's own default
governs, the same model-faithful rule
:meth:`resolved_enable_thinking` follows. ``enable_thinking`` is
excluded because it rides its own resolved channel: the route needs
the bool for the thinking regime, not just for the render.

The VALUES are not vetted here — the loaded checkpoint's template is
the authority on what it accepts, so
:meth:`arbi_serve.tokenizer.Tokenizer.check_chat_template_vars`
refuses an unsupported one at render time.

Reject the logprobs combinations that cannot be served.

Rejected rather than clamped: a request that silently got fewer
alternatives than it asked for produces a plausible response
whose top-k is wrong, and a caller scoring against it has no
signal that anything happened.

The prompt-logprobs depth this request actually asks for.

``echo`` with ``logprobs`` needs a logprob for every echoed
prompt token, which is the prompt-logprobs computation under
another name — so it resolves to the same depth rather than
being served by a second mechanism.

Engine build + long greedy decode for convert-to-exl3 post-conversion validation.

Builds the real arbi-serve engine on a converted checkpoint and decodes a
LONG greedy continuation. This exercises the exact runtime weight-load +
quant-swap path (trellis dequant for exl3, packed dequant for AWQ, ...) under
the SAME compile + cudagraph-capture conditions production uses — a few-token
"Paris" test exits before decode reads the cache back many times, so it
green-lights a checkpoint whose decode collapses after token ~2.

The decode is therefore long (>= 64 tokens) AND screened for degeneracy:

  * collapse to a single repeated token (the "!!!!" failure),
  * entropy loss after the first couple of tokens (a short loop), and
  * — when the source bf16 model is available — a hard divergence from the
    bf16 greedy continuation (teacher-forced top-1 agreement below a floor).

Any of these FAILS loudly. This is the acceptance gate for every exl3 arch.

Construct (and build) the arbi-serve engine on ``model_dir``.

Uses the bf16 KV backend (orthogonal to the weight-quant path under
test) so the check doesn't depend on a TKV calibration bundle, and the
DEFAULT compile + cudagraph-capture path (no eager override) so the
validation exercises exactly what production runs — including the
whole-forward prefill / decode capture that a checkpoint must survive.

Raise if ``tokens`` looks like a collapsed / low-entropy continuation.

Screens the tail (dropping the first two tokens, which a coherent reply
often legitimately shares across models — e.g. a leading "\n\n") for:

  * single-token domination (the "!!!!" wall — ~100% one token), and
  * too few distinct tokens (a tight repeating loop).

A healthy >= 64-token continuation easily clears both bars.

Teacher-forced top-1 agreement of ``eng`` vs ``ref_tokens``.

For each position ``t`` feed ``prompt_ids + ref_tokens[:t]`` and check the
test engine's single greedy next-token equals ``ref_tokens[t]``. Avoids the
compounding divergence of comparing two free-running greedy decodes (quant
noise makes those split early even when both stay coherent).

Build the engine on ``model_dir``, long-greedy-decode, screen degeneracy.

Decodes ``max_tokens`` (default :data:`DEFAULT_VALIDATE_TOKENS`, >= 64) so
decode reads the cache back many times, then FAILS via
:func:`_assert_not_degenerate` if the continuation collapsed to a repeated
token or lost entropy. When ``source_model_dir`` is a loadable bf16 model,
additionally fails when the teacher-forced top-1 agreement vs the bf16
greedy continuation falls below :data:`_MIN_BF16_AGREEMENT`.

Returns the decoded text of the quantized model's own greedy continuation.
Raises on error or degeneracy.

Convert any HF checkpoint to arbi-serve's canonical EXL3 format.

Wraps :mod:`exllamav3.conversion.convert_model` with arbi-serve-style
flags + sensible defaults + post-conversion validation that the result
loads through arbi-serve's existing EXL3 backend.

Why convert? arbi-serve's preferred quantized weights backend is EXL3
(trellis kernel, sm_89-tuned, well-validated). Operators bringing
arbitrary HF checkpoints (bf16, AWQ, GPTQ, ...) can convert to EXL3
once and serve at full TKV+EXL3 speed forever after, instead of
chasing the long tail of compressed-tensors variants. Block-wise FP8,
exotic AWQ layouts, GPTQ, GGUF — all collapse to "convert to EXL3."

The conversion itself is layer-by-layer with checkpoint/resume, so a

Usage::

    arbi-serve convert-to-exl3         --model /path/to/models/Qwen3-4B         --out /path/to/models/Qwen3-4B-exl3-4bit         --bits 4

Resume an interrupted conversion::

    arbi-serve convert-to-exl3 --resume         --work-dir ~/.cache/arbi-serve/exl3-conversions/Qwen3-4B/

After conversion, the output directory is a regular EXL3 checkpoint
that arbi-serve loads via the standard EXL3 backend::

    arbi-serve --model /path/to/models/Qwen3-4B-exl3-4bit         --kv-cache-dtype tkv ...

Stub exllamav3's grammar-filter transitives so ``import exllamav3``
works on the ``--no-deps`` slim install.

``exllamav3/__init__`` eagerly pulls in the generator/filter chain
(``formatron`` → ``kbnf``), which the slim image does NOT install. The
serving path sidesteps this via :mod:`arbi_serve.weight_quant.exl3.shim`;
the conversion path imports exllamav3 directly, so it must stub them too.
Reuses the shim's stub list (single source) but SKIPS ``flash_attn``
when it is REAL in the environment — the calibration forward prefers it.
When ``flash_attn`` is absent (the serving image ships without it on
sm_89 — arbi-serve has its own attention backends), it is stubbed like
the rest so the import chain completes, and
:func:`_prune_stubbed_attn_fns` then removes the flash-attn candidates
from exllamav3's attention dispatch so the calibration forwards fall
back to its real triton / torch-SDPA implementations.

Drop exllamav3 attention candidates whose backing package is a stub.

Runs AFTER exllamav3 imported successfully. When ``flash_attn`` was
stubbed by :func:`_stub_grammar_modules`, the ``fn_flash_attn_*``
candidates in ``exllamav3.modules.attention_fn.dispatch.attn_fns`` are
bound to stub classes — calling one during the calibration forward
would raise ``TypeError`` mid-conversion. Remove them so dispatch falls
through to the self-guarded triton / torch-SDPA fallbacks (xformers
candidates already self-guard via ``has_xformers``).

Warn when ``--bits`` is integral, which silently disables the allocator.

exllamav3's ``create_q_strategy`` starts every eligible Linear at
``base_bpw = int(math.floor(bpw))`` and then spends **the remainder** one
bit at a time by group priority. At an integral ``--bits`` that remainder
is exactly zero, so the allocator runs and has nothing to allocate: every
layer keeps the floor and the bake is uniform. Nothing errors and nothing
in the output says so — the resulting bundle simply carries no per-layer
differentiation, and ``region_bits`` stays ``None``.

That is how our production 27B came to be baked flat while the 4.15bpw
checkpoint next to it is allocated: the difference is entirely that one
was asked for a fractional bitrate and the other was not.

Returns the warning text (also for tests), or ``None`` when the value
already carries surplus.

Per-model work directory under ~/.cache/arbi-serve/exl3-conversions/.

Uses the model directory's basename so multiple in-flight conversions
don't collide. The work_dir holds the calibration data + per-layer
checkpoints + the final compiled output before move-to-final.

Fail fast with an actionable message if exllamav3 isn't installed.

The runtime check happens here rather than at module import time so
arbi-serve users who never convert anything don't pay the import
cost. exllamav3 install is two-step (per the [exl3] extra docs in
pyproject.toml); raise with the exact two commands they need.

Translate arbi-serve flags to exllamav3's argparse namespace.

exllamav3.conversion.convert_model uses argparse with hyphenated long
names that map to underscored attributes. Build the namespace manually
so we don't have to subprocess exllamav3 (saves ~3 s import overhead
+ lets us catch + reformat errors).

Build the real engine on the converted checkpoint and long-decode.

This is the part that catches a mismatch the metadata checks miss: a
checkpoint can dispatch to a registered arch and ship valid
safetensors shards yet still fail to load (a quant-blind ``.weight``
key check that never sees the exl3 packed tensors) or load but decode
wrong (a non-default codebook seed the loader ignored, or a decode
path that collapses after the first token). Building the engine
exercises the exact runtime weight-load + quant-swap path under the
default compile + cudagraph-capture conditions, and the LONG greedy
decode + degeneracy screen proves the dequantized weights actually
compute coherently past the first few tokens — a non-loadable or
mis-decoding conversion (including the "!!!!" decode collapse) fails
loudly here.

``source_model_dir`` is the original bf16 checkpoint; when loadable it
drives a teacher-forced top-1 agreement check against the quantized
decode (best-effort — skipped if absent / not bf16).

Needs CUDA — exl3 trellis dequant is a CUDA kernel. When CUDA is
unavailable the caller falls back to the metadata-only smoke.

Validate the converted output loads through arbi-serve's loader.

Catches the most likely silent failure — a successful convert run that
produces a checkpoint arbi-serve's EXL3 backend can't actually open or
forward (or that loads but decodes garbage). Metadata checks (arch
dispatch + safetensors present) run first; then, when CUDA is
available and the arch is wired, the real weight load + a LONG greedy
decode + degeneracy screen run so a non-loadable OR mis-decoding
conversion fails loudly instead of passing a broken checkpoint.

``source_model_dir`` is the original bf16 checkpoint, threaded through
to the bf16 teacher-forced agreement check (best-effort).

Return True if any of ``archs`` is a live (non-stub) arbi-serve arch.

Reads the registry's arch-table keys directly. Resolving the public
``ARCHITECTURES`` map would import every model class (and the CUDA /
kernel modules they pull in), which perturbs the libcudart init order
the exl3 kernel relies on; the raw table is import-free.

Carry the source checkpoint's bundled MTP block into the EXL3 output.

``exllamav3.conversion`` only emits the decoder + lm_head it knows how
to quantize and silently drops the ``mtp.*`` multi-token-prediction
block. arbi-serve builds its bundled speculative-decode head iff the
checkpoint carries ``mtp.fc.*`` (see
:meth:`Qwen3_5ForCausalLM.from_safetensors` — detection is purely
tensor-presence; the head's dims come from the main config and
``n_draft`` is fixed at 1, so NO config.json field is needed). Without
this the head is never built on an EXL3 body and spec-decode is
unavailable.

The block is a single transformer layer + an ``fc`` projection — tiny
— so we pass it through DENSE (bf16, untouched) rather than quantizing
it; its embed_tokens / lm_head are tied to the main model's (now EXL3)
weights at load time. We write it to a standalone ``mtp.safetensors``
shard: :class:`SafetensorsCollection` indexes every ``*.safetensors``
in the directory by header scan, so no ``model.safetensors.index.json``
edit is required.

Best-effort: logs + skips when the source is absent (e.g. a ``--resume``
with no ``--model``) or carries no ``mtp.*`` tensors.

Entry point. Returns 0 on success, non-zero on failure.

Three phases:
  1. Pre-flight — exllamav3 import + arg validation + free-disk check.
  2. Convert — call exllamav3.conversion.convert_model.{prepare, main}.
     exllamav3 streams the model layer-by-layer + checkpoints every
     ~120 s, so the conversion is RAM/VRAM-bounded (one decoder layer
  3. Post — validation that the output loads through arbi-serve's
     existing EXL3 backend.

OpenAI-compatible HTTP surface.

``build_app`` is intentionally exposed via lazy ``__getattr__`` rather
than a top-level ``from .app import build_app`` so that subpackages
which only need the routers (``server.api`` / ``server.health``) can
be imported in light environments — notably the ``dump_openapi`` CI
smoke, which has neither ``opentelemetry`` nor torch installed.
Importing ``app`` eagerly would pull the OTEL instrumentation
chain on every ``import arbi_serve.server.*``.

OpenAI /v1 surface aggregator.

The endpoint handlers live in the :mod:`arbi_serve.server.routes`
package — one shared :class:`~fastapi.APIRouter` decorated across
``models`` / ``embeddings`` / ``openai_completions`` / ``openai_chat`` /
``loras`` / ``admin`` (+ the torch-free ``_helpers`` / ``_request_ctx``
glue). See the package docstring for the chat / completions / models /
admin / tool-calling / constrained-decoding behaviour those modules carry.

This module is the thin public-API aggregator: it re-exports the shared
``router`` (consumed by :mod:`arbi_serve.server.app` and
:func:`arbi_serve._openapi_cli.build_schema_only_app`) plus the small set
of private symbols that :mod:`arbi_serve.server.batch_api`, the bench
scripts, and the test suite import from ``server.api``:

  * ``_run_chat`` / ``_run_completion`` / ``_run_embeddings`` — the
    request cores the Batch API worker drives directly.
  * ``completions`` — the HTTP route handler (imported by the tokenize
    microbench + ``test_completions_async_tokenize``).
  * ``_completion_stream`` — the SSE generator (``test_timeout``).
  * ``_auth_fields`` — the pre-submit auth/tenant resolution
    (``test_tenant_scoping``).

Like the route modules, this aggregator stays torch-free at import (no
top-level ``Engine`` / torch) so the OpenAPI smoke can build the schema
without booting the engine stack.

FastAPI app wiring + engine lifecycle hooks.

Composes:
  - Request-ID middleware (innermost; runs first on inbound).
  - Auth middleware (after request-id so 401/403/429 logs carry the
    id).
  - ``/health`` + ``/health/ready`` routes (auth-exempt).
  - SIGTERM / SIGINT handlers attached on lifespan startup.
  - OTEL ``MeterProvider`` + ``TracerProvider`` configured on lifespan
    startup; FastAPI auto-instrumented for inbound HTTP spans; stdlib
    logging instrumented to inherit trace context. Telemetry is
    pushed to an OTLP collector and also exposed for direct scrape at
    the in-process ``/metrics`` Prometheus route (see
    ``server/routes/models.py``).

Build the engine, attach it, and open the serving gate.

Everything between "the process started" and "this server can answer a
completion". Runs either inline in the lifespan startup (the port stays
unbound until it returns) or as a background task on the already-serving
HTTP loop — :func:`_lifespan` picks, and the sequence is identical
either way except for which thread runs the CUDA-bound build.

Three shapes of engine reach this: a single-rank in-proc
:class:`~arbi_serve.engine.engine.Engine`; a multi-rank driver this
coroutine builds itself (``driver_factory`` — rank 0 under torchrun,
where the build is COLLECTIVE and is the longest, most opaque part of
a TP boot); and an already-built driver (``prebuilt_driver``), which
has no build left to watch.

On failure the process hard-exits: a boot refusal (e.g.
``MemoryBudgetError``) must not unwind through uvicorn's lifespan into
interpreter finalization, where a captured CUDAGraph still pinning a
cuMem pool faults its destructor and masks the clean non-zero exit as a
SIGSEGV. Same guard the ``driver.build()`` path carries in
``cli/bootstrap``.

Construct a FastAPI app pre-bound to a built engine (in-proc modes).

The engine is built INSIDE the app's lifespan startup (so the
per-process CUDA context belongs to the uvicorn worker that's
going to drive it), and the run-loop task is spawned there too.

``driver_factory`` is the multi-rank entry path: under torchrun with
``world_size > 1``, rank 0 hands in a factory and the LIFESPAN builds
the :class:`DistributedEngineDriver` — on the boot thread, after
uvicorn is already serving. That ordering is what makes a TP boot
watchable: the driver's build is collective and, cold, is the whole
boot, so building it before uvicorn left rank 0's port unbound for
all of it. Worker ranks build on their own main thread; they bind no
HTTP, so they need no lifespan.

``prebuilt_driver`` takes an ALREADY-BUILT driver instead. It has no
build left to watch, so it is the one path that cannot start HTTP
early. Either way the driver delegates every Engine attribute read
via ``__getattr__``, so server code keeps accessing
``app.state.engine`` uniformly.

Process mode (``ARBI_ENGINE_PROC``) does NOT use this — the API
child builds via :func:`build_api_app` and the engine never hosts
HTTP (docs/engine_core_process.md §1).

Construct the API-CHILD FastAPI app (engine_core_process.md §1).

Same middleware / routers / tags as :func:`build_app`, but the app
is engine-FREE: ``app.state.engine_client`` (the
:class:`~arbi_serve.engine.proc.client.EngineClient`) replaces the
live engine, the tokenizer pool is built from the child's OWN
``Tokenizer(cfg.model.path)`` (§4 — tokenize lives API-side), and
the lifespan neither builds an engine nor installs drain signal
handlers (§7: the child's SIGTERM defers to the engine sequencing,
wired in ``engine/proc/api_child.py``). ``/health`` serves 503
``build_in_progress`` from the client caches until the engine's
``ReadyMsg`` lands (§3 of the P4 map — the boot-inversion contract).

``startup_hook(loop)`` / ``shutdown_hook()`` are the api_child
runtime's IO-thread lifecycle (they need the running HTTP loop).
Importing this module in the child must never pull the engine/torch
chain — the Engine import above is TYPE_CHECKING-only and every
in-proc engine import is function-local.

API-child lifespan: telemetry + tokenizer + client, NO live engine.

Startup completes immediately (uvicorn then accepts traffic), so
``/health`` is reachable while the engine process is still building
— the §1 boot inversion. ``app.state.engine`` is the P4b
:class:`~arbi_serve.engine.proc.api_engine.ProcEngine` facade — the
routes' engine surface re-expressed over the wire (submit → SubmitMsg
+ SubmitAck, cancel → CancelMsg, admission from the StatsMsg cache) —
NEVER the engine/torch chain. Engine-derived state (semaphore
capacity from ``ReadyMsg.inflight_limit``, tokenizer-fingerprint
validation) is event-driven off the sink's ``on_ready``.

Shared FastAPI assembly: middleware, routers, semaphore, tags.

Used by BOTH :func:`build_app` (in-proc engine lifespan) and
:func:`build_api_app` (API-child lifespan) so the HTTP surface can
never drift between the modes — only the lifespan differs.

Authn-z + per-token rate limiting + cumulative quotas.

Auth backends:
  - **tokens-file**: newline-delimited bearer tokens; one token per
    line. Optional ``rps=...,tps=...,daily=...`` annotations after a
    whitespace separator (e.g. ``mytoken rps=10 tps=2000 daily=1000000``).
    For dev / smoke testing.
  - **verifier-module**: dotted path ``mod.sub:callable``; the
    callable receives the raw bearer token and returns either:
      - ``None`` / raises → token rejected (401);
      - ``UserAuth`` instance → token accepted with attached
        per-user limits + quotas.

Status codes (binding contract):
  - 401: ``Authorization`` header missing, malformed, or token unknown.
  - 403: token valid but daily quota exhausted.
  - 429: token valid but rate (req/s or tokens/s) limit exceeded.

The state dict (per-token TokenState) is asyncio.Lock-protected end-
to-end; multiple concurrent requests for the same token serialize
on that token's lock. Different tokens are independent.

Quota persistence is best-effort: writes go through a single shared
``asyncio.Lock`` and re-serialize the full state dict on every check.
That's fine for dev-scale volumes; production can swap in a Redis or
SQL backend by writing a custom verifier module that hands the
engine a different ``QuotaStore``.

Admin privilege separation (the security model)
-----------------------------------------------
A valid bearer token is NOT automatically an operator. Admin power
(``/v1/admin/*`` — model reload, memory release, config mutation,
flat-weight / snapshot dumps, profiler control) is gated behind an
explicit ``is_admin`` scope on the auth identity:

  - **Authenticated mode** (a verifier is configured): a token may
    reach an admin route only if its identity carries ``is_admin=True``.
    The tokens-file backend sources this from an explicit ``admin=true``
    annotation on the token's line; verifier modules set
    ``UserAuth.is_admin`` themselves. Everything defaults to
    non-admin — a plain serving token gets ``403`` on every admin path.
  - **Open mode** (no verifier configured): admin routes are
    FAIL-CLOSED. They are refused (``403``) UNLESS the request comes
    from loopback (local dev/bench tooling on the same host) OR the
    operator has explicitly opted in by setting
    ``ARBI_ALLOW_UNAUTHENTICATED_ADMIN=1``. A loud warning is logged the
    first time an unauthenticated admin call is allowed.

Enforcement is centralized in :class:`AuthMiddleware` keyed on the
``/v1/admin`` path prefix (defense-in-depth: a newly-added admin route
is gated by construction, not by remembering a per-route dependency).

Filesystem-path allowlist
-------------------------
Admin endpoints that take an operator-supplied filesystem path
(``/v1/admin/model`` ``path``, ``/v1/admin/dump_flat_weights``
``out_dir``, ``/v1/admin/memory/record/dump`` ``path``) route it through
:func:`resolve_admin_path`, which resolves symlinks + ``..`` and rejects
anything escaping the configured allow-listed base directories
(safe defaults; override with ``ARBI_ADMIN_ALLOWED_DIRS``, a
``:``-separated list). This closes the arbitrary-path read/write hole
even for a legitimately-admin caller. Rejections are a loud ``400``.

Map an auth result to the request's :class:`TenantContext`.

Resolution order for ``tenant_id``:
  1. ``explicit_tenant_id`` (``X-Cache-Tenant`` header) — operator
     opt-in for cross-token cache sharing or deliberate isolation.
  2. ``user.user_id`` — auth-backed default.
  3. ``bearer_token`` — fallback if auth is open-mode but a token
     was presented (legacy behaviour: tenant == bearer token).
  4. ``"anon"`` — no auth, no header → :data:`ANONYMOUS_TENANT`
     is returned as-is.

The :class:`QuotaBudget` lifted from ``user.rate_limits`` /
``user.quotas`` is what the scheduler's admission filter
consults. ``rps`` is intentionally NOT folded in — that gate
runs in the HTTP middleware (``AuthManager.authorize``); the
scheduler enforces ``tps`` (output tokens/sec) which the
middleware can't know upfront.

Returns a fresh :class:`TenantContext` (frozen). The
``backend_overrides`` slot stays at its empty default; no
deployment-level A/B routing is applied.

Load a tokens file and return a verifier that looks tokens up.

Each non-blank, non-comment line is a token, optionally followed
by ``key=value`` annotations: ``rps``, ``tps``, ``daily``, ``admin``.
Example::

    # dev tokens
    sk-dev-alice rps=5 tps=1000 daily=100000
    sk-dev-bob
    sk-operator admin=true

A token without any annotations gets unlimited rate and quota and
is NON-admin (suitable for smoke testing; production sets all three
limits). ``admin=true`` (also ``1``/``yes``/``on``) grants operator
scope — the ONLY way a tokens-file token reaches ``/v1/admin/*``.

Open-mode admin policy: allow only loopback or an explicit opt-in.

Used only when NO verifier is configured. Returns ``True`` when the
request originates from loopback OR
``ARBI_ALLOW_UNAUTHENTICATED_ADMIN`` is set; ``False`` otherwise
(fail-closed — the public hole is closed by default while local
dev/bench tooling keeps working).

Resolve an operator-supplied admin path and enforce the allowlist.

Resolves symlinks + ``..`` and confirms the result is one of, or
nested under, the allow-listed base directories
(:func:`admin_allowed_base_dirs`). Returns the resolved absolute
path string on success; raises ``HTTPException(400)`` (loud) when
the path is empty or escapes every allowed base.

Stable, non-reversible id for a bearer token.

First 12 hex chars of SHA-256. Lets traces/logs/dashboards group
and search by which credential made a request WITHOUT ever
exposing the secret. Stable across requests for the same token, so
a Grafana/Tempo query on ``arbi.key_id`` returns that key's whole
request history.

Best-effort: stamp the active inbound HTTP span with attributes.

The FastAPI auto-instrumentation span (see
:func:`arbi_serve.server.app._server_request_hook`) is the current
span when this middleware runs, so these attributes attach to the
per-request trace Tempo/Grafana already render. No-op if no
recording span is active (e.g. tracing disabled). Never raises into
the request path.

Validate Authorization: Bearer <token> against the engine's AuthManager.

Skips paths in ``_AUTH_EXEMPT``. When auth is not configured on
the engine, it short-circuits to allow all (with the startup
warning printed once at boot in cli.py).

A **pure ASGI** middleware, deliberately NOT ``BaseHTTPMiddleware``:
the latter pumps every streamed response chunk through an anyio
memory-object-stream + task group, and stacked across the chain that
starves the engine thread of the GIL at concurrency. See
:class:`arbi_serve.server.middleware.RequestIDMiddleware`.

Decide whether ``token`` may proceed.

Returns ``(user_or_None, status_code, reason)``. Status codes:
  - 200 → allowed
  - 401 → unknown / invalid token
  - 429 → rate limited
Quota (403) is checked after the request finishes via
:meth:`record_completion` which doesn't gate admission.

Update tps bucket + daily quota after a request finishes.

Called by the engine wrapper at finish; doesn't reject — the
hard reject was at admission time. This bookkeeping informs
the next admission decision.

OpenAI-compatible Batch API: ``/v1/files`` + ``/v1/batches``.

The "submit → poll → fetch" async surface. A client uploads a JSONL
file (one request per line), creates a batch over it, polls the batch
until ``completed``, then downloads the output (and error) file. One
endpoint covers chat + text completions + embeddings via each line's
per-request ``url`` field.

Architecture
------------
* **Reuse, not reimplementation.** Each input line is dispatched to
  the same callable cores the single-request HTTP routes use —
  :func:`arbi_serve.server.api._run_chat` /
  :func:`~arbi_serve.server.api._run_completion` /
  :func:`~arbi_serve.server.api._run_embeddings`. Those submit through
  ``eng.asubmit``, so the engine **continuous-batches** the lines
  alongside any live traffic. No duplicate generation / embedding
  logic lives here.

* **Task validation.** One model/task per arbi instance. A
  ``--task generate`` instance serves lines whose ``url`` is
  ``/v1/chat/completions`` or ``/v1/completions``; ``--task embed``
  serves ``/v1/embeddings``. The batch-level ``endpoint`` and each
  line's ``url`` are validated against the instance task. A
  mismatched batch ``endpoint`` is rejected at create time (400); a
  mismatched per-line ``url`` becomes a per-line error row.

* **Concurrency.** Lines are processed with a bounded
  ``asyncio.Semaphore`` (default :data:`_MAX_INFLIGHT_LINES`) so the
  worker doesn't fan thousands of coroutines into the engine intake at
  once — but real throughput is governed by the engine's continuous
  batcher, which schedules as many of those in-flight requests as the
  KV budget allows. Results are reassembled in input order.

* **Durability.** The batch / file registries are in-memory and
  process-local (see :mod:`arbi_serve.server.batch_store`). State does
  NOT survive a restart, and ``completion_window`` is accepted but
  processing starts immediately (no 24h deferral). OpenAI's service is
  durable + deferred; this is the documented single-instance gap.

Upload a file. For the Batch API, ``purpose='batch'`` + a JSONL body.

Accepts an OpenAI-style ``multipart/form-data`` upload (``file`` part +
``purpose`` field). The multipart body is parsed in-process (see
:func:`_parse_multipart`) so the image does not need the optional
``python-multipart`` dependency. As a convenience, a raw JSONL body with
``Content-Type: application/x-ndjson`` (or ``application/jsonl``) is also
accepted, with ``purpose`` taken from the ``?purpose=`` query param.

The JSONL is not parsed here — validation happens when a batch is
created over the file, so a malformed body surfaces as a batch error
rather than an upload 400 (mirrors OpenAI).

Parse a ``multipart/form-data`` body into ``(filename, purpose, file_bytes)``.

Minimal RFC 7578 parser — enough for the OpenAI ``files.create``
shape (one ``file`` part + scalar text fields like ``purpose``). Avoids
the optional ``python-multipart`` dependency. Raises 400 on a malformed
body or a missing boundary.

Create a batch over an uploaded input file and kick off processing.

Validates that the input file exists and that the batch ``endpoint``
matches this instance's serving task; then spawns the background
worker and returns the batch object in the ``validating`` state.

Request cancellation of an in-flight (or pending) batch.

Cooperative: sets a flag the worker checks between lines. Already
terminal batches are returned unchanged. A batch that hasn't started
its lines yet flips straight to ``cancelled``.

Run one input line through the matching engine core, returning its body.

Builds the appropriate request model from the line's ``body`` (filling
a missing ``model`` with the instance's served name) and calls the
shared core. ``HTTPException`` from the core propagates to the caller,
which records it as a per-line error.

Every batch line is submitted at ``"batch"`` priority: the scheduler
admits it only with spare capacity and yields it the moment interactive
traffic arrives. This is what makes the Batch API "top up the GPU when
free" without ever degrading real-time latency.

Background worker: process every input line and finalise the batch.

Reads the input JSONL, validates + dispatches each line through the
shared engine cores (bounded concurrency), then writes the output and
error JSONL files, sets ``output_file_id`` / ``error_file_id`` and the
request counts, and flips the batch to ``completed`` (or ``failed`` /
``cancelled``). Order + ``custom_id`` are preserved.

In-process storage + registry for the OpenAI-compatible Batch API.

Two object kinds, modelled after the OpenAI surface:

  * **File objects** (``file-...``) — uploaded JSONL inputs plus the
    output / error JSONL blobs the batch worker writes. Bytes live on
    disk under :attr:`BatchStore.root`; the metadata (id, filename,
    purpose, size, ``created_at``) lives in an in-memory dict.
  * **Batch objects** (``batch_...``) — the async job. State machine:
    ``validating → in_progress → finalizing → completed`` (or
    ``failed`` / ``cancelled`` / ``expired``). Tracked entirely in
    memory.

Durability caveat
-----------------
Unlike OpenAI's Batch API (backed by durable object storage + a job
DB), this registry is **process-local and volatile**: the in-memory
maps are lost on restart, and orphaned blobs are left on disk. This is
a deliberate single-instance simplification — see the module that owns
the worker (:mod:`arbi_serve.server.batch_api`) for the rationale. For
a durable deployment, swap these dicts for a small KV / SQLite table
and re-scan ``root`` on boot.

Concurrency
-----------
A single :class:`asyncio.Lock` guards mutations to the registries. The
worker holds it only for the brief status flips / counter bumps, never
across the (slow) per-line engine work — so polling (``GET``) stays
responsive while a batch runs.

Process-local registry for Batch API files + batch jobs.

Not durable across restart (see module docstring). Thread-safety is
via a single asyncio lock; the store assumes a single event loop
(the uvicorn worker's), which holds for arbi-serve's one-process
deployment.

The serving gate: HTTP is up, serving is not.

The HTTP surface binds before the engine exists so an operator can watch
the boot on ``/health`` and ``/metrics`` instead of an unreachable port.
That leaves a window in which the process answers but cannot serve a
token, and a request that arrives in it must be REFUSED — never queued,
never hung on a half-built engine.

:class:`ServingGateMiddleware` is that refusal. It turns every non-exempt
request into a 503 naming ``engine_booting`` and the phase the boot is in,
so a caller (or a load balancer) can act on it. The exempt set is the
instrument surface — liveness / readiness / the Prometheus scrape — which
is the whole point of binding early.

The gate is INERT unless a lifespan installs one: an app whose
``state.serving_gate`` is unset (every schema-only / test-built app) is
never gated.

Latch that opens once the engine can serve.

Starts CLOSED. The lifespan opens it after the engine is attached and
its run loop is up; nothing re-closes it (a drain / reload is the
engine's own admission gate, not this one).

One line naming what the boot is doing and how far in it is.

The position wording is :func:`boot_progress.describe_position` — the
single renderer this and ``/health/ready``'s detail both use, so a
caller's 503 and the probe an operator curls can never tell different
stories. It reports only what the registry observed: a message that
reads as a measurement and is a guess is worse than no message.

Refuse every non-exempt request until the serving gate opens.

Raw ASGI (not ``BaseHTTPMiddleware``): it stays mounted for the life
of the process, so the open-gate path must cost one attribute read
and a jump, not a per-request task.

A refusal is counted every time and LOGGED once per path — a probe
that retries every second through a long boot would otherwise bury
the boot log it exists to make readable.

Boot-progress gauges, registered before any engine exists.

Exports gauges:

  - ``arbi_serve.boot_elapsed_seconds`` — seconds since the boot started,
    labelled by boot state
  - ``arbi_serve.startup_phase_active_seconds`` — seconds the currently-open
    boot phase has been running, labelled by phase

Both read :mod:`arbi_serve.boot_progress`, which is torch-free and
process-global, so they carry no engine dependency and register at
telemetry-install time. That is what makes them present on the FIRST
scrape of a booting process — including a TP>1 rank 0, whose engine is
built by the lifespan and therefore does not exist to bind against while
the collective rendezvous runs.

They are the in-flight companion to ``arbi_serve.startup_phase_seconds``,
which reports a phase only once it closes.

Cancellation primitives for streaming SSE handlers.

Two complementary patterns:

  - :func:`safe_run`: an async context manager that wraps a streaming
    engine call. On :class:`asyncio.CancelledError` (the FastAPI /
    Starlette signal when the client disconnects mid-stream), the
    cleanup coroutine is invoked under :func:`asyncio.shield` so it
    survives the cancel. The error is re-raised as
    :class:`SafeRunException` so the caller distinguishes "client gave
    up cleanly" from an unhandled engine bug.

  - :class:`DisconnectHandler`: a background ``asyncio.Task`` that
    polls the FastAPI request's ``is_disconnected()`` at ~20 Hz and
    triggers a cleanup callback the instant the client TCP socket goes
    away. The streaming generator can be deep inside an ``await`` on a
    token event when that happens — without an active poller, FastAPI
    does not propagate the disconnect into the generator coroutine
    until the NEXT yield, which can be many seconds away on a stalled
    generation.

Both are infrastructural: the engine is unaware. The handler / wrapper
calls :meth:`Engine.cancel` (existing surface) on the request handle
to free pages + drop the request from the run loop.

Sentinel raised by :func:`safe_run` after a handled cancel.

Re-raised in lieu of :class:`asyncio.CancelledError` so callers can
distinguish "the client disconnected and we cleaned up" from
"something inside the engine raised an unrelated exception".

Streaming HTTP handlers typically swallow this — the connection is
already closed; there is no client to inform.

Async context manager guarding a streaming engine call.

Usage:

    async with safe_run(cleanup=lambda: engine.acancel(req)):
        async for token in engine.iter_tokens(req):
            yield sse_chunk(token)

On :class:`asyncio.CancelledError` (client disconnect, server
shutdown, etc.):

  1. Run ``cleanup()`` under :func:`asyncio.shield` so the cancel
     doesn't propagate into the cleanup coroutine itself. Cleanup
     is the engine's chance to drop in-flight pages, release LoRA
     refs, and remove the request from the queue — losing it
     leaks GPU memory.
  2. Re-raise as :class:`SafeRunException` (NOT the original
     ``CancelledError``). Starlette's ``StreamingResponse``
     interprets a re-raised ``CancelledError`` as "the framework
     is also shutting down"; raising a regular exception lets
     normal exception-handling flow.

Raising any other exception from inside the ``async with`` body
propagates unchanged — :func:`safe_run` is purely a cancel
handler.

Internal helper — invoke cleanup, swallow the inner cancel.

Even with :func:`asyncio.shield`, the wrapped task can still
receive a cancel if the *outer* task was the one that set the
signal. ``SafeRunException`` itself.

Background poller that watches a FastAPI request for disconnects.

The streaming generator coroutine doesn't naturally see a
disconnect until its next ``await`` returns to the event loop AND
the framework happens to check the receive channel.
For long-running generations (multi-second TPOT, no inter-token
yields), that means the engine keeps running a generation whose
client is gone.

Solution: a sidecar task polls
:meth:`fastapi.Request.is_disconnected` at ~20 Hz. When the client
is gone, it invokes the cleanup callback (typically
:meth:`Engine.cancel`) and exits.

Lifetime: caller spawns one per streaming request, awaits / cancels
it when the generation finishes. ``async with`` is the typical
shape::

    async with DisconnectHandler(request, on_disconnect):
        async for chunk in stream():
            yield chunk

Or call :meth:`start` / :meth:`stop` manually.

Semaphore-first concurrency control for OpenAI handler endpoints.

Every inbound request acquires a permit from a fixed-size semaphore at
the entry point. If the semaphore is empty, the server returns HTTP
429 **immediately** (no queueing — backpressure is the client's
problem). The permit is released even on exception.

Pair with the engine-level admission gate: :func:`_enforce_admission`
already returns 429 when the engine queue is too deep. The semaphore
is a **second**, smaller cap — it bounds the count of concurrent
*coroutines* in HTTP-handler land, regardless of whether they're
sitting in the engine queue. That matters because each open SSE
connection holds open file descriptors, a chunked-encoder, an
iterator state machine, etc.; we want a hard ceiling we can size
independent of the engine.

The middleware path (registered via :func:`build_app`) is the
canonical entry. The :class:`RequestSemaphore` itself is a thin
wrapper around :class:`asyncio.Semaphore` so tests + future
direct-dependency consumers can use it without going through
middleware.

Priority-aware bounded permit pool with fail-fast acquire.

Standalone primitive usable from middleware AND from inline
endpoint bodies (FastAPI dependency, manual ``async with`` block).

**Priority-awareness (the whole point).** A single counter bounds
total in-flight at ``capacity``, but the two priority classes see
*different* admission thresholds against that one counter:

  * **interactive** acquires whenever ``in_flight < capacity`` —
    it may take the full pool, including the reserved headroom band.
  * **batch** acquires only while ``in_flight < batch_capacity``
    (a lower watermark, :data:`DEFAULT_BATCH_CAPACITY_FRACTION`
    of ``capacity`` by default). Once batch has filled its
    watermark, further batch requests are refused (HTTP 429) even
    though interactive can still get in.

The band ``[batch_capacity, capacity)`` is therefore
interactive-only headroom: a 128-wide batch fan-out fills the
batch watermark and gets 429'd, but every interactive probe still
acquires immediately. Total in-flight never exceeds ``capacity``,
so memory stays bounded.

Failure mode is **not blocking**: when no permit is available for
the request's class, :meth:`try_acquire` returns ``False``
immediately and the middleware converts it into HTTP 429.

Attributes
----------
capacity
    Maximum concurrent permits (interactive ceiling).
batch_capacity
    Watermark above which batch requests are refused.

Pure-ASGI middleware that fronts every request with a semaphore.

On overload returns ``HTTP 429`` with a ``Retry-After`` header. The
body matches the OpenAI error envelope used by the rest of the
server (``{"error": {"message": ..., "type": "backpressure"}}``)
so existing client retry logic keeps working.

Excluded paths bypass the semaphore — health probes and metric
scrapes must never trip on 429 even at saturation.

A **pure ASGI** middleware, deliberately NOT ``BaseHTTPMiddleware``:
the latter pumps every streamed response chunk through an anyio
memory-object-stream + task group, and stacked across the chain that
starves the engine thread of the GIL at concurrency (see
:class:`arbi_serve.server.middleware.RequestIDMiddleware`). The
semaphore acquire/release wraps ``self.app`` directly here.

Build the pool.

Parameters
----------
capacity
    Positive total permit ceiling (the interactive ceiling).
batch_capacity
    Explicit batch watermark. When omitted, derived as
    ``round(capacity * batch_capacity_fraction)``, clamped to
    ``[1, capacity]``. Must satisfy ``1 <= batch_capacity <=
    capacity``; a watermark equal to ``capacity`` disables the
    reservation (batch == interactive, legacy behaviour).
batch_capacity_fraction
    Fraction used to derive ``batch_capacity`` when it is not
    given explicitly.

Acquire one permit for the given priority class.

The fail-fast variant. Returns ``True`` on success, or
``False`` immediately when the class's threshold is exhausted.

Parameters
----------
interactive
    ``True`` (the safe default) acquires against the full
    ``capacity``. ``False`` (batch) acquires only while
    ``in_flight < batch_capacity`` — so batch can never consume
    the interactive-reserved headroom band.

Re-point the pool at a new total capacity (boot-time only).

Used to reconcile the HTTP semaphore to the engine's resolved
``inflight_limit`` (the engine builds AFTER the middleware is wired,
and an embed/rerank fan-out inflates the limit past the config's own
``max_batch`` arithmetic).
Rebuilds the underlying counter — safe ONLY before any traffic /
while no permits are outstanding (server lifespan startup, before the
app accepts requests), which is the sole call site. The batch-priority
watermark is re-derived from the new capacity.

Try-acquire context manager for the given priority class.

Yields ``True`` if the permit was acquired, ``False`` otherwise.
Always releases on exit (only when acquired). Use::

    async with sem.slot(interactive=is_interactive) as got:
        if not got:
            return Response(status_code=429)
        ...

Wrap ``app`` with a priority-aware request semaphore, 429 retry delay, and exempt-path prefixes.

Pass a pre-built ``semaphore`` to share the instance with the caller
(so it can ``set_capacity`` to reconcile against the engine's resolved
``inflight_limit`` at boot); otherwise one is built from ``capacity``.

Resolve the admission priority class from request headers.

Mirrors :func:`arbi_serve.server.api._resolve_priority`'s
header surface: ``X-Arbi-Priority: batch`` (case-insensitive)
marks the request as batch; anything else (including an absent
header) is interactive. The header can only LOWER priority, so
a client on an untrusted edge cannot use it to grab the
interactive-reserved headroom — exactly the same guarantee the
scheduler relies on.

The body ``priority`` field is not visible at the middleware
layer (the body is still an unread stream), but that is safe:
the body can also only set ``batch`` (never raise), and any
request that wants batch admission treatment can carry the
header. The Batch API worker submits in-process (bypassing this
middleware entirely), so it is unaffected.

Acquire a permit for non-exempt requests; emit HTTP 429 on overload, else forward.

Interactive requests acquire against the full capacity; batch
requests (``X-Arbi-Priority: batch``) acquire only against the
lower batch watermark, so a batch fan-out can never 429 an
interactive request out of admission.

The container healthcheck: readiness to BECOME healthy, liveness to STAY.

A container healthcheck has two jobs that want different probes, and asking
one probe to do both is what forced the retry budget to be a second boot
budget.

BECOMING HEALTHY is readiness' question. ``/health`` answers 200 from the
moment uvicorn binds and stays 200 through the whole weight load, so a boot
that HANGS reads healthy forever once ``start_period`` expires; only
``/health/ready`` proves the model is built and the loop is stepping.

STAYING HEALTHY is liveness' question, and readiness cannot answer it. After
the container is healthy, a 503 from ``/health/ready`` is ambiguous: a wedged
run loop, a drain, and an operator's ``POST /v1/admin/model mode=reload``
rebuilding the model all produce it, and this repo budgets a reload at 900 s.
So a healthcheck that treats every readiness 503 as a fault must tolerate a
window longer than a reload -- and a window that long is also how long a
server that will never answer another request goes on presenting as fine.
``/health`` has no such ambiguity: it is 200 through boot, reload and sleep,
and 503 only for a dead engine thread or a run loop that has not ticked in
``LIVENESS_STALL_S``.

Hence the latch. The first readiness pass writes a marker in the container's
own filesystem, and from then on this probe asks liveness alone. The marker
dies with the container, so a recreated or restarted container gates on
readiness again from scratch.

What this gives up, deliberately: the healthcheck no longer polices a
deliberate, operator-initiated rebuild. It could never do that correctly --
over HTTP a reload's 503 and a wedge's 503 are the same answer, which is the
reason the window was being asked to tell them apart. The operator's own
request is what is waiting on a reload.

Exit 0 healthy, 1 unhealthy. Nothing here can exit 0 without having read a
200 out of the server in this invocation.

Engine build / boot-guard / warm-cache helpers for the app factory.

Covers the engine-side build path (:func:`abuild_engine_core`, shared
with process mode), the single-rank / prebuilt-driver attach helpers,
stable-VA residency + pool preparation, the audio-decoder preload, the
MTP-on-tkv wiring guard, and the MTP-K boot sweep. Imported back into
``server/app.py`` so the public surface (and the process-mode
``engine_bindings`` imports) are unchanged. Kernel-module pre-load is
the "serve-kernel warmup" boot phase in ``engine/build.py`` (pre-ready),
not a task here.

Resolve the tokenizer pool worker count (config > env knob > default).

Flag-truth (ARBI_TOKENIZER_POOL_WORKERS): fires only when the env knob
actually sets the count — a config value shadows it (refused), and 0
(the compiled default) derives the built-in default.

``(entries, max_ids)`` for the pool's growing-prefix cache.

Neither number is this module's to invent. ``entries`` is how many
CONVERSATIONS should stay warm at once, and the server already states that:
it is the number of requests it admits concurrently, the same
``inflight_limit`` the HTTP semaphore and the engine admission gate share.
``max_ids`` follows from that and the context those requests are admitted
at — the cache can hold at most one whole prompt per warm conversation, so
the id budget is ``entries x max_context``, i.e. the same order as the ids
the server is already holding for the requests themselves.

A boot whose ``max_context`` has not resolved to a concrete number yields
the SINGLE-entry cache this shipped with, and refuses its counter with the
reason: an id budget it cannot derive is one it would have to guess, and a
guessed budget is a memory bound nobody can check.

Point the NVML true-driver-floor gauge at this engine's sleep pool.

Lets nvml_metrics subtract cuMem-mapped model memory from driver
overhead without importing the engine (it only calls the lambda).
Best-effort: an engine without a ``sleep_pool`` simply leaves the
provider unset, so true_driver_floor == driver_overhead.

Only torch-UNTRACKED cuMem bytes belong here — bytes that sit inside
``framebuffer − torch.cuda.memory_reserved`` (the "driver overhead"
the floor formula starts from):

  * ``sleep_pool.mapped_bytes`` — per-tensor stable-VA registrations,
    allocated directly via the cuMem driver and wrapped through a CAI
    shim, invisible to the caching allocator;
  * the B2 growable KV slab (``GrowableRegion``) — its own VA
    reservation mapped via ``map_range``, also outside torch's stats.

The cuMem PLUGGABLE-ALLOCATOR named pools (weights / kv / graph / …)
must NOT be added: they are ``torch.cuda.MemPool``-backed, so torch's
caching allocator already counts their mapped bytes in
``memory_allocated`` / ``memory_reserved`` — the same bytes
``CuMemPoolAllocator.mapped_bytes()`` reports. Including
``CuMemPoolAllocator.mapped_bytes()`` here double-subtracts
them from driver overhead, clamping ``true_driver_floor`` to 0 and
inflating the reconciliation residual by ~the pool footprint on every
default (cuMem) boot.

Start the engine run-loop, returning ``(engine_thread, run_task)``.

Own-thread mode (default, ``ARBI_ENGINE_OWN_THREAD``) drives the loop
on a dedicated :class:`EngineThread` so the engine's per-step Python
bookkeeping never holds the GIL on the HTTP/SSE loop. The bridge is
always ``eng._loop_bridge`` (pass ``driver.engine`` as ``engine_obj``
for the distributed-driver path). ``ARBI_ENGINE_OWN_THREAD=0`` falls
back to an in-line ``create_task`` on the HTTP loop (warning-logged,
never silent). Exactly one of the returned pair is non-``None``.

Both halves route a run-loop exit through the SAME classifier
(:mod:`arbi_serve.engine.run_loop_exit`): own-thread from the thread's
own teardown, in-line from the task's done-callback. Neither container
ends the process by itself -- a dead thread and an abandoned task both
leave uvicorn answering for an engine that can never step again -- so
the terminal case exits deliberately.

In-process TKV calibration hook — the SAME hook on every boot path.

If boot deferred calibration (it booted ``tkv-bypass`` because the requested
codec had no complete bundle), calibrate on the LIVE engine now — its
run-loop is up — then hot-swap to the codec. One model load, no child, fails
safe (bypass keeps serving) on a bad bundle. No-op when boot did not defer.

Called from BOTH boot paths so the behaviour is identical at every topology:
:func:`_build_and_attach_engine` (single-rank) and
:func:`_attach_prebuilt_driver` (torchrun multi-rank, rank 0). The
calibration itself is rank-0-driven and issues no out-of-band collective —
its forwards reach the worker ranks as ordinary slate deltas and its swap as
a control op, so the workers never leave their serving loop (see
``calibration.inprocess``). ``driver`` is the distributed driver on the
multi-rank path, used to stop the workers cleanly on ``--calibrate-only``.

Attach an already-BUILT distributed driver (multi-rank rank 0).

The app-side half of the TP>1 boot, shared by both routes into it:
the driver the lifespan just built on the boot thread
(``driver_factory``), and one built before the app existed
(``prebuilt_driver``). Binds metrics + the cuMem provider, stashes
the driver on ``app.state``, starts its run-loop, and installs signal
handlers. Returns ``(engine_thread, run_task, eng)``.

Engage stable-VA residency for the freshly-booted model (+ pool).

Wraps the boot model as the active resident record so the routing
layer (``_maybe_switch_model``) can auto-switch to other prepared
residents recapture-free. With declared ``pool_members`` the boot
model PLUS each member is prepared (built + captured under its own
tag namespace, then parked in host RAM — single-VRAM-resident).
No-op when ``stable_va_residency`` is off — the default deployment is
byte-identical.

Eagerly build the token2wav decoder when speech output is prepared.

Gated on the affirmative ``ARBI_ENABLE_AUDIO`` flag (the single audio
modality switch — input tower AND speech output). When it is off (the
default) a text-only boot SKIPS this entirely: no token2wav weights, no
vocoder warmup, no DiT-graph capture — saving the boot time and the VRAM
the decoder would otherwise reserve before engine build. When it is on, the
preload runs BEFORE the engine builds so the decoder's resident + captured
footprint is already accounted when measure-then-fit sizes the KV pool (this
is why ``--gpu-memory-utilization 0.99`` boots without OOM even with audio).

Consulted signals (beyond the flag): the checkpoint ships ``token2wav/``
assets AND a voices dir with at least one prepared ``.pt`` cache (either
``ARBI_AUDIO_VOICES_DIR`` or the default ``<model>/token2wav/voices``).
On any miss this is a no-op and the decoder stays lazy (an audio
request then 400s with the prepare-a-voice message).

Run ``fn`` on a dedicated DAEMON thread; return an awaitable future.

Not a :class:`~concurrent.futures.ThreadPoolExecutor`: its workers are
non-daemon and ``concurrent.futures`` joins them at interpreter exit,
so a SIGTERM half-way through a boot would block shutdown until the
build finished. A daemon thread lets the process leave whenever the
HTTP side is done, which is the whole point of being able to stop a
boot. Named, so a stack dump of a slow boot says which thread is
building.

Construct the Engine, preload the audio decoder, build, engage residency.

The blocking, CUDA-bound half of :func:`abuild_engine_core`, factored
out so it can run either on the calling loop's thread or on the boot
thread. Ordering is load-bearing:

* the token2wav decoder preloads BEFORE the build fills VRAM, so the
  engine's KV-pool sizing and cudagraph capture budget around it
  instead of leaving a lazily-loaded decoder to OOM at the first
  audio request. Engaged only when a prepared voices dir exists;
  otherwise nothing loads and text boots are untouched;
* stable-VA residency engages AFTER the build, wrapping the booted
  model as the active resident so the routing layer can auto-switch
  to other prepared residents recapture-free. No-op when
  ``stable_va_residency`` is off — the default deployment is
  byte-identical.

Construct the multi-rank driver and run its COLLECTIVE build.

The blocking half of the TP>1 rank-0 boot, factored out so it can run
on the boot thread while rank 0's HTTP loop already serves the probes.
Every rank walks the same ``Engine.build``; the rendezvous inside it
(``init_distributed_environment``) is what the ranks meet at.

Running this off the main thread is sound and already the norm here:
the driver's own run loop issues its per-step NCCL collectives from
:class:`~arbi_serve.engine.engine_thread.EngineThread`, so a
communicator in this process is already driven from a non-main thread.
What matters is that a communicator is never used by two threads at
once — and it is not: the build completes before the run loop starts,
the serving gate is shut for the whole build, and no ``/metrics``
gauge callback issues a collective.

ONE engine-side build path shared by BOTH serve topologies.

In-proc modes call it from the lifespan
(:func:`_build_and_attach_engine`, ``audio_host = app.state``);
process mode (P4b, :mod:`arbi_serve.engine.proc.engine_bindings`)
calls it without a FastAPI app — ``audio_host`` is a plain
namespace there (the token2wav decoder is engine-process state, §4:
it is CUDA). Covers: Engine construction, the token2wav preload
(BEFORE the build fills VRAM so measure-then-fit budgets around
it), ``abuild``, stable-VA residency, metrics + cuMem provider
binding, and the CLI ``--lora`` preloads. Returns the built engine.

``offload_to_thread`` runs the blocking half (:func:`_construct_and_build`)
on the boot thread instead of the calling loop, so a caller that has
already started serving HTTP keeps its event loop responsive for the
whole boot. The engine is identical either way — only which thread
runs the CUDA-bound build differs.

Single-rank path: construct + build the Engine, attach + start it.

Delegates the engine-side build to :func:`abuild_engine_core`
(shared with process mode), then does the app-side half: attach to
``app.state``, start the run-loop, install signal handlers. Returns
``(engine_thread, run_task, eng)``.

``offload_to_thread`` is the early-HTTP path: the build runs on the
boot thread so this coroutine's loop keeps serving ``/health`` and
``/metrics`` throughout.

Bring the HTTP concurrency semaphore to the engine's resolved bound.

``max_inflight = INFLIGHT_PER_BATCH × max_batch`` is the SINGLE
admission number; the engine finalizes it at build() once
``--max-batch auto`` resolves, so the semaphore (wired before the
engine built) is brought into lock-step here — guaranteeing the HTTP
door and the engine queue gate enforce the SAME value, never a second
divergent cap. Runs at startup before any request, so no permits are
outstanding. Best-effort: a test/mock engine without ``inflight_limit``
leaves the construction-time value untouched.

Fail-loud MTP-on-tkv wiring guard (boot critical path).

When MTP is enabled on a tkv backend the split-K verify kernel
(``MTPFusedAttend`` / ``core._mtp_attend``) MUST be wired onto the
built cores, otherwise the (B, K+1) verify shape silently falls
through to the heavier Turbo prefill path — MTP runs, returns 200 OK,
captures a cudagraph, and is *slower* than no-MTP with no crash to
flag it. The cores are eagerly built during ``abuild``
(ensure_tkv_cores_built), so the wrapper is observable here. Run it
synchronously so a regression aborts the boot rather than degrading
throughput in production. No-op when MTP is off or the backend isn't
tkv.

MTP-K auto-pick sweep-on-miss (``--mtp-n-draft auto`` default).

Runs ONLY here — after the engine loop (run_forever) is live, so a
rank-0 ``asubmit`` drives every TP rank via the per-step plan
broadcast, and after cudagraph capture, so each swept K hits a
captured graph. ``build_mtp_driver`` flagged the engine + built the
driver at the sweep ceiling on a cache miss. No-op when the cache
hit, K was pinned explicitly, or MTP is off. Best-effort: a sweep
failure logs + falls back to the ceiling default rather than
aborting the boot.

Classify the in-line run task's exit (own-thread's counterpart).

The lifespan cancels this task on teardown, which is why
``cancelled()`` returns early rather than reporting a fault; every
other outcome carries the same meaning as the dedicated thread's,
so it goes to the same classifier. ``task.exception()`` also
retrieves the exception, so a terminal fault can never end up as a
bare "Task exception was never retrieved" at interpreter exit.

Liveness + readiness probes.

Both endpoints are auth-exempt; the auth middleware skips these paths.

  - ``GET /health``       — liveness; 200 once the FastAPI app is
    serving, INCLUDING while the engine is still building (that carries
    ``"booting": true`` plus the live boot snapshot). 503 only when the
    run loop is dead / wedged or the boot refused.
  - ``GET /health/ready`` — readiness; 200 only when the engine has
    finished :meth:`Engine.build`, the pool exists, a scheduler is
    attached to step the queue, the run-loop thread is alive, and the
    engine is not draining; otherwise 503. A boot in progress reports
    ``build_in_progress`` with the phase it is in, so a load balancer,
    a sibling, or the admin UI can tell "up but building" from "ready".

READINESS DOES NOT IMPLY LIVENESS, and a caller pointing an LB or a k8s
readinessProbe at ``/health/ready`` alone must know it. Readiness answers
every question that is DECIDABLE from state — including the two terminal
ones, a dead engine process (process mode) and a dead run-loop thread
(in-proc): both are permanent, and no load, however heavy, can produce
them. It does NOT answer the one question that needs a clock: a run loop
that is alive but has not TICKED for a while. That distinction is a
threshold, and the only threshold available (:data:`LIVENESS_STALL_S`)
cannot tell a wedged loop from a slow one — see :func:`health_ready`. So
a WEDGED loop reads 200 here and 503 on ``/health``. Probe BOTH (every
shipped compose recipe does; ``tests/test_compose_healthchecks.py``
holds it) — healthy means 200 from each.

One gate, both sources: :class:`_EngineView` + :func:`_not_serving` is
the single readiness implementation, and :func:`_loop_wedged` the single
wedge implementation. ``StatsMsg`` confirmation in process
mode.

Two engine sources (docs/engine_core_process.md §6):

  * **in-proc modes** (thread / inline) — ``app.state.engine`` is the
    live engine; probes read its attributes directly.
  * **process mode** — ``app.state.engine_client`` is set (API child);
    probes read the client's cached ``StatsMsg`` / ``ReadyMsg`` +
    engine liveness. NO RPC on the probe path; heartbeat staleness
    (> ~3 s) degrades to 503. During boot the child serves 503 (the §1
    boot inversion — in-proc modes are simply unreachable during
    build). Readiness is fail-closed: the boot ``ReadyMsg`` alone (a
    boot-once "build announced" message, re-emitted from cache on
    API-child respawn) does NOT flip ``/health/ready`` to 200 — a
    confirming ``StatsMsg`` from the engine loop (``ready`` +
    ``pool_built`` +

The engine-state answers both probes ask, from either engine source.

In-proc modes read the live engine object; process mode reads the API
child's cached ``StatsMsg``. The two adapters below
(:func:`_view_of_engine`, :func:`_view_of_stats`) are the only place
that difference is expressed, so a question asked of this view is
asked on BOTH deployments — one gate, two sources, nothing to keep in
sync by hand.

``engine_unhealthy_reason`` recomputed from a ``StatsMsg``.

Same precedence (sticky fatal -> sticky memory-exhausted -> sticky swap
-> self-healing infra);
the infra latch ships RAW and the API child applies the
:data:`INFRA_UNHEALTHY_WINDOW_S` decay itself from
``infra_failure_ts`` (the ``proc/stats.py`` builder contract) —
``monotonic`` is same-host comparable (§2).

The readiness 503 this engine state owes, or ``None`` when it serves.

THE readiness gate — both probe paths route through it, so every state
is answered identically whichever side of the process split the engine
is on. Each reason names a state something routing traffic would act on
differently.

True when an ACTIVE engine's run loop has not ticked in time.

The active gate and the staleness comparison, once, for both sources.
Boot (``ready`` False), reload (no scheduler), sleep
(``memory_released``) and drain all legitimately pause stepping, so they
are excluded and can never false-trip a restart. ``last_loop_tick`` is
CLOCK_MONOTONIC and same-host comparable either side (§2: the API child
is always same-host).

LIVENESS ONLY — ``/health/ready`` does not call this, deliberately:

  * The signal is a THRESHOLD on the age of a tick stamped at the top
    of each loop iteration, so that age IS the duration of one
    iteration. "Wedged" and "mid-prefill-wave under load" differ only
    in how large that number is, and how large it legitimately gets is
    a function of the model, the GPU, the chunk size and the batch —
    none of which :data:`LIVENESS_STALL_S` is parametrized on. It is
    sized as a ceiling no legitimate iteration approaches, which is
    the right shape for "restart it" and the wrong shape for "route
    away from it now".
  * The two consumers act on the answer differently. An orchestrator
    reading liveness restarts after ``retries`` consecutive failures,
    under a ``start_period`` sized for the boot; a load balancer
    reading readiness pulls the node on the first 503, with no such
    budget. And load that lengthens an iteration lengthens it on
    EVERY replica at once, so a readiness flip on this signal sheds
    the whole fleet together — a self-inflicted outage in place of a
    hang one probe pairing already catches.
  * The graded form is already exported for whoever wants to alert on
    a developing hang: ``arbi_serve.engine.loop_tick_age_seconds``,
    which shows the climb instead of answering yes/no at one age.

Closing the gap is therefore the compose/LB pairing (probe both), not
a second threshold reader — ``tests/test_compose_healthchecks.py``
holds the pairing for every shipped recipe.

One line saying where the boot is and how long it has been running.

The position wording is :func:`boot_progress.describe_position` — the
single renderer this and the serving gate's refusal both use, so the
probe and the 503 a caller gets can never tell different stories.

Liveness answer for a process whose engine has not been built yet.

200 with ``booting`` set while a boot is running: the process is alive
and doing the right thing, so an orchestrator's liveness probe must not
restart it. The boot snapshot rides along so the same probe an operator
curls tells them WHICH phase. A refused boot, or a boot that claims to
have finished with nothing attached, is 503.

Readiness answer for a process whose engine has not been built yet.

Always 503, and the reason distinguishes every state something routing
traffic would act on differently: ``engine_not_built`` (no boot has
started — a bare app), ``build_in_progress`` (one is running, with the
phase), ``boot_failed`` (it refused), ``engine_missing`` (it claims to
have finished, yet nothing is attached).

Liveness — 200 unless the engine run loop is dead or wedged.

Reports 503 (so an orchestrator with a healthcheck + restart policy
recycles the container) in two cases:

  * the dedicated engine thread has died (own-thread modes) — the run
    loop crashed out while the HTTP server stayed up;
  * the engine is ACTIVE (built, not sleeping, not draining) but its
    run loop has not ticked within :data:`LIVENESS_STALL_S` — a hang.

A hard process crash needs no probe — the container exits and the
``restart`` policy recycles it directly. Boot / reload / sleep are NOT
treated as unhealthy: :func:`_loop_wedged` requires an active engine,
and the loop ticks throughout those states anyway.

``engine_thread_dead`` is the one reason here that ALSO ends the
process: a terminal run-loop fault exits with
``EXIT_CODE_UNRECOVERABLE`` (see
:mod:`arbi_serve.engine.run_loop_exit`), so the restart happens on the
``restart`` policy's timescale rather than on ``interval * retries``.
That is deliberate, and it is what lets the deployment's healthcheck
window stay wide enough for a legitimate admin model reload: over HTTP
a reload's 503 and a dead engine's 503 are the same answer, so the
window cannot be the thing that tells them apart. This probe reports
the dead engine; the process exit acts on it.

This is the ONLY probe that inspects the run loop's tick: readiness
answers a wedged loop 200 (see :func:`health_ready` for why), so a
caller that acts on health must read this one too.

Process mode (API child): the same semantics from the cached
``StatsMsg`` + engine liveness — see :func:`_proc_health_live`.

Readiness — 200 only when a request would be served right now.

503 otherwise, with a ``reason`` naming the state a caller routing
traffic would act on: ``engine_not_built`` / ``build_in_progress``
(carrying the boot phase) / ``pool_not_built`` /
``scheduler_not_present`` / ``draining`` / ``engine_thread_dead`` /
``engine_process_dead`` / ``engine_infra_failure`` / ``boot_failed``.

READINESS DOES NOT IMPLY LIVENESS. It refuses every not-serving state
that is decidable without a clock — including a dead run-loop thread
(in-proc) and a dead engine process (API child) — but a run loop that
is alive and has simply STOPPED TICKING reads 200 here and 503 on
``/health``. Telling a wedged loop from a busy one is a threshold
judgement, and it belongs to the probe with a retry budget rather
than to one a load balancer acts on immediately
(:func:`_loop_wedged` carries the argument). So point an LB or a k8s
probe at BOTH endpoints: healthy is 200 from each.

Process mode (API child): the same gate from the cached ``StatsMsg``
+ engine liveness — see :func:`_proc_health_ready`.

Friendly, stable, unique per-process instance identity.

Co-resident arbi-serve instances share one OTEL collector, and nothing
in the default telemetry uniquely identifies a process — ``instance``
(the Prometheus scrape target), ``job`` and ``service.name`` all collapse
to one value, so every aggregated dashboard panel silently sums across
instances. We fix that by stamping a unique ``instance.name`` resource
attribute on every process.

The id is human-memorable (Docker-style ``adjective_noun_xxx``) so it reads
well in an ``$instance`` dropdown. A short random suffix means two instances
that roll the same word pair still never collide. The value is regex-safe
(``=~"$instance"`` in PromQL without escaping).

Resolution order:
  1. ``ARBI_INSTANCE_NAME`` env (explicit operator choice, used verbatim).
  2. ``adjective_noun_xxx`` generated from the curated word lists below.

Build this process's instance id: a short memorable name.

``ARBI_INSTANCE_NAME`` if set+non-empty, else a generated
``adjective_noun_xxx``. No timestamp — boot time is already obvious
from the Uptime field. The value is regex-safe for ``=~`` matching.

Graceful SIGTERM / SIGINT drain.

On SIGTERM (and SIGINT outside tests), set the engine's
``_terminating`` flag; admission then refuses fresh requests with
HTTP 503 and ``Retry-After: 30``. Wait for in-flight requests to
finish or for ``--shutdown-grace-s`` to elapse, then ``os._exit(0)``.

The signal handler is installed on the running asyncio loop via
:meth:`asyncio.AbstractEventLoop.add_signal_handler`, so the drain
coroutine runs as a normal task — no thread, no signal-context work.

Mode scope: these handlers belong to the IN-PROC modes (engine thread /
inline), where the HTTP loop and the engine share a process. In process
mode (``ARBI_ENGINE_PROC``, docs/engine_core_process.md §7) SIGTERM
sequencing is engine-side — ``engine/proc/supervisor.py::
run_sigterm_sequence`` implements the §7 order and the API child's
handler defers to it (``engine/proc/api_child.py``); neither installs
this module's drain.

Drain in-flight requests, then ``exit_fn(0)`` (default ``os._exit``).

Idempotent: calling twice is fine; the second call returns once
the first has completed because :attr:`Engine._terminating` short-
circuits the second pass before sleeping.

``exit_fn`` is the process-exit seam: ``None`` (the default) resolves
``os._exit`` AT CALL TIME — late binding, so tests' ``patch("os._exit")``
keeps working — and keeps today's in-proc behaviour; the P4b
engine-side §7 sequence reuses the drain logic with a non-exiting
hook (it must still flush + notify the API child before the process
may die).

Log-formatting helpers — print a number WITH just enough context.

The house rule (see ``docs/logging_clarity.md``): a server log line should be
readable by someone who has never seen arbi-serve internals. Numbers that beg
the question "out of what?" carry their denominator; everything else stays
short. The aim is INSIGHT, not verbosity — one contextualized number per
clause, plain words, no internal codenames.

These helpers keep call sites short and consistent instead of hand-rolling
``%``-arg arithmetic. Pure, no I/O, no torch.

    >>> gib_of(5.77 * 1024**3, 24.0 * 1024**3)
    >>> gib(1.6 * 1024**3)

Loguru-based structured logging.

Loguru integrates with the OTEL ``LoggingInstrumentor`` (which
decorates stdlib ``LogRecord`` objects with ``otelTraceID`` /
``otelSpanID`` attributes); loguru-only call sites additionally pull
the active span via ``opentelemetry.trace.get_current_span()`` so
they also pick up trace context.

CLI surface:

  ``--log-json`` → JSON-Lines sink on stdout (one record per line).
  (default)      → human-friendly console sink.

Stdlib ``logging`` is bridged into loguru via an
``InterceptHandler`` so every ``logging.getLogger(...).info(...)`` —
including uvicorn's ``uvicorn.error`` / ``uvicorn.access`` loggers —
flows through the same loguru sink.

Patcher: enrich each loguru record with trace + request context.

Mutates ``record["extra"]`` in place. Keys added:

- ``request_id`` — from the contextvar populated by RequestIDMiddleware
  (None outside a request scope, in which case the key is omitted).
- ``trace_id`` / ``span_id`` — from the active OTEL span (only when
  a recording span exists; otherwise omitted).

Strip SGR escapes so a machine sink stores the readable body.

A colour escape is for a terminal. Loki, a JSON line and an OTEL log
attribute all want the text, and the glyph survives on its own.

Loguru sink that emits one JSON object per record on stdout.

Standard shape: ``timestamp`` (ISO-8601 UTC), ``level``,
``logger`` (the loguru ``name`` field — the module / function),
``message``, plus ``request_id`` / ``trace_id`` / ``span_id`` when
present, plus any user-supplied ``extra`` fields (e.g.
``logger.bind(backend=...).info(...)``).

Loguru → OpenTelemetry Logs SDK bridge.

Takes each loguru record, builds a structured OTEL log entry with
severity / attributes / trace context, and emits it through the
SDK's ``LoggerProvider`` so the OTLP collector ships it to Loki.

Attached as a sink only when an OTLP endpoint is configured; the
JSON / console sinks remain on stdout in parallel so logs are
still visible in the local terminal.

``True`` when an OTLP collector endpoint is available — same
rule as the metrics / traces side. Without this, attaching the
OTEL log sink would just queue records into a black hole.

Reconfigure loguru + bridge stdlib ``logging`` into it.

Idempotent: removes existing loguru handlers + stdlib root
handlers before installing fresh, so repeated calls (e.g. from
tests) don't stack.

The OTEL ``LoggingInstrumentor`` is installed by the FastAPI
lifespan (see :func:`arbi_serve.server.app.build_app`) and is
independent of this function — it instruments stdlib ``LogRecord``
objects regardless of which sink ultimately renders them.

OTLP log forwarding (loguru → OTEL Logs SDK → collector → Loki):
automatically attached when an OTLP endpoint is configured. The
stdout sink stays in parallel so terminal output is unchanged.

OpenTelemetry metrics + tracing wiring.

Telemetry is pushed to an OTEL collector (OTLP HTTP by default, OTLP
gRPC opt-in via env); the collector fans out to Prometheus / Tempo /
Loki and Grafana renders it. The same instruments are also exposed for
direct scrape at the in-process ``/metrics`` Prometheus route (served
from ``server/routes/models.py`` via the ``PrometheusMetricReader``).

Surfaces:

- Counters: ``arbi_serve.requests.total{status,backend,kind}``,
  ``arbi_serve.tokens.in.total{backend}``,
  ``arbi_serve.tokens.out.total{backend}``,
  ``arbi_serve.hot_swap.total{kind}``,
  ``arbi_serve.cal_mismatch_refusal.total``,
  ``arbi_serve.mtp.proposed.total{backend}``,
  ``arbi_serve.mtp.accepted.total{backend}``,
  ``arbi_serve.prefix_cache.hit_total{kind}``,
  ``arbi_serve.prefix_cache.tokens_saved_total{kind}``,
  ``arbi_serve.prefix_cache.eviction_total{reason}``,
  ``arbi_serve.residency.switch_total{from_model,to_model}``,
  ``arbi_serve.admission.shed.total{layer,reason,status}`` — EVERY
  refused / shed request, from any gate (engine admission, HTTP
  concurrency semaphore, auth rate-limit / quota, per-tenant scheduler
  filter). Pairs with the WARNING log line each refusal also emits
  (:mod:`arbi_serve.server.shed`), which carries the governing
  threshold; without both, a run that turned away part of its offered
  load is indistinguishable from one that served all of it.
  Accept-rate is the ratio
  ``mtp.accepted.total / mtp.proposed.total`` — derive in the
  collector / Grafana, no separate gauge so the value stays
  numerically consistent across export windows. Prefix-cache hit-rate
  is similarly derived from
  ``prefix_cache.tokens_saved_total / tokens.in.total`` over the
  same window.
- Histograms: ``arbi_serve.request.queue_wait_seconds{backend}``,
  ``arbi_serve.request.first_token_seconds{backend}``,
  ``arbi_serve.request.decode_window_seconds{backend}``,
  ``arbi_serve.request.total_seconds{backend}``,
  ``arbi_serve.tokens_per_second{backend}``,
  ``arbi_serve.batch.in_batch_shared_tokens`` (per scheduled batch,
  emitted only when the prefix grouper actually clusters something),
  ``arbi_serve.mtp.verify_slate_rows{k}`` (per MTP verify forward: the
  rows ``B x (K+1)`` the verify GEMM ran at, padded rows included on a
  pad-up replay; even 2-row buckets through 64 so each captured rung
  at one draft depth ``k`` is its own bucket — the served distribution
  the int8 verify ladder and ``exl3_int8_verify_min_rows`` are decisions
  about),
  ``arbi_serve.residency.switch_duration_ms{from_model,to_model}``
  (park+wake wall time of a stable-VA model switch).
- Per-label first-call cost counter:
  ``arbi_serve.first_call.ms_total{backend,label}``.
- Up-down counter: ``arbi_serve.requests.in_flight``.
- Observable gauges (callbacks polled at export interval):
  ``arbi_serve.kv_pool.pages_used{kind,backend}``,
  ``arbi_serve.kv_pool.pages_total{kind,backend}``,
  ``arbi_serve.residency.active_model{served_name}`` (info-style: 1 for
  the active resident, 0 for each parked pool member; in the default
  single-model mode exactly one series with value 1),
  ``arbi_serve.residency.parked_models`` (count of host-parked members),
  ``arbi_serve.residency.parked_host_bytes`` (their pinned host-RAM
  footprint).

Counter / histogram / up-down-counter operations on the SDK are
thread-safe; call sites do not need an :class:`asyncio.Lock`.
Observable callbacks run on the SDK's exporter thread and only read
the engine's pool state (plain int reads on
``free_pages`` / ``num_pages``).

This module is split into submodules by concern: ``_meter`` (meter /
tracer resolution), ``_buckets`` (histogram boundaries), ``_instruments``
/ ``_bind`` / ``_collectors`` (the :class:`Metrics` mixins), and
``_metrics`` (the composed class). The OTEL provider configuration and
its process-global state live here so callers and tests address them at
``arbi_serve.server.metrics``.

``PeriodicExportingMetricReader`` that opens ONE collection scope.

The scope is what makes an observable gauge's expensive source read
once per collection instead of once per instrument that projects it
(:mod:`arbi_serve.server.metrics._collection_scope`). It is opened
here, on the reader, because the reader is the only object that knows
where one collection begins and ends — a callback cannot know whether
it is the first of its cycle or the third.

Build a ``PrometheusMetricReader`` that opens a collection scope.

Subclassed lazily (inside the function) because the base class is an
optional import: ``opentelemetry-exporter-prometheus`` may be absent,
and a module-level ``class X(PrometheusMetricReader)`` would then fail
at import time rather than at the mount, where the caller already
raises a message that says what to install.

OTEL ``Resource`` describing this process to the collector.

The base resource carries ``service.name`` + ``service.version``;
the optional ``extra`` dict layers additional attributes on top
(model identity, dtype, head count, etc.) so every emitted metric
/ log / span gets them as labels via the collector's
``resource_to_telemetry_conversion``. ``OTEL_RESOURCE_ATTRIBUTES``
env-set values still merge on top via the OTEL SDK default.

Pick the OTLP exporter based on env.

- ``OTEL_EXPORTER_OTLP_PROTOCOL`` ∈ {``http/protobuf``, ``grpc``};
  default ``http/protobuf``.
- ``OTEL_EXPORTER_OTLP_ENDPOINT`` overrides the default
  (``http://localhost:4318`` for HTTP, ``http://localhost:4317``
  for gRPC).
- ``OTEL_METRICS_EXPORTER=console`` short-circuits to a stdout
  exporter (test / smoke).

Pick the OTLP log exporter mirroring metrics / traces selection.

Honours ``OTEL_LOGS_EXPORTER=console`` for short-circuiting to
stdout in tests. Otherwise picks gRPC vs HTTP from the same
``OTEL_EXPORTER_OTLP_PROTOCOL`` env that the other signals use.

The address the OTLP exporters will dial, as the SDK will resolve it.

Signal-specific endpoints win over the shared one (the SDK's precedence);
with none set the SDK falls back to localhost on the protocol's port.

One TCP connect to the collector, ONCE per process, at boot.

This is what makes telemetry-on-by-default safe to ship: with the
exporter attached and nothing listening, the SDK's periodic export
turns into ``Connection refused`` on stdout every few seconds for the
life of the process. Deciding once, here, converts that into a single
logged line and no exporter.

A boot-time decision on purpose. The alternative — noticing per export
that the collector is gone — is a health subsystem, and the failure it
would catch (a collector that dies mid-run) costs a dropped batch, not
a served request.

Should we attach an OTLP push exporter?

On by default (``ARBI_SERVE_ENABLE_OTEL``), and safe to be, because the
answer is not taken from the flag alone: a single boot probe confirms
something is actually listening at the endpoint the exporters would
dial. Without that confirmation the periodic exporter is never
installed, so a box with no collector pays nothing and says nothing
beyond one line naming the address it tried.

``ARBI_SERVE_ENABLE_OTEL=0`` switches the channel off outright, without
probing.

Install OTEL ``TracerProvider`` + ``MeterProvider`` globally.

Idempotent: subsequent calls are no-ops so import-time + lifespan
startup doesn't double-install. Tests that need a clean slate can
set ``arbi_serve.server.metrics._otel_configured = False``.

OTLP push is **on by default** and gated on a live collector (see
:func:`_otlp_push_configured`): the periodic exporter is installed only
when a boot probe finds something listening at the endpoint it would
dial, so a box without a collector stays silent on the OTEL channel and
pays nothing for it. The Prometheus pull-side reader mounts
independently under ``ARBI_SERVE_ENABLE_PROMETHEUS=1`` (default), so
``/metrics`` keeps working with or without a collector.

The ``OTEL_METRICS_EXPORTER=console`` / ``OTEL_TRACES_EXPORTER=console``
short-circuits in ``_build_*_exporter`` are honoured: setting either
counts as "user opted in" and a console exporter is wired up.

Build a retroactive per-request span tree and end it immediately.

Parent span ``request`` spans ``submit_time → t_complete``; child
spans ``queued`` (submit→first_scheduled_step), ``prefill``
(first_scheduled_step→first_token), and ``decode``
(first_token→complete) carve up the lifecycle. A prefix-cache hit is
reported as ``cached_tokens`` / ``prompt_tokens`` ATTRIBUTES on the
``prefill`` span, never as its own child span — the cached tokens are
skipped at admission, so they occupy no part of the prefill window and
any span drawn there would be a duration nothing measured. Every
boundary is guarded against the ``0.0`` /
``None`` "unstamped" sentinel — a request that errored before it
was scheduled emits only the spans whose boundaries exist (e.g.
just ``queued``).

``cached_tokens`` is always stamped, so a reader can distinguish a
measured 0% hit from an unmeasured one. ``reasoning_tokens`` is
stamped whenever the request was in reasoning mode — 0 included — and
omitted only when it could not have reasoned; the same absent-vs-zero
rule the API's ``usage.completion_tokens_details`` follows.

``t_prefill_start`` is ``RequestTiming.t_first_scheduled_step`` —
the moment the scheduler picked the request up for its first
forward (the queued→prefill boundary). All times come in as
Unix-seconds floats (0.0 / None = unstamped). Runs ENTIRELY on the
off-thread emitter; never on the engine loop.

Single-thread queue-fed OTEL event drain.

Public surface: :meth:`submit` enqueues a ``(callable, args,
kwargs)`` tuple; the bg thread pops + runs them. :meth:`shutdown`
drains the queue with a bounded wait then joins. Any exception
raised inside a submitted callable is logged + swallowed so one
bad emit doesn't kill the drain.

Lifecycle: built once in :func:`configure_otel` and stored at
module level (:data:`_event_emitter`). Tests that need a fresh
state can call :func:`reset_event_emitter`.

Return the module-level emitter, lazily starting it in the serving path.

Engine call sites use :meth:`OtelEventEmitter.submit` to push the
per-finish metric burst off the engine step coroutine. The burst is
NOT free even without OTEL push configured — the Prometheus
counters / histograms are live, so ~10 inline records per finish is
real work on the single engine loop that starves the GPU at high
request rates (bulk embedding / rerank).

So auto-start the off-thread emitter the first time it's requested
from inside a running event loop (i.e. the server). The sync test
path (no running loop) keeps ``None`` → inline emission, so metric
assertions stay deterministic.

Estimate VRAM fragmentation in bytes from cheap counters.

The base term is the caching allocator's reserved-but-unused span
(``torch.cuda.memory_reserved() - torch.cuda.memory_allocated()``):
bytes the allocator holds from the driver that no live tensor
occupies. When ``eng`` carries a cuMem sleep pool, the gap between
its mapped physical bytes and the bytes the caching allocator
accounts for is added on top (cuMem-mapped tensors live outside the
caching allocator, so the driver holds them but ``memory_reserved``
does not see them).

All reads are O(1) counter lookups — no ``torch.cuda.synchronize``,
no allocation. Returns ``None`` when CUDA is unavailable (CPU runs)
so the caller emits nothing rather than a fabricated zero, and when an
allocator query would be unsafe.

O(1) is not the same as safe. ``memory_reserved`` / ``memory_allocated``
take the caching allocator's mutex WITHOUT releasing the GIL, and this
runs on the event-loop thread (the serving heartbeat's tick) while an
engine-thread allocation into a cuMem-backed pool holds that same mutex
and wants the GIL. Opposite order, hard deadlock — so the read goes
through :func:`~arbi_serve.runtime.named_pool.allocator_query_lease`,
which grants it only when no allocation window is open and holds that
state for the duration.

Observable-gauge registration against a live engine OR a StatsMsg cache.

Mixin half of :class:`arbi_serve.server.metrics.Metrics`: wires the
pool-occupancy / memory-accounting observable callbacks once the engine
is fully constructed (:meth:`MetricsBindMixin.bind_engine`, in-proc
modes), or the StatsMsg-sourced subset against the API child's
``EngineClient`` cache (:meth:`MetricsBindMixin.bind_stats_client`,
process mode; see docs/engine_core_process.md §6).

Process-mode gauge availability (the explicit contract)
-------------------------------------------------------
The API child registers ONLY the gauges whose data rides the cached
``StatsMsg`` (never engine memory, never an RPC):

  PRESENT in the API child (StatsMsg-sourced, series-identical):
    - ``arbi_serve.kv_pool.pages_used``   (kv_total_pages − kv_free_pages)
    - ``arbi_serve.kv_pool.pages_total``  (kv_total_pages, mapped-aware)
    - ``arbi_serve.engine.loop_tick_age_seconds`` (last_loop_tick)
    - ``arbi_serve.gpu.sleep_state``      (memory_released + mode label)
    - ``arbi_serve.gpu.sleep_mode``       (sleep_pool_mode)
    plus the per-process NVML/psutil gauges ``_install_telemetry``
    registers (they read hardware / the local process, not engine
    memory).

  ABSENT in the API child (data is NOT in StatsMsg — weight
  attribution, activation_profile, named-pool snapshots, residency
  registry, CUDA allocator walks, boot phases, ...): every other
  ``_OBSERVABLE_GAUGES`` entry. These stay LOCAL to the engine process
  — they exist there iff the engine process configures its own OTEL
  exporter; with no engine-side exporter they are simply absent in
  process mode. This is deliberate: shipping tensor/allocator walks
  through the status snapshot would bloat the >=1 Hz heartbeat for
  dashboard-only data.

Wire pool-occupancy observable callbacks to a live engine.

Call once the engine is fully constructed (pool may still be
None if called pre-build — the callback handles that).
Repeated calls re-bind to the most-recent engine, useful in
tests that build multiple engines in one process.

Wire the StatsMsg-sourced gauge subset to an ``EngineClient``.

The process-mode (API child) counterpart of :meth:`bind_engine`:
registers ONLY the ``STATSMSG_GAUGES`` subset (see the module
docstring's present/absent table), with instrument name / unit /
description resolved from the SAME ``_OBSERVABLE_GAUGES`` table
the in-proc path uses — the two sources cannot drift on
instrument identity. Callbacks read the client's cached
``StatsMsg`` only (§6: no RPC on the export path). Idempotent
per Metrics instance, like :meth:`bind_engine`.

One metric collection = one reading of each expensive source.

WHY THIS EXISTS. An OTel observable gauge is one instrument with one
callback, and several of our gauges are three different projections of a
single physical reading — the caching allocator's segment walk gives
``allocated``, ``reserved`` and ``fragmentation`` for every named pool at
once. Registered as three instruments, the SDK invoked the reading three
times per collection, so the process paid three ``torch.cuda.memory.
_snapshot`` walks to answer one question about one instant. MEASURED on
the shipped 27B: a ``/metrics`` scrape spent 52% of its engine-process
time inside those walks.

WHAT IT IS NOT. It is not a cache with a lifetime, and there is no
interval to tune. The scope is opened by the READER, once, immediately
before it drives the callbacks (:func:`begin_collection`), and every
memoised source is recomputed in the next collection. So a value is
never older than the collection that reports it — the reading is exactly
as fresh as it was before, taken once instead of N times. A source that
nothing asks for in a given collection is not read at all.

CONSISTENCY, NOT JUST COST. Three walks described three instants, and
their sums were therefore never required to reconcile: a pool could be
reported allocating bytes that the fragmentation row, taken 40 ms later,
had already seen freed. One walk per collection makes the rows describe
one allocator state, which is the only state they can honestly add up in.

Thread-safe: the pull reader collects on whatever thread serves
``/metrics`` while the periodic push reader collects on its own timer
thread. Each ``begin_collection`` starts a scope private to its caller,
so a scrape and an export never share (or invalidate) each other's
readings.

Open a fresh collection scope on this thread.

Called by the metric readers immediately before they drive the
instrument callbacks. Discards whatever the previous collection on
this thread memoised, so nothing is ever carried across collections.

Close the scope, dropping every memoised reading.

Explicit rather than implicit so a reading cannot outlive the
collection that took it — a snapshot left behind on a thread-local
would be handed to the NEXT collection as if it were current, which
is the exact failure this module exists to make impossible.

Return ``build()``, computed at most once in the current scope.

Outside a scope (no reader opened one — an ad-hoc caller, a test)
``build`` runs every time: memoising with no defined end would be a
cache with an unbounded lifetime, which is a different and worse
thing than the one this function provides.

Observable-gauge callbacks polled by the OTEL exporter thread.

Mixin half of :class:`arbi_serve.server.metrics.Metrics`: every method
reads the bound engine's pool / memory / residency state and returns the
observations for one gauge. All fail closed (empty tuple) so a malformed
observation never raises into the exporter thread.

Two gauge sources exist (see ``docs/engine_core_process.md`` §6):

* **in-proc modes** — :class:`MetricsCollectorsMixin` below, reading the
  live bound engine.
* **process mode (API child)** — :class:`StatsMsgCollectors`, defined in
  the sibling :mod:`._stats_collectors` module and re-exported here so
  it stays importable as ``...metrics._collectors.StatsMsgCollectors``.
  It reads the :class:`~arbi_serve.engine.proc.client.EngineClient`'s
  cached ``StatsMsg`` (never engine memory, never an RPC — §6 "no RPC
  on the probe path"). Only the gauges whose data the ``StatsMsg``
  carries exist in the child; see ``_bind.py::bind_stats_client`` for
  the explicit present/absent table.

Allocatable KV-pool capacity for the watermark/used metrics.

For a GrowableRegion-backed PAGED_KV pool ``num_pages`` is the reserved
VA CEILING; only the physically MAPPED prefix is allocatable (the
free-list spans it). Reporting the ceiling as ``pages_total`` shows the
pool ~full at idle on Grafana (the unmapped tail booked as used) — the
same misleading accounting admission divides against. Report the mapped
page count for a growable pool; the non-growable path keeps
``num_pages`` (gated on ``kv_is_growable`` so MagicMock test doubles do
not auto-satisfy the mapped accessor).

Observe page-locked HOST bytes booked in the host ledger, by owner.

Read from the process-wide ledger
(:mod:`arbi_serve.runtime.pinned_host_budget`), which is the same
source the boot block prints — not from the engine — so the number
exists whether or not anything is watching and whichever subsystem
made the reservation. An owner holding nothing emits no series.

The bound engine's StableVaResidencyController, or ``None``.

Mirrors the ``_maybe_switch_model`` guard: only a real
``StableVaResidencyController`` counts, so a test double (e.g. a
``MagicMock`` engine auto-vivifying ``stable_va``) stays on the
default single-model path.

One observation per pool member: 1 = active resident, 0 = parked.

Default single-model mode (no controller) emits exactly one series
with value 1 for the served model — no parked noise.

Pinned host-RAM bytes held by parked members' sleep-pool backups.

Sums each parked record's snapshotted ``sleep_pool.staging_bytes``
— a plain attribute read per record (the same cheap source the
``gpu.pinned_host_parked_bytes`` gauge reads for the active
model), never a tensor walk.

Observe the serving device's PHYSICAL capacity.

Absent — not zero, and not the addressable total — when NVML cannot be
asked. A consumer must be able to say the card's size is unknown.

Observe seconds since the engine run loop last ticked.

Mirrors the ``/health`` liveness signal as a Grafana series so a
developing hang is visible + alertable before the probe trips an
orchestrator restart. The loop stamps ``_last_loop_tick`` every
iteration (busy / idle / quiesced, <=0.5s cadence), so a low value
is healthy in every legitimate state (boot / reload / sleep
included) and only a sustained climb signals a wedge.

Resolve the phase map to observe.

Prefer this instance's recorded phases; fall back to the bound
engine's (covers the case where OTEL keeps the FIRST registered
callback for a given instrument name across multiple Metrics
instances sharing one MeterProvider — the bound engine is the
single source of truth either way).

One observation per visible CUDA device. Silently empty when
CUDA isn't available so the gauge degrades cleanly on CPU runs.

Reads through :func:`~arbi_serve.server.nvml_metrics.
device_memory_used_free_bytes`, which reports every visible device
WITHOUT making one current: ``torch.cuda.mem_get_info(d)`` does make
``d`` current, so asking it per device leaves this process holding a
primary CUDA context on cards it does not serve from.

Observe per-component CUDA weight bytes for this rank.

One observation per component label (text_layers / lm_head /
embed / vision / mtp / norm / quant_scratch / other). The
underlying walk is the SAME CUDA param+buffer enumeration as
:meth:`_observe_weights_bytes`, bucketed by qualified-name
prefix, so Σ(component) == weights_bytes for this rank. Every
component is emitted even when zero so the label set is stable
across models.

Observe PHYSICALLY-mapped KV-slab bytes (per rank/backend).

``kv_pool_bytes`` reports ``view.bytes_total()`` — the full VA span
(the slab tensors span ``num_pages`` whether or not every page is
backed). For an Option-B growable slab, only ``kv_mapped_pages`` of
``num_pages`` are physically mapped during capture; this gauge scales
the span by the mapped fraction (every page is equal bytes) so the
Phase-1 prefix → post-capture grown size is visible and reconciles
against NVML used. Non-growable slabs map all pages ⇒ equals the span.

Evaluate the freeze's identity against the device now.

``None`` before the armed freeze — ``post_freeze_floor_bytes`` is set
only by a freeze that capped the allocator, while an unfrozen boot
raises ``_serving_frozen`` without one — and whenever the reading
cannot be taken safely: the cuMem term takes the allocator's lock, so
it is read under the allocator query lease, and a tick that lands
inside a pool ``use()`` window is a gap in the series rather than a
wedged engine.

Advance the violation episode and stamp ``inv`` with its start.

The first violating tick opens an episode (logged once, at ERROR, with
the figures); every later violating tick keeps its timestamp; the
first healthy tick closes it. Idempotent within a tick, so the gauges
that share one collection cycle may each call it.

Observe 1 per KNOWN active degradation (component + reason labels).

Reads the engine's ``_serving_degraded`` registry (``component ->
reason``): ``capture`` when ARBI_ALLOW_DEGRADED_CAPTURE downgraded
failed capture buckets, ``drafter`` while the MTP drafter chain is
cold-pathing. Empty registry ⇒ no series (full-speed serving).

The live ring's own accounting, or ``None`` when no ring is attached.

The ring keeps the numbers; this only reads them. A pure-attention
model, a pool built before the ring existed, and the host-arena route
all legitimately have no ring, and each reads as absent rather than as
a zero that would look like a ring that never wrapped.

Observe PHYSICALLY BACKED bytes of the DeepSeek-V4 sparse state pool.

Not the reservation: the stores reserve their whole virtual span at
boot and back rows as live requests reach the contexts that need
them, so a gauge reporting the span would read flat at the maximum
and say nothing about occupancy.

Sum cos/sin table bytes attributed to ``scratch.rope``.

The cos/sin tensors are the only bytes under that tag, so the tag's
total IS the rope-cache footprint. Read through the cuMem tag ledger
rather than the pool's block view so the number survives the
persistent-scratch fold, after which the tables are a sub-tag carve of
a slab in another pool and the ``scratch.rope`` pool holds nothing.

Observe caching-allocator reserved bytes per named pool.

Also feeds ``scratch.forward_arena``'s reading to the serving watch
(:mod:`arbi_serve.engine.arena_watch`) — the reading is already being
taken for the gauge, so passing it on costs nothing. It is a second
trigger only: the watch is OWNED by the engine's post-step seam,
because a floor that only recorded on scrape recorded nothing at all
in a deployment nobody scrapes.

Hand the reserved-bytes gauge's arena reading to the serving watch.

A SECOND TRIGGER, NOT AN OWNER. The watch lives in
:mod:`arbi_serve.engine.arena_watch` and is driven by the engine's
post-step seam, because the reading exists to record that SERVING grew
the pool — and an export callback fires only when something scrapes
(and, under the no-op meter provider, never at all), so a floor that
depended on it was a floor that depended on being watched.

This path is still worth keeping: the reading is already in hand for
the gauge, it costs nothing, and it covers a process whose pool grew
while the engine loop was not the thread that noticed. Both triggers
call the same function and share the engine's high-water, so there is
one recording, not two copies of one.

Hand a driver-residency reading to the serving watch.

A SECOND TRIGGER, NOT AN OWNER — the same demotion its arena sibling
above carries, and for the same reason. The watch lives in
:mod:`arbi_serve.engine.arena_watch` and is driven by the engine's
post-step seam, because the reading exists to record that SERVING added
driver-resident physical. An export callback fires only when something
scrapes, and under the no-op meter provider never at all, so a floor
term that depended on this path was a term that read 0 on every
deployment nobody watched — which is the exact defect
``transient.serving_step.driver_growth`` was added to end, and it would
have re-introduced it one layer down.

Kept because it costs nothing and covers a process whose driver grew
while the engine loop was not the thread that noticed. Both triggers
call the same function and share the engine's high-water, so there is
one recording, not two copies of one.

Emit one observation per named pool, reading ``snapshot_key`` from each pool's snapshot.

ONE registry read per collection, shared by the three gauges that
project it (allocated / reserved / fragmentation). They are three
views of a single ``torch.cuda.memory._snapshot`` walk, so taking
the walk per gauge cost three times over for one instant's worth
of information — and produced three instants that were not
required to reconcile. See
:mod:`arbi_serve.server.metrics._collection_scope`.

Uses :meth:`NamedPoolRegistry.snapshot_with_residual` (NOT
``snapshot_all``) so ``memory_reserved(device) - Σ named-pool reserved``
is emitted as first-class series alongside the named pools and the
per-pool breakdown reconciles to ``torch.cuda.memory_reserved`` instead
of silently dropping the remainder.

That remainder arrives as THREE labels, not one, because it is three
different things and one name for them cannot be true of all three:
``unpooled.torch_default_pool`` (live tensors that escaped every
``NamedMemPool`` context — unbudgeted and invisible to the budget
enforcer, but real resident memory), ``address_space.released_pool_va``
(address space of a pool this boot destroyed, whose physical the driver
already took back) and ``unpooled.unregistered_pool`` (an unregistered private
``MemPool`` — the routing-bug signal, ~0 on an honest boot; NOT the boot
ledger's ``driver.residual``, which is the whole-card remainder and not a
leak signal). Emitting them under one label is what made the gauge
disagree with the boot VRAM ledger. ``snapshot_with_residual`` degrades to the named-only list
when CUDA is unavailable, so CPU metric fixtures are unaffected.

Emit a WARNING log when the arena's high-water fraction
stays above 0.95 for 5+ minutes.

Tracks the first observation that crossed 0.95 in
``self._arena_warn_armed_at``; resets when frac drops back
below the threshold. The warning fires once per crossing —
a sustained-above-0.95 condition does not spam the log.

Backend label for aggregate ``counter_deltas`` (§5.2).

``counter_deltas`` are summed across backends with no per-backend
breakdown, but the API counters carry a ``{"backend": <spec>}``
attribute. When the engine is serving a single spec that spec is the
unambiguous label; a mixed-spec engine gets ``"aggregate"`` (the
bench regex matches regardless of labels — this only keeps the label
honest, never fabricating a spec the delta cannot be attributed to).

Every canonical startup-phase slug.

The alias table maps both the human phase names and the slugs onto the
slug, so its VALUES are the canonical set. Exposed so a caller that names a
phase (a boot cache declaring which phase it gates) can be checked against
the same source the metrics use, instead of a second copy that drifts.

Consumer half of the engine→API counter bridge (§5.2).

In process mode the engine step loop records its MTP / generated-
token counters on the ENGINE process's registry, which has no
``/metrics`` HTTP surface — so the API child's scraped counters
would be missing every engine-step count. The heartbeat ships the
per-tick DELTAS on every ``StatsMsg`` (``counter_deltas``); the API
child re-emits each straight through into THIS (API-process)
counter bundle so ``/metrics`` reflects engine-step throughput.

Delta key → counter:

  - ``mtp_accepted``     → ``mtp_accepted_total``  (arbi_serve_mtp_accepted_total)
  - ``mtp_drafted``      → ``mtp_proposed_total``  (arbi_serve_mtp_proposed_total)
  - ``mtp_spec_disabled_steps`` → ``mtp_spec_disabled_steps_total``
    (arbi_serve_mtp_spec_disabled_steps_total)
  - ``generated_tokens`` → ``tokens_out_total``    (arbi_serve_tokens_out_total)
  - ``mtp_draft_rows``   → INTENTIONALLY NOT bridged: no public
    counter mirrors it (it only feeds the heartbeat line's
    tokens/decode-step figure), and minting a series no dashboard
    reads would be gold-plating.

Straight-through by contract: the producer drains its accumulator
exactly once per ``StatsMsg`` (``take_counter_deltas``), so each
delta is ``add``-ed here exactly once — never cached or re-
accumulated. Non-positive deltas are skipped (mirrors the
producer's ``_accumulate`` guard). ``counter_deltas`` are aggregate
across backends with no backend info, so they carry the sole
active spec as the ``backend`` label when unambiguous, else
``"aggregate"``. Never raises — an observability path must not
perturb the HTTP loop.

Record one cold-boot phase's wall-clock duration.

``name`` is matched (case-insensitively) against
:attr:`_STARTUP_PHASE_ALIASES`, first whole and then with a
trailing parenthetical qualifier stripped; unknown phases are
slugified and kept as-is so nothing is silently dropped. Repeated
phases (e.g. ``compile_warmup`` charged twice on the hybrid path,
``decode_autotune`` charged on both sides of the resize)
accumulate.

Record one backgrounded boot task's overlap.

``duration`` is submit→worker-return; ``blocked`` is the main
thread's wall inside the join. ``overlapped`` is the difference —
what the overlap bought. A task whose ``overlapped`` is ~0 ran
serially behind a thread, and the phase that joins it is a wait.

Phase totals as they stand now, to measure a later build against.

Phases ACCUMULATE (see :meth:`record_startup_phase`), so after the
first build the dict describes every build the process has run. A
caller timing ONE build takes a snapshot when that build starts and
hands it back as ``since``.

Return ``(phase_sum, unattributed)`` for a build of ``total_seconds``.

The phases telescope: each closes the window that opened when the
previous one closed, so their sum IS the build minus whatever runs
after the last phase closes. ``unattributed`` is that remainder and is
the number that must stay near zero — a window no phase names is a cost
nothing can be held to.

``since`` scopes the sum to one build. Without it a SECOND build (a
live backend swap prepares a member through the same code) sums every
build the process has run against only its own elapsed time, which
printed sums larger than the total and a NEGATIVE remainder — the
invariant this function exists to check, reported as violated by the
measurement rather than by the boot.

Record one byte value in the cold-boot capture→grow VRAM lifecycle.

``stage`` is one of :attr:`_BOOT_MEM_STAGES`. Last write wins per stage.
Recorded during build(); replayed every export interval (see
:meth:`_observe_boot_mem`). Negative / None values are dropped LOUDLY
via a warning rather than silently — a bad measurement is a real bug.

Per-request prompt + completion length distributions.

Fires on every finished request (not gated on the per-step timing
being enabled). The underlying instruments are histograms; the
Prometheus exporter emits ``_bucket`` / ``_count`` / ``_sum`` so
consumers can compute population averages cheaply.

Inter-token latency = decode_window / max(1, n - 1).

Fires on every finished request with at least 2 output tokens
(single-token outputs have no inter-token gap). Independent of
per-step timing configuration.

Emit per-request OTEL events for the request-timing block.

Called from :meth:`Engine._on_finished` exactly once per request
whose ``timing`` is non-None. Records raw component-duration
histograms, prompt-/completion-length distributions, the
derived inter-token latency, and per-label first-call cost
counter increments. Frontends compose any "TTFT minus
first-call cost" or "steady-state TPOT" view from these
instruments themselves.

Per-Engine bundle of OTEL instruments.

Engine builds one of these in ``__init__`` and increments / records
via ``engine.metrics.requests_total.add(...)`` etc. Each instrument
is created against the *currently-installed* MeterProvider; if
OTEL hasn't been configured (test path), the SDK silently returns
no-op handles.

Observable gauges (KV pool occupancy) need a live engine reference
to query — :meth:`bind_engine` wires the callbacks once the
engine is fully constructed.

StatsMsg-sourced observable-gauge callbacks (process-mode API child).

These names stay importable via re-export from
:mod:`arbi_serve.server.metrics._collectors`. Two gauge sources exist
(see ``docs/engine_core_process.md`` §6):

* **in-proc modes** — ``MetricsCollectorsMixin`` (in ``_collectors``),
  reading the live bound engine.
* **process mode (API child)** — :class:`StatsMsgCollectors` below,
  reading the :class:`~arbi_serve.engine.proc.client.EngineClient`'s
  cached ``StatsMsg`` (never engine memory, never an RPC — §6 "no RPC
  on the probe path"). Only the gauges whose data the ``StatsMsg``
  carries exist in the child; see ``_bind.py::bind_stats_client`` for
  the explicit present/absent table.

StatsMsg-sourced observable-gauge callbacks (process-mode API child).

Holds the ``EngineClient`` (duck-typed: anything with a
``latest_stats`` attribute yielding a
:class:`~arbi_serve.engine.proc.messages.StatsMsg` or ``None``) and
reproduces the SAME instrument series the in-proc callbacks emit —
identical names, values, and label sets — from the cached snapshot.
Before the first snapshot every callback fails closed (empty tuple),
matching the in-proc "no engine bound" convention.

``engine.loop_tick_age_seconds`` from ``last_loop_tick``.

``last_loop_tick`` is CLOCK_MONOTONIC and the API child is
always same-host (§2), so the subtraction is valid across the
process boundary. 0.0 = loop never ticked → no observation
(mirrors the in-proc "attribute is None" case).

``gpu.sleep_state`` from ``memory_released`` + the mode label.

An empty ``sleep_pool_mode`` means the engine has no sleep pool
(mid-build / non-cuMem boot) → no series, mirroring the in-proc
"engine without a sleep_pool" convention.

Request-ID propagation middleware.

Accept ``X-Request-ID`` from the client (generate a UUID4 if absent),
propagate via a :class:`contextvars.ContextVar` so async coroutines
spawned by the handler observe the same id, and round-trip the id on
the response. JSON log records read the same contextvar.

First value of an ASGI request header (lowercased ``bytes`` key).

Shared by the pure-ASGI middlewares (this module + concurrency) so
they can read request headers off the scope without materialising a
``starlette.Request`` (which BaseHTTPMiddleware forced).

Attach ``x-arbi-degraded`` naming any active serving degradation.

Opt-in via ``ARBI_DEGRADED_HEADER=1`` (the app only mounts this
middleware when the flag is set — default OFF, zero wire-format
change). Reads the engine's ``_serving_degraded`` registry
(``component -> reason``, e.g. ``capture`` under
ARBI_ALLOW_DEGRADED_CAPTURE, ``drafter`` during a drafter cold-path
window) per response and emits the sorted component names as a
comma-separated header value. No active degradation ⇒ no header.

Pure ASGI for the same SSE/GIL reasons as
:class:`RequestIDMiddleware`; the per-response cost is one attribute
read of a (usually empty) dict.

Bind X-Request-ID to a contextvar for the duration of the request.

A **pure ASGI** middleware, deliberately NOT ``BaseHTTPMiddleware``:
the latter wraps every response — including SSE streams — in an anyio
memory-object-stream + task group, so each streamed chunk crosses
extra GIL-held plumbing. Stacked across the middleware chain at
concurrency, that starves the engine thread of the GIL. Pure ASGI
hands ``send`` straight through, so streamed chunks never touch that
machinery.

Hardware-truth VRAM telemetry via ``pynvml``.

Why this exists alongside ``arbi_serve.server.metrics``:

The existing ``arbi_serve.gpu.memory_{used,free}_bytes`` gauges in
:mod:`arbi_serve.server.metrics` read ``torch.cuda.mem_get_info()`` —
which reflects PyTorch's caching-allocator view of the world. That view
can disagree with what NVIDIA's driver actually has booked because:

* NCCL allocates raw ``cuMalloc`` buffers outside PyTorch's allocator.
* Other in-process libraries (cuBLAS workspace, cuDNN scratch, third-
  party CUDA modules) also bypass the caching allocator.
* PyTorch's *reserved* != PyTorch's *allocated* != hardware *used*.

``pynvml.nvmlDeviceGetMemoryInfo()`` is the hardware-truth reading —
the same number ``nvidia-smi`` shows. The two disagree by hundreds of
MB on a steady-state worker; the difference is exactly what gets
attributed to "external-to-PyTorch" allocations on the dashboard.

Module shape: a small init function called once by the engine /
lifespan startup, plus observable-gauge callbacks registered against
the global OTEL meter. Boot does not break if the host has no NVIDIA
driver, no GPU, or pynvml fails to import — the registration is a
no-op in that case.

Register a callable returning physical cuMem-mapped bytes.

Wired from the engine-bind path with ``lambda: eng.sleep_pool.mapped_bytes``.
Used by :func:`_observe_true_driver_floor` to carve cuMem-mapped model
memory back out of the driver-overhead number so the true driver floor
stays stable awake AND asleep. Pass ``None`` to clear (test isolation).

Resolve ``nvmlDeviceGetMemoryInfo_v2`` from libnvidia-ml.so.1.

The Python wrapper (``nvidia-ml-py``) ships ``nvmlMemory_v2``
(the version constant) and ``c_nvmlMemory_v2_t`` (the struct
layout) but, in some releases, omits the function binding. The
underlying NVML library always exports the symbol on R450+
drivers, so we bind it directly via ctypes and feed it the struct
class the wrapper already provides.

Sets module-level ``_v2_ctypes_func`` / ``_v2_struct_cls`` /
``_v2_struct_version`` on success and returns True; returns False
on any failure (no library, no symbol, no struct).

Pick a stable ``rank`` label value.

The OTEL convention used by every other arbi-serve gauge (see
:mod:`arbi_serve.server.metrics`) is to read ``RANK`` from the
environment — set by torchrun on multi-rank launches and absent
(defaulting to ``"0"``) on single-process runs. Same default the
distributed launcher uses internally.

Initialize NVML and probe device handles. Idempotent.

Returns ``True`` if NVML is now live (handles cached, gauges may
register), ``False`` if init failed (no driver, no GPU, pynvml not
importable). On ``False`` the caller should skip gauge registration;
the engine will keep booting fine.

``(addressable total, addressable used, physical total)`` for THIS
process's serving device — ``None`` when there is no CUDA device.

ONE READING, so the three figures can be quoted together. The card has two
totals and they are hundreds of MiB apart:

* ``torch.cuda.mem_get_info`` reports what a process can ADDRESS, and it is
  the total the memory ledger closes against and every engine-side
  percentage is a fraction of. Both of the first two figures come from that
  single call, so a used/total ratio built from them cannot mix sources.
* NVML reports the card's PHYSICAL capacity — the ``nvidia-smi`` figure.
  Larger, because the driver keeps a carve-out of it that no allocation can
  reach. It is here to be NAMED beside the other total, never to be divided
  into. ``None`` when NVML cannot be asked; unknown is what a caller says.

Note that NVML's own ``reserved`` field does NOT reconcile the two (measured
~29 MiB apart on an Ada card), which is why the carve-out is the difference
between the two totals here rather than a third number from a third source.

Current device only. Reading another device's ``mem_get_info`` makes it
current and instantiates a primary CUDA context there — hundreds of MiB of
resident VRAM on a card this process does not serve from. NVML is asked for
the same index, so at TP>1 each rank publishes its own card exactly once and
the three series keep the same multiplicity as each other, which is what
makes summing them across ranks a fleet total rather than a ratio bug.

``{device index: physical capacity}`` — the figure ``nvidia-smi`` prints.

NOT the total the engine allocates against; that one is
``torch.cuda.mem_get_info``'s, and it is smaller. The whole reason this
helper exists is so a caller can name the difference between the two
instead of quietly swapping one for the other, so it returns one number and
leaves the pairing to the caller that has both halves.

Deliberately not NVML's ``reserved`` field beside it: that does not
reconcile the two totals (483 MiB against a 454 MiB gap on an Ada card), and
a caller offered it would read it as the carve-out.

Costs no CUDA context — NVML is a driver query, so this is safe to ask about
a device this process does not serve from. Empty when NVML is unavailable:
the card's capacity is then unknown, and unknown is what a caller must say.

Normalize an NVML/torch device UUID to a bare lowercase hex string.

NVML returns ``"GPU-<uuid>"`` (str or bytes); torch's
``get_device_properties(d).uuid`` is a ``uuid.UUID``. Strip the
``GPU-`` prefix and dashes, lowercase, so the two compare equal.

Map device UUID → bytes THIS process has resident, per NVML.

Sums ``usedGpuMemory`` across compute + graphics running-process
records whose ``pid`` is ours. Returns only devices where our PID was
found with a usable byte count.

Device-wide ``(used, free)`` bytes per visible GPU, keyed by device index.

THE CONTRACT: reading another device's memory must not make this process a
tenant of it. ``torch.cuda.mem_get_info(d)`` reports the CURRENT device, so
the binding makes ``d`` current, and that instantiates ``d``'s primary CUDA
context — hundreds of MiB of permanently-resident VRAM. A telemetry loop
over it therefore takes up residence on every card the container can see; at
TP>1 that is one context per rank on every PEER's GPU, memory the owning
rank cannot use and books as ``driver.foreign_process``.

NVML answers the same device-wide question with no context at all
(``nvmlDeviceGetMemoryInfo`` is a driver query, not a CUDA call), so it is
the source here whenever NVML initialised. Without NVML the reading is taken
for the CURRENT device ONLY: a device this process already owns costs
nothing to ask about, and a missing gauge row is cheaper than a resident
context.

This is a different NVML surface from the one the per-process attribution
path works around: ``nvmlDeviceGetComputeRunningProcesses`` reports HOST pids
inside a PID namespace and is unusable there, while device-wide totals are
unaffected by the namespace.

Device-wide ``mem_get_info.used - memory_reserved`` per torch device.

Used when per-process NVML attribution is unavailable — most
importantly inside a PID-namespaced container, where
``nvmlDeviceGetComputeRunningProcesses`` reports HOST pids that don't
match ``os.getpid()``. The number is DEVICE-wide (it folds in
any other process on the card), which equals this process's overhead
on the common dedicated-GPU inference pod; the ``scope="device"``
attribute flags that it is an upper bound on shared GPUs.

The device-wide half comes from :func:`device_memory_used_free_bytes`, which
does not make a peer device current; ``memory_reserved`` is this process's
own allocator state and needs no context either, so a device this rank does
not own contributes its whole used bytes — which is what ``scope="device"``
already declares.

Yield True only while querying the CUDA caching allocator is SAFE.

``torch.cuda.memory_reserved`` / ``memory_stats`` take the allocator's
mutex without releasing the GIL. These callbacks run holding the GIL — on
the OTel exporter thread, or on the HTTP thread answering a ``/metrics``
scrape. An engine-thread allocation into a cuMem-backed pool takes that
same mutex and then WANTS the GIL (the pluggable-allocator hook calls back
into Python). Opposite acquisition order, so the two deadlock outright,
not merely contend. Observed at spec-decode drafter attach, and again at
the cuBLAS workspace warmup with a scraper running: every thread in
``futex_wait``, 0% CPU, idle GPU, the boot wedged until it was killed.

Delegates to :func:`arbi_serve.runtime.named_pool.allocator_query_lease`,
which makes the safety check and the query ONE step — a bare predicate
read is check-then-act, and the gap between check and query is where the
measured hang landed. A gap in a memory graph is a cosmetic loss; the
alternative is a dead server. NVML-sourced numbers are unaffected — they
never touch the torch allocator.

Observe non-allocator CUDA memory: NVML framebuffer - memory_reserved.

Prefers PER-PROCESS attribution (``scope="process"``) — our NVML
framebuffer minus ``torch.cuda.memory_reserved`` for each device we're
resident on, matched by UUID — which isolates this process on a shared
GPU. When NVML can't attribute our pid (PID-namespaced containers; MIG;
insufficient perms) it falls back to a DEVICE-wide
``mem_get_info``-based number (``scope="device"``). Never raises.

Observe the irreducible driver floor with cuMem memory carved out.

``true_driver_floor = max(0, driver_overhead - cumem_mapped)`` where
``driver_overhead`` is the same per-UUID ``framebuffer - memory_reserved``
computed by :func:`_observe_driver_overhead`. Subtracting the cuMem-mapped
model/KV bytes (which live outside the torch allocator and therefore fall
into driver_overhead) leaves just the CUDA-context + cubins floor,
which is stable whether the engine is awake or asleep.

Falls back to a DEVICE-wide number (``scope="device"``) when per-process
NVML attribution is unavailable, mirroring :func:`_observe_driver_overhead`.
Never raises.

Observe the VRAM-reconciliation residual; ~0 when the model holds.

The invariant
``framebuffer == memory_reserved + cumem_mapped + true_driver_floor``
holds by construction here (true_driver_floor is defined as the slack),
so this series should read ~0 on every scrape. A persistently non-zero
value flags a bug in one of the underlying reads (e.g. cuMem bytes that
aren't actually inside driver_overhead). Per-UUID, ``scope="process"``;
skipped entirely when per-process NVML attribution is unavailable (the
device-wide fallback can't be reconciled against a per-process floor).
Never raises.

Register hardware-truth VRAM gauges against the global meter.

Idempotent on the meter side via the underlying SDK
(``create_observable_gauge`` is fine to call multiple times against
the same meter — duplicates resolve to the same instrument). Caller
is expected to have run :func:`init_nvml` first; this re-runs it
defensively in case of out-of-order startup.

Returns ``True`` if gauges are now wired, ``False`` if NVML init
failed and no gauges were registered.

Frozen sizes for the catalogue's token-targeted prompts.

A workload that names a token count has to be sized against a real
tokenizer, and :func:`~arbi_serve.admin_console.ab_probe.canonical_prompt`
does it by bisecting the target's own ``POST /tokenize``. That search is
correct and it is slow in the one place it is least affordable: a dozen
probes, each carrying the whole candidate body -- ~700 KB at the long tier
-- run on EVERY launch, before the first request reaches the engine.
Measured from the console, pressing launch on the 128k preset left the
swimlane blank for ~6 s while a short preset started instantly.

Nothing in that search depends on the moment it runs. The catalogue's
prompts are FIXED: the corpus is this package's own source, the seed is
the default, the kind and the task line are pinned by the spec. The same
search, run at every launch, returns the same answer every time.

So it is run once and the answer is written here, by
``tools/freeze_playground_prompt_sizes.py``. A launch reads the size
instead of searching for it, and pays no tokenize call at all.

WHAT IS STORED IS A RECIPE, NOT A PROMPT. The prompts are hundreds of
kilobytes each and are perfectly reproducible from the few numbers the
search returned, so those are what is kept: the word budget the bisection
landed on, and the character cut the refinement took off the body. Feeding
them back through the same builder reproduces the prompt byte for byte
(``rebuild`` in ``ab_probe``), which is also what makes this table
checkable -- a test rebuilds from it and compares.

WHEN IT GOES STALE, AND WHY THAT IS SURVIVABLE. A size is a fact about one
tokenizer AND one corpus, and both move. Swap the model for one that
tokenizes differently and these budgets still build a prompt of very
nearly the same length, but the exact count shifts. So does the corpus:
the material is this package's own source, so every commit to it moves the
text a word budget draws. Measured across two nearby trees, the same
recipe came out 127,969 and 128,131 tokens -- 0.13% apart, and neither is
wrong.

That is survivable because the frozen number never has to be the number a
reader believes. Every entry records the model it was measured against,
and a run reports the prefill the ENGINE measured for the request it
actually sent (``prompt_tokens`` on each ``request_done``). The frozen
figure sizes the prompt; the engine's figure is the one any tok/s is
divided by, and it is measured on every run precisely because this table
is not. Re-freeze with the tool when the served model changes, or whenever
the tier names should be tight again.

A miss -- a depth, seed or task line with no entry, which is what a CUSTOM
prefill depth typed into the UI is -- falls back to the live bisection.
An operator who asks for an arbitrary number gets it measured; the
catalogue does not pay for that.

The table key for one prompt's inputs.

Every input the builder reads is in it. A key that left one out would
match an entry measured for a different prompt, and a prompt of the
wrong size is worse than a slow launch: it silently renames the
measurement.

Canonical admin-UI playground workloads.

Not invented from scratch: these mirror the traffic shapes treated as
canonical elsewhere in the repo —

  * ``matrix_a_*`` / ``matrix_b_mtp_*`` — the concurrency-swept
    TTFT/TPOT/tok-s cells rendered as "Matrix A" (cross-config
    throughput/latency) and "Matrix B" (the MTP-aware variant);
    applied here to ONE live config instead of a cross-engine grid.
  * ``staggered_ramp`` / ``mixed_coadmission`` — the two realistic
    traffic shapes from ``tools/flag_lane_probe.py``'s ``_drive_smoke``
    (private there; generalized here into something importable).

Distinct-prompt-per-request and usage.completion_tokens-based token
rate are both hard-won bench methodology (identical prompts
let the radix cache / prefix-grouper share prefill across "concurrent"
requests, producing a phantom-low TTFT that isn't real concurrent
load). This module is UI-responsiveness-scoped, not a full bench
harness: one warm-up discard + one measured pass,
not full warm-to-stability convergence — good enough to see live in a
browser, not a substitute for a research-grade sweep harness when a
real publishable number is needed.

A digest of the SOURCE the corpus would be built from.

Content, not mtimes: a checkout or a bind-mount rewrites timestamps on
files whose bytes never changed, and a corpus keyed on those would be
rebuilt for nothing -- while an edited file that kept its timestamp
would serve a stale corpus, which is the failure that matters. Reading
all of it costs ~34 ms against the ~1.9 s it decides whether to spend.

Best effort. A corpus that could not be written is rebuilt next time,
which is slow and correct; failing the launch would be neither.

Written through a temporary file in the same directory and renamed, so a
second console starting at the same moment reads either the whole
corpus or nothing -- never half of one.

Real text for ``corpus``, as LINES. Built once per process, and for
``prose`` kept on disk so it is built once per SOURCE TREE.

THE PROSE CORPUS IS THE EXPENSIVE ONE, and it is expensive in the place
that hurts: an operator's first long-tier launch on a fresh console.
Collecting docstrings means parsing every module -- measured here, 931
files and 19.5 MB through ``ast.parse`` is 1.89 s, against 0.03 s to
read the same bytes and 0.07 s to build the ``code`` corpus, which
needs no parse at all. That 1.9 s sat in front of the first request of
every console process, and a console is restarted often. Short
workloads never touched it (they draw the synthetic word list), which
is exactly why a long workload felt broken next to them.

So the built corpus is written beside the other caches and keyed by a
digest of the source it was built from. A tree that has not changed
reads it back; one that has rebuilds and rewrites it. ``code`` is not
cached: rebuilding it costs less than reading 19.5 MB back would.

Sourced from this installed package — its ``.py`` source for ``code``, the
long-form docstrings inside it for ``prose``.

LINES, not words. Flattening source into a space-joined run of tokens
destroys the indentation and line breaks that make it code at all, and a
model handed thousands of words of that does not answer it normally — it
rambles, which shows up as an absurd output length and a decode rate that
describes the rambling rather than the server. Keeping the line structure
is what makes the prompt realistic rather than merely long.

A realistic ``PROMPT_KINDS[kind_idx]`` prompt of about ``n_words`` words.

Whole lines from a seed-dependent offset until the word budget is met, so
the material reads as what it is and no two requests share a prefix — a
prefix-cache hit would serve one request's prefill out of another's KV and
measure nothing. The prompt states the task and stops: the decode length
is the model's own, never something the prompt asked for.

The concurrency this run should use.

``batch_fraction`` is resolved against the target's live ``max_batch`` so
the ladder tracks the scheduler the server was actually sized for: at
``max_batch=4`` the rungs are 1 / 2 / 4, at 64 they are 1 / 16 / 64. A
fraction floors at 2 — a middle rung of 1 would silently duplicate the c1
rung and report it as a batch. With no readable ``max_batch`` the declared
``concurrency`` stands, so an unreadable probe cannot quietly run a
different shape than the workload's name claims.

Prompt word-count per request.

``fill_context_frac`` scales the target's live ``max_context`` across the
workload's concurrency, so each request carries an equal share and the
batch together approaches a full pool. Falls back to the fixed ``n_words``
when the fraction is unset or the context is unknown.

The request's token cap, or ``None`` for "decode until the model stops".

A workload declaring ``max_tokens=0`` is open-ended by design: capping it
would truncate the long reasoning answers that are most of the decode in
the thinking regime, and measure the cap instead of the model.

Host-process RAM telemetry via ``psutil`` plus optional ``tracemalloc``.

Exports gauges:

  - ``arbi_serve.process.rss_bytes``  — resident set size
  - ``arbi_serve.process.vms_bytes``  — virtual address space committed
  - ``arbi_serve.process.pss_bytes``  — proportional set size (Linux),
    registered only under ``ARBI_SERVE_METRICS_PSS`` (default off)
  - ``arbi_serve.process.shared_bytes`` — shared memory (Linux)
  - ``arbi_serve.process.num_threads``
  - ``arbi_serve.process.num_open_fds``

PSS / shared are Linux-specific (``psutil.Process.memory_full_info``
on Windows / macOS doesn't surface them). Missing fields are silently
dropped from the per-callback observation list rather than emitted as
zero — Prometheus' "no data" rendering is unambiguous in the dashboard.

PSS is the one gauge here that is not free. It reads
``/proc/<pid>/smaps_rollup``, whose cost the kernel pays by walking every
mapping — 96.7 ms on the shipped 27B against 0.05 ms for everything else
in this module — so it is registered only when
``ARBI_SERVE_METRICS_PSS`` asks for it.

``tracemalloc`` is opt-in via the ``--enable-tracemalloc`` CLI flag.
When enabled at boot, :func:`enable_tracemalloc` calls
``tracemalloc.start(25)``; the admin endpoint
``/v1/admin/memory/python_heap`` then returns the top-N statistics by
traceback. Always-on tracemalloc adds 10–25% memory overhead to every
allocation, so it is firmly off by default.

Start ``tracemalloc`` so the admin endpoint can take snapshots.

Idempotent. Adds 10–25% allocator overhead, so the CLI gates this
behind ``--enable-tracemalloc`` (default off). ``frames`` controls
how many stack frames each allocation captures; 25 is enough to
walk through the FastAPI handler down into the engine without
truncating in our hottest paths.

Return a JSON-serialisable summary of the top-N allocators.

Each entry has ``traceback`` (newline-joined frames),
``size_bytes`` and ``count``. Sorted by size descending.

Raises :class:`RuntimeError` if tracemalloc was never started — the
admin endpoint translates that into a 409.

Incremental structural scan of a ``/v1/chat/completions`` request body.

A long chat body arrives over the wire as a sequence of chunks whose
boundaries are arbitrary byte offsets — possibly inside a UTF-8 sequence,
inside a JSON escape, or between the two halves of a surrogate pair. This
module turns that byte stream into the only two things the admission path can
act on before the body is complete:

* the chat turns that have arrived WHOLE, and
* a prefix of the turn still arriving, cut where it can be decoded exactly.

It is a STRUCTURAL scan, not a parser. Between two structural characters the
bytes are skipped in C by a compiled regex, and every byte of the body is
visited at most once across the whole scan no matter how it is chunked — which
is what lets a 700 KB body be rescanned on every arriving chunk without the
scan itself becoming the cost. In particular the escape walk is FORWARD (each
backslash jumps over what it escapes), never a backward parity count, so a
body that is one long run of backslashes costs the same as any other.

Nothing here decides anything about tokenization; it hands
:mod:`arbi_serve.server.prompt_prefetch` a snapshot and that module decides
whether a speculative encode is worth attempting. Exactness of the ids is not
this module's to guarantee either — a prefix it reports wrongly costs one
wasted encode, because the consumer proves the prefix relation against the
final text before reusing anything (see ``TokenizerPool._encode_from_prefix``).

What the scan can prove about the body received so far.

``fields``
    Top-level members of the request object that have arrived COMPLETE,
    excluding ``messages``. A member still arriving is absent, so a body
    that puts ``tools`` after ``messages`` yields a snapshot without them —
    the consumer's speculation is then wrong and is rejected downstream,
    which costs a wasted encode and never a wrong id.
``messages``
    The chat turns that have arrived complete, parsed.
``open_message``
    The completed members of the turn still arriving (``role`` and
    friends), or ``None`` unless the value still arriving is that turn's
    ``content`` string.
``open_content``
    A prefix of that ``content`` string, decoded exactly.
``complete``
    Whether the root object has closed, i.e. the body is whole. A consumer
    that sees this has nothing left to overlap with.

Feed body bytes in; read a :class:`BodySnapshot` out.

Stateful and incremental: each :meth:`feed` resumes where the last one
stopped. A body that is not a JSON object, or that the scan cannot follow,
sets :attr:`failed` and every later snapshot is empty — the consumer then
does nothing and the request is served exactly as it is today.

The end of the last COMPLETE UTF-8 character in ``buf[lo:hi]``.

UTF-8 is self-synchronizing, so this is a bounded look back from ``hi``:
at most three continuation bytes, then the lead byte that owns them. A
sequence whose trailing bytes have not arrived is excluded whole.

Bytes in the JSON escape starting at the backslash at ``p``, or -1 if incomplete.

A ``\uXXXX`` naming a HIGH surrogate counts its low half too: decoding the
high half alone yields a lone surrogate, which cannot be encoded back to
UTF-8 and would poison every later step of the prefix. The pair is either
both here or neither.

``(end, safe, cur)`` for the string body being read from ``cur``.

``end``
    Offset of the closing quote, or ``-1`` while the string is still
    arriving.
``safe``
    The furthest offset the string body may be CUT at — a character
    boundary that is neither inside a UTF-8 sequence, nor inside an escape,
    nor between the halves of a surrogate pair.
``cur``
    Where a later call should resume. Escapes already jumped over are never
    revisited, so the walk costs one step per escape for the whole scan
    rather than one per chunk.

Tokenize a long chat prompt WHILE its body is still arriving.

The pre-engine path for a long prompt is dominated by one encode, and today
that encode cannot start until the last byte of the request has landed: the
server reads the whole body, parses it, renders the chat template, and only
then tokenizes. From a browser on another host the body takes real time to
arrive, and the tokenize is stacked on top of that time rather than hidden
inside it.

This module overlaps the two. As chunks arrive it re-renders the chat template
over the turns received SO FAR — including a prefix of the turn still arriving
— and encodes that text. The rendered prefix and its ids land in the
:class:`~arbi_serve.server.tokenizer_pool.TokenizerPool`'s growing-prefix
cache, which is the same seam a second chat turn already uses; when the body
finally completes, the route's own encode finds the head already done and pays
only for the tail.

WHY THIS IS EXACT
-----------------
It does not need to be right. The pool proves the prefix relation against the
FINAL rendered text before reusing anything (length plus digest, then a
round-trip-verified backoff window), so a speculation that guessed wrong —
because ``tools`` arrived after ``messages``, because the template is not
concatenative, because the client sent a shape this module cannot follow — is
rejected there and costs one wasted encode on a worker thread. There is no
path by which a wrong guess becomes a wrong id, and therefore none by which it
becomes a wrong radix-cache key. Every gate below is a performance gate.

Two things this module DOES have to get right, because the pool cannot check
them:

* the prefix it hands over must be text the client actually sent — hence the
  exact-decode discipline in :mod:`arbi_serve.server.prompt_body_scan`;
* the cut must be somewhere the shipped backoff window can repair, which is
  whitespace. A run with no whitespace in it is one long merge the window
  cannot straddle, so such a prompt is DECLINED and served exactly as today.

WHAT IT COSTS WHEN IT DOES NOTHING
----------------------------------
A one-shot body (the whole prompt in a single chunk, i.e. every localhost
client) never triggers a speculation: the gate needs a chunk with more body
still to come. A short prompt never reaches the size floor. A non-chat route
never enters the middleware at all. In all three the request is served by
byte-identical code to today's.

Arm or disarm the overlap for this process; ``None`` restores the flag.

Arming also re-closes the circuit breaker, so this is the deliberate
re-arm a latched-open breaker waits for. Returns the value now in force.

The rendered-prompt prefix implied by ``snap``, or ``None`` to decline.

Renders the chat template over the completed turns plus a whitespace-cut
prefix of the turn still arriving, with :data:`_SENTINEL` marking where the
known text stops, and returns everything before the sentinel.

Declines — quietly, and the caller then does nothing — whenever the shape
is one this cannot speculate on: no turn in flight, a turn whose role has
not arrived yet, multimodal content, a cut with no whitespace to anchor to,
or a template that did not place the content verbatim (the sentinel is then
absent or repeated).

Read a chat body as it arrives and tokenize the part already here.

A **pure ASGI** middleware, for the reason
:class:`~arbi_serve.server.middleware.RequestIDMiddleware` gives: this one
sits on the request path of every streamed completion, and
``BaseHTTPMiddleware`` would wrap each of their responses in an extra anyio
task group.

Mounted INNERMOST, below auth, the serving gate and the admission
semaphore, so nothing is tokenized for a request that was never going to be
served. The body it drains is buffered and replayed verbatim, so the route
below reads exactly the bytes it would have read.

Start a speculative encode if enough new prompt text has landed.

At most one is in flight: a second would race the first for the single
cache entry they both write, and the later-finishing one would win with
the SHORTER prefix. Chaining them instead means each encode only covers
what the previous one did not.

Render the prefix implied by ``snap`` and encode it, off the loop.

The render is Jinja and the parse is orjson — both hold the GIL, and
the whole point of this module is that the event loop stays free to
keep reading the socket while they run, so neither happens on it.

The encode goes through the ordinary ``encode`` surface on purpose: the
ids, the parallel split for the first (cold) step, the prefix reuse for
every step after it, and the cache write are all the pool's, so this
module owns no second copy of any of them — and stores no prompt text
of its own, matching the pool's no-plaintext-retention property.

Let the last speculative encode land before the route encodes.

Awaited rather than abandoned, and that is the cheaper choice either
way: the work is a prefix of the text the route is about to encode, so
finishing it makes the route's encode smaller by exactly what it cost.
Abandoning it would leave the route doing the whole encode anyway AND
the worker thread doing this one.

Read the whole body, speculating between chunks; return what was read.

The messages are kept verbatim (``more_body`` flags and any
``http.disconnect`` included) and replayed to the app below, so nothing
downstream can tell the body was read here first.

In-process ring buffer of recent per-request stage timelines.

This is the data source behind the ``GET /admin/request_timeline``
endpoint, which feeds Grafana's multi-request *swimlane* (one
horizontal stage-coloured bar per request lifecycle phase, many
requests overlapped on a shared timeline -> concurrency is visible as
vertical overlap).

The exact same per-request timestamps that
:func:`arbi_serve.server.metrics.emit_request_trace` ships to Tempo as
an OTel span tree are *also* appended here on finish. Tempo keeps the
OTLP-native trace; this ring keeps a flat, query-cheap table so a panel
(Gantt / State Timeline / ECharts) can render the overlay without the
nested-trace flattening that the Tempo datasource makes awkward.

NOTE (data path): the endpoint is endpoint-sourced, NOT OTLP-sourced.
It reads from this ring, which the engine populates from the same
``req.timing`` stamps, whether or not any exporter is configured -- the
ring never leaves the process, so an export flag is not a precondition
for writing it. It is deliberately additive and never raises into the
engine loop -- :func:`record` swallows everything.

Each recorded entry is one request, expanded into one *row per stage*
(``queued`` / ``prefill`` / ``reasoning`` / ``answer`` / ``decode``) with
absolute millisecond start/end so a flat ``(request, stage, start, end)``
frame drops straight into a swimlane panel.

The expansion happens ON THE FINISH PATH, not on the endpoint. A finished
request is immutable, so its rows can only ever be built to one value;
building them per poll made the cost of LOOKING at the server proportional
to how long it had been up, and charged that cost to the process running
the forward pass, once per observer. Now the engine pays it once per
request and a poll concatenates.

This ring only ever holds FINISHED requests (appended on finish). For
requests still in flight -- still queued or mid-prefill/decode --
:func:`in_flight_rows` reads the scheduler's live ``waiting``/``running``
state directly, on demand, and the admin resolver merges both into one
response. That is a deliberate per-request live view, not a single
aggregate "how long has something been stuck" gauge: it costs nothing
on the hot path because it only runs when the endpoint is polled.

The tenant a live scheduler request is being served under.

Read off the request rather than passed in, because the live producer is
handed the scheduler's own deques and nothing else. The canonical id is
``req.tenant.tenant_id``; the flat ``req.tenant_id`` is the older scalar
spec and is read as a fallback so a Request built either way reports the
same namespace. Never raises: this is the observability path.

The correlation block every timeline row carries, from either producer.

ONE function for both producers (the finished-request ring and the live
scheduler read) because a field added to one and not the other is a row
that renders differently in flight than it does once finished — which is
a defect a reader sees as data.

``request_id`` is the engine's own per-process counter: it names a lane in
the swimlane but means nothing outside this process. ``client_request_id``
is the ``X-Request-ID`` the HTTP layer bound for the call — the value a
caller may supply, sees echoed on the response and reads in the logs — so a
bar is attributable to the exact call that caused it.

``None`` when the request reached the engine outside an HTTP scope (bench,
calibration, in-proc callers): a row that carried no client id says so by
absence, never by synthesising one that would correlate to nothing.

``tenant`` is the cache namespace the request was SERVED under
(:class:`~arbi_serve.tenant.context.TenantContext`), which is what makes
session affinity legible: with a sticky load balancer in front, whether two
turns of one conversation shared a prefix cache is a question about this
value. ``"anon"`` is a real answer and not a missing one — it means no
``X-Cache-Tenant`` header and no auth identity reached this engine, which
is also what a header stripped in transit looks like from here. The row
carries the value so a caller that DID send one can see it did not arrive.

``instance`` is the replica that served the request: the same name the
Prometheus ``instance_name`` label carries, so a bar and a dashboard series
can be pointed at the same process. Every row a process emits carries its
own name by construction — the ring holds only what this process served —
and it is on the ROW rather than on the response because a row is what gets
copied into a screenshot, a paste, or a bug report.

One finished request's timeline rows — built ONCE, when it finishes.

A finished request is a set of stamps that have already been taken, so
these rows can only ever be built to the same value. Building them on
the finish path (which already holds the entry) rather than on every
poll is what makes an admin read cost what CHANGED rather than what the
ring has accumulated since boot.

Stages emitted, only where their boundaries are stamped and the
interval is positive:

  - ``queued``: submit -> first scheduled step
  - ``prefill``: first scheduled step -> first token
  - ``reasoning`` / ``answer``: first token -> the measured
    reasoning->answer boundary -> completion, when the model reasoned
    and closed the block
  - ``decode``: first token -> completion, otherwise

NO ``prefill_cache_hit`` time row. It used to emit ``[t_prefill_start,
t_prefill_start + window * cache_frac]``, which was wrong twice over on
a time axis:

  * it scaled a TOKEN ratio by a TIME window, so a 60%-cached prompt
    was drawn as 60% of the prefill DURATION — a number nothing ever
    measured; and
  * the cached tokens are not in that window at all. The radix match
    resolves at ADMISSION, where the scheduler sets ``prompt_consumed
    = match_len`` and reuses the matched pages by pointer, so those
    tokens are SKIPPED. The prefill window contains only the uncached
    remainder — the cache hit's whole value is that it made this bar
    SHORTER.

There is no cache-restore stamp to draw honestly (a pure-attention hit
is a radix lookup ~ 0; the hybrid savepoint resume that does cost
something also runs pre-prefill). So the hit stays on the prefill row
as DATA — ``cached_tokens`` / ``cache_hit_frac`` — and the swimlane
marks the bar rather than inventing a sub-interval.

Thread-safe bounded ring of recent finished-request timelines.

Writes come from the engine's off-thread finish path; reads come
from the FastAPI admin handler (a different thread). We guard both
append and snapshot read with a lock so a concurrent append can't
tear the iteration.

When ``req``'s CURRENT queue residency began (Unix seconds).

``submit_time`` for a request that was never preempted; the
preemption stamp for one that was returned to the queue and is
re-prefilling from scratch, so its second wait is not drawn from the
original submission.

``req``'s queued->prefill boundary, or ``None`` when unavailable.

Reads ``RequestTiming.t_first_scheduled_step`` — the once-only stamp
the scheduler latches when it first admits the request to a forward
pass, and the same field the finished-request path uses as
``t_prefill_start``. ``None`` when the request carries no timing
object, when the latch has not fired, or when it predates the current
queue residency: the latch cannot express a re-admitted request's
second boundary.

When ``req`` stopped reasoning and began answering, or ``None``.

Reads ``RequestTiming.t_first_answer_emitted`` — the once-only stamp
latched on the first answer token, the same field the finished-request
path records. ``None`` while the model is still inside the reasoning
block (the answer has not started, so there is no boundary yet) and for a
request that never reasoned at all. Both draw as one ``decode`` bar; the
split appears the moment the stamp fires, which is what makes the
reasoning/answer distinction available DURING the request rather than
only from its trace.

Live per-request rows for requests still WAITING/PREFILLING/DECODING.

    This is the per-request answer to "is anything stuck right now" --
    deliberately NOT a single aggregate gauge. It reads live off the
    scheduler's own ``waiting``/``running`` deques at call time (the same
    collections the engine already maintains for admission), so it adds
    zero cost anywhere on the request hot path: this function only runs
    when something polls ``/v1/admin/request_timeline`` (Grafana, on its
    5s panel refresh), never during request processing.

    Each row's ``end_ms`` is pinned to ``now`` rather than a real
    lifecycle stamp, so the bar visibly keeps growing across repeated
    polls for as long as the request stays in that stage -- that growth
    IS the live signal. Rows merge into the same ``(group, stage,
    start_ms, end_ms)`` shape :meth:`RequestTimelineRing.rows` produces
    for finished requests, so one swimlane panel shows both without a
    second query. ``live: True`` distinguishes them for panel styling.

    The queued->prefill boundary is the SAME stamp the finished path
    uses: ``RequestTiming.t_first_scheduled_step``, latched once by the
    scheduler when it first admits the request to a forward pass. A
    running request therefore splits into a CLOSED ``queued`` segment
    (``submit_time`` -> first scheduled step) plus its live stage
    segment, exactly as :meth:`RequestTimelineRing.rows` renders the
    same request once it finishes.

The stamp is not always available -- the latch has not fired for a
    request still waiting, and a Request built outside admission carries
    no timing object. Without it the stage start is unknown and the bar
    anchors at ``submit_time`` with ``queued_split: False`` -- the row
    then states that it still contains the queue wait instead of implying
    a measured boundary.

    A DECODING request emits THREE rows: the closed ``queued`` and
    ``prefill`` segments and the live ``decode`` one; PREFILLING and
    SUSPENDED emit the closed ``queued`` segment plus their own live
    one. All carry the same
    ``request_id``, so a swimlane groups them into one lane and the bar
    only ever GROWS. Emitting just the current stage re-anchored the
    request's single row at the prefill->decode flip, collapsing a long
    bar back to the right edge so it read as a RESET -- and it made the
    same request render differently in flight than once finished, where
    the ring emits one row per stage.

    Preemption: a victim returned to the queue is re-prefilled from
    scratch, and ``requeued_time`` marks the start of that second queue
    residency. The first-schedule latch is once-only and cannot express
    the second boundary, so a re-admitted request anchors at
    ``requeued_time`` with ``queued_split: False`` rather than reusing a
    stamp that predates its current residency.

The open decode span, split at the answer boundary once it exists.

While the model is still inside the thinking block there is no boundary, so
the whole span is one live ``reasoning`` bar that grows. The instant the
first answer token lands the stamp latches, that segment CLOSES at it and
an ``answer`` segment becomes the live one — the same shape the finished
request will have, so a bar does not change form when the request
completes.

A request that is not reasoning draws as ``decode`` throughout. The two are
told apart by the request's own state, not by the absence of a stamp: see
the comment on the ``reasoning_tokens`` test below.

The recent entries' rows, one per lifecycle stage.

Each row is a dict with absolute epoch-millisecond ``start_ms`` /
``end_ms`` so a Grafana swimlane (Gantt / State Timeline /
ECharts) can lay bars on a shared time axis. ``group`` is the
request id (one swimlane per request); ``stage`` colours the bar.

The rows themselves were built by :func:`_rows_for_entry` when
each request FINISHED. This method selects and concatenates them,
so a poll costs what the window holds rather than rebuilding
every row in the ring from scratch.

Per-(model, backend, config) aggregates over the finished requests.

One group per distinct *model load* — model + backend + the compact
config fingerprint — so an A/B (a hotswap, a backend change, or a pure
settings flip on the same model) surfaces as separate rows. Each group
reports the averages that characterise throughput: prefill/decode length
(tokens) and prefill/decode rate (tok/s), plus the mean cache-hit share.
Rates are per-request means (each request's own tokens÷seconds), so a
few huge requests don't dominate a token-weighted pool.

One row for ``req``'s ``stage``.

``end=None`` pins the row to ``now`` — the growing live segment.
Passing an explicit ``end`` emits a CLOSED earlier segment of a
still-running request (see the prefill segment below).
``queued_split=False`` marks a bar whose start is the request's
queue entry rather than a measured stage boundary.

Per-resident tokenizer resolution for stable-VA multi-model residency.

A request for a NON-active resident is not switched to synchronously — it is
parked in the graceful queue and the run loop performs the fast park/wake swap
later. So the route tokenizes the prompt while a DIFFERENT model is still
active. For a cross-architecture pool (e.g. Qwen3.5 248k-vocab ↔ gemma
256k-vocab) the active model's tokenizer produces the wrong ids for the target
→ the target model decodes garbage.

The fix: the route encodes with the RESOLVED target resident's own tokenizer,
not the active one. This module builds the mode-agnostic provider that maps a
resolved resident key to its tokenizer:

* **in-proc**: the target member's own tokenizer is already loaded and lives in
  its residency snapshot (``rec.extra["state"].attrs["tokenizer"]``) — return
  it directly, no rebuild.
* **process mode** (or any snapshot miss): the API child has no engine
  snapshots, but it DOES share ``cfg``, so the member's checkpoint path is
  derivable (boot: ``served_name``/``path``; each member: ``served_name`` →
  ``path``). Build the tokenizer from that path; the pool caches it per key.

Map each stable-VA resident routing key to its checkpoint path.

The keys match the residency registry keys ``_maybe_switch_model`` resolves
to: the boot model keyed by ``served_name or path``, each declared pool
member by its ``served_name``. Both processes share ``cfg``, so the map is
reconstructible route-side without any engine round-trip.

The member's OWN tokenizer from the live in-proc residency snapshot.

Returns ``None`` in process mode (the cached routing view carries no
engine-state snapshots) or when the key has no captured state — the caller
then falls back to building from the checkpoint path. Mirrors the defensive
snapshot read of ``residency_view.resident_model_path``.

Build the resolved-key → tokenizer provider for a residency pool.

In-proc returns the member's already-loaded tokenizer from the snapshot;
process mode / snapshot miss builds it from the member's checkpoint path.
Returns ``None`` for an unknown key so the pool keeps the active tokenizer.

OpenAI /v1 surface, split across one shared :class:`APIRouter`.

The whole HTTP surface (chat / completions / embeddings / models / loras /
admin) decorates a SINGLE shared ``router`` created here — NOT per-group
sub-routers — so FastAPI's route-match order is exactly the import order of
the submodules below.

Each route module does ``from arbi_serve.server.routes import router`` and
decorates it; importing this package imports every route module (the
side-effect that registers the routes). ``arbi_serve.server.api`` is the
thin aggregator that re-exports ``router`` plus the handful of private
symbols ``batch_api`` and the test suite pull in.

Torch-free import discipline: NO ``routes/*.py`` may import
torch / ``Engine`` / multimodal at module top level, so
``_openapi_cli.build_schema_only_app`` stays engine-free. ``Engine`` is
``TYPE_CHECKING``-only; heavy / engine imports are function-local.

Torch-free request/response normalizers + SSE framing + sampling builders.

Pure translation helpers shared across the OpenAI route modules. NO
top-level torch / Engine / multimodal imports — this module participates
in the engine-free ``dump_openapi`` path, so the only imports here are
FastAPI, orjson, the request schemas, and the engine *request dataclasses*
(themselves torch-free).

Render tool defs as the dict shape Qwen3's chat template expects.

``tool_choice="none"`` suppresses the tool-calling affordance entirely so
the model answers in plain text — OpenAI semantics guarantee no tool is
called, which the template can only honor if the call format is absent.

Coerce an assistant tool_call's ``arguments`` to the dict the template needs.

OpenAI carries ``arguments`` as a JSON *string*; the Qwen3 template iterates
it as a mapping (one ``<parameter=k>v</parameter>`` per item) when a prior
assistant tool-call turn is replayed. Parse the string back to a dict; a
dict passes through, and an unparseable value degrades to ``{}`` so a
malformed prior turn never 500s the render.

Warn once when an explicit request value departs from the published profile.

The checkpoint publishes a sampling recipe per reasoning regime, and an
explicit request value always wins over it — deliberately, since a caller
who states a value means it. But silently winning is how a whole
measurement campaign ends up off-recipe without anyone noticing: every
request looked normal, and nothing named the deviation.

Fires only for a field the caller SET, and only where the checkpoint
actually publishes a value to depart from, so a model with no published
profile is silent. Reported at WARNING because it is actionable and
one-shot, not per request.

Resolve each sampling field to the request value or the model default.

For each backfillable field, an EXPLICIT request value (present in
``req.model_fields_set``) always wins. Only an OMITTED field — left at
the schema's hardcoded default — picks up the model's
``generation_config.json`` value, and only if that model default is
set (non-``None``). ``gen_defaults=None`` (or an EMPTY defaults object)
means "no backfill" — every field keeps the request value.

``thinking`` is the request's RESOLVED regime
(:func:`resolve_thinking_regime`). A hybrid-reasoning checkpoint ships
only its thinking profile in ``generation_config.json``, so the regime
selects which published profile backfills
(:meth:`~arbi_serve.generation_defaults.GenerationDefaults.for_thinking`);
``None`` means the surface has no reasoning block to resolve.

Build engine :class:`SamplingParams` from a chat-completion request body and resolved LoRA id.

``gen_defaults`` (a :class:`~arbi_serve.generation_defaults.GenerationDefaults`)
backfills sampling params the client OMITTED. ``None`` keeps the schema
defaults. ``thinking`` is the resolved regime that selects which of the
checkpoint's published profiles backfills.

The engine's ``max_tokens`` for a completions request.

``max_tokens=0`` (score the prompt, return no completion) still runs
one decode step: the prompt-logprobs pass rides the prefill, and a
zero-step request would never reach a terminal state. The response
builder drops that token.

Build engine :class:`SamplingParams` from a text-completion request body and resolved LoRA id.

``gen_defaults`` backfills OMITTED sampling params from the model's
generation_config.json; explicit request values always win. ``None``
keeps the schema defaults.

Resolve the concrete per-request thinking regime the drafter routes on.

``explicit`` is the client's merged directive
(:meth:`arbi_serve.schemas.openai.ThinkingResolutionMixin.resolved_enable_thinking`):
``True`` / ``False`` when the request pinned thinking on either API
surface, ``None`` when it left the knob at the model default.

In the ``None`` (model-default) case the model's OWN chat-template
default governs — read off the actual rendered prompt via
``opens_in_reasoning`` (:func:`arbi_serve.server.routes.openai_chat._opens_in_reasoning`),
the fabrication-free, per-checkpoint source of truth (Qwen3.6
undefined→ON, Qwen3.5 undefined→OFF, Qwen3-4B-Thinking always ON). This
is NEVER a hardcoded default — resolving ``None`` to a literal ``True``
would substitute an invented value (banned; the model default IS the
product decision).

Returns a DEFINITE bool: the concrete thinking fact for this chat
request, resolved once at admission so the drafter needs no re-resolution
per draft-row. (The ``None`` "no regime info" sentinel is reserved for
non-chat paths, which never call this.)

Build the ``usage.timing`` extension dict, or None if no timing.

The dict is the serialised form of :meth:`RequestTiming.to_payload`.
Callers attach it to ``usage`` for non-streaming and to the final
streaming chunk's ``usage`` field.

Only the per-step tier reports: every request carries lifecycle stamps,
but the block's step breakdown and first-call costs exist solely for a
request that asked for them (``return_timing`` / ``cfg.timing_debug``),
so a request that did not ask gets no ``usage.timing`` key.

Attach both usage detail blocks to ``usage`` in place, and return it.

ONE call site for the pair so the prompt-side and completion-side
breakdowns cannot drift apart across the six places a usage block is
assembled (chat non-streaming, chat final chunk, chat include_usage,
completions non-streaming, completions final chunk, completions
include_usage). Each block is still omitted when it has nothing to say.

Build ``usage.prompt_tokens_details``, or None to omit it.

``cached_tokens`` is stamped UNCONDITIONALLY (0 included), unlike the
completion-side details which are omitted when they have no producer.
A client cannot distinguish an absent field from a zero, so omitting a
cache miss would make it indistinguishable from a server that does not
report prefix caching at all — and a caching probe cannot tell "the cache
axis is disengaged" from "this build has no such field". Reporting the
zero is what makes the number falsifiable.

Build ``usage.completion_tokens_details``, or None to omit it.

``None`` (omit) unless the request actually reasoned OR ran speculative
decode — the engine sets ``reasoning_token_count`` to an int only for
requests that entered a ``<think>…</think>`` block (else ``None``, so a
non-thinking response never carries a misleading ``reasoning_tokens: 0``),
and ``mtp_proposed`` is > 0 only when MTP speculation was active.

``accepted_prediction_tokens`` / ``rejected_prediction_tokens`` are OpenAI's
predicted-outputs counters; here they carry the MTP speculative-decode
accounting — accepted = draft tokens that verified and were kept, rejected =
proposed − accepted. Omitted (not 0) when speculation was inactive. Shared
by the non-streaming usage and both streaming terminal usage chunks.

orjson-backed :class:`~fastapi.routing.APIRoute` for request bodies.

FastAPI/Starlette parse the inbound request body with the stdlib
``json.loads`` (GIL-held, slower than ``orjson`` on the hundreds-of-KB
chat payloads we see) **before** the handler runs and before
tokenize-offload. The response side already serialises with
``orjson`` (``_orjson_response`` / ``_sse_frame``); only the request
parse was still on the slow stdlib path. HTTP loop, each O(body bytes).

:class:`ORJSONRoute` reads the raw body once and pre-parses it with
``orjson.loads`` straight from bytes (C speed, GIL released for the
bulk), stashing the result on ``request._json`` so Starlette's
``Request.json()`` (which FastAPI calls to extract the Pydantic body)
returns the orjson dict without re-parsing. ``request._body`` is set too
so any later ``await request.body()`` is also a cache hit.

Error / content-type parity with stock FastAPI is preserved exactly:

* ``orjson.JSONDecodeError`` **is a subclass of** ``json.JSONDecodeError``
  and carries the same ``.pos`` / ``.doc`` / ``.msg`` attributes, so on
  malformed JSON we simply do *not* pre-stash ``_json`` and let the
  default handler re-read the cached body and raise the identical
  ``RequestValidationError`` (HTTP 422, ``type: json_invalid``). The
  response is byte-for-byte what stock FastAPI produces.
* The content-type gate below mirrors FastAPI's own logic
  (``application/json`` or ``*+json``, or no content-type when the
  endpoint is not strict): we only pre-parse when FastAPI itself would
  have called ``request.json()``. GET routes, form posts, and non-JSON
  content types are untouched (no body read, no parse).

Mirror FastAPI's gate for when it would call ``request.json()``.

Returns ``True`` only for content types FastAPI treats as JSON
(``application/json`` or an ``application/*+json`` suffix), or for a
missing content-type header (FastAPI defaults that to JSON for
non-strict endpoints). Anything else (forms, ``text/plain``, etc.)
returns ``False`` so we leave the body untouched.

``APIRoute`` that pre-parses JSON request bodies with ``orjson``.

Wire it onto a router via ``APIRouter(route_class=ORJSONRoute)``; the
behaviour is otherwise identical to the stock route (it delegates to
the default route handler after stashing the parsed body).

Per-request glue: LoRA / swap / admission / auth / timing / streaming.

The request-scoped collaborators the OpenAI route handlers share — none
of which belongs to a single endpoint. Multimodal preprocessing and the
group selector defer their heavy imports to call-time so this module stays
engine-/torch-free at import (it participates in the ``dump_openapi`` path).

``Engine`` is annotation-only here (``from __future__ import annotations``
keeps it stringified); importing it eagerly would drag torch into the
torchless OpenAPI smoke.

Lazily build + cache the model's image processor on ``app.state``.

The model's image :class:`MediaBinding` declares its preprocessor
family in ``extra["image_processor"]``; each family maps to one
processor class. An undeclared family raises rather than guessing —
a mismatched preprocessor silently garbles every image.

Lazily build + cache the model's audio processor on ``app.state``.

Dispatches on the served model's architecture: each audio-capable
model family ships its own mel/feature shape (Step-Audio-2's
Whisper-style processor vs NemotronVoiceChat's NeMo-style 128-mel
one), and feeding one model's audio through the other's processor
would silently produce garbage features rather than an error. Reuses
:func:`~arbi_serve.realtime.nemotron_voicechat_turn.uses_nemotron_voicechat_turn_loop`
as the dispatch signal — the same "is the served model
NemotronVoiceChat" question that module already answers for the
realtime turn-loop selection, so this stays one source of truth
rather than a second ad hoc architecture check.

Lazily build + cache the token2wav decoder on ``app.state``.

Normally the API child preloads it right after the engine reports ready
(``server/app.py``); this is the fallback when that hook did not run. The
lock makes it a true singleton: two concurrent first-audio requests would
otherwise each load flow+HiFT and each claim the vocoder reserve.

Resolve the request's audio-output plan → dict or ``None``.

``None`` when the request did not ask for audio. Raises 400 when it
did but the model / assets / voice cannot serve it.

Strip the RIFF header, returning raw little-endian PCM16 samples.

Walks the chunk table rather than slicing a fixed 44 bytes or searching for
the literal ``b"data"``. Both shortcuts are wrong: an extra ``LIST`` chunk
shifts the offset, and a naive search matches the ``data`` inside a
neighbouring chunk's payload. Either way the caller gets header bytes framed
as audio — a click at the start of every utterance, and a "did it return 200"
test would never notice.

Expand media placeholders + preprocess media → ``(ids, features)``.

Walks the model's ``mm_bindings`` (image / audio): preprocesses each
modality present in the message content parts and expands its
placeholder markers. Returns the (possibly expanded) token ids and
the per-modality features dict, or ``(ids, None)`` for text-only
requests or text-only models.

Resolve the per-request LoRA selection.

Surfaces accepted (in increasing precedence):

  1. ``model: "<base>:<lora_id>"`` — OpenAI-style colon split.
  2. body field ``lora: "<lora_id>"``.
  3. ``?lora=<lora_id>`` query param (the canonical
     per-request-override surface, consistent with
     ``?attention_backend=``).

Higher takes precedence on conflict. When the chosen LoRA id is not
in ``eng.lora_store`` we raise ``HTTPException(400)``.

Returns ``(base_model, lora_id_or_None)``; the ``model`` field is
rewritten to drop the colon suffix when used.

Wrap a streaming generator with safe_run.

:func:`safe_run` shields the engine cancel so the GPU work gets
cancelled even if Starlette tears down the request task while the inner
generator is mid-yield.

A 20 Hz ``is_disconnected()`` poller (``DisconnectHandler``) is
deliberately NOT used — ``safe_run`` alone is sufficient for
correctness (see the NOTE in the body).

The HTTP path cancels BY REQUEST ID (``Engine.cancel_by_id`` — P3
CancelMsg semantics; P4b maps it onto the wire): routes never hand
the engine ``Request`` back across the boundary, only its id. The
wrapper is safe to drop into existing ``StreamingResponse(...)``
call sites without any other change.

Resolve a request ``model`` string to a registered resident key.

The model-name contract (see :func:`_maybe_switch_model`). Returns the
resolved resident key, or ``None`` when the name does not map to any
resident (the caller turns that into a 404).

Two steps, in order:

  1. the conventional default alias (``"default"``) → the ACTIVE resident,
     checked ahead of the shared map so "serve whatever is resident" cannot
     be captured by a resident whose basename happens to be ``default``. An
     exact resident key still wins over it.
  2. everything else → :func:`~arbi_serve.engine.residency_view.
     resolve_resident_key`, the ONE identity map shared with the runtime
     load/switch path (exact key → unambiguous basename → the checkpoint
     path the resident was loaded from). Its docstring carries the order and
     the reasoning; routing adds no alias of its own, so a name that routes
     to a resident also loads to it.

True when ``resolved`` and the ``active`` resident are the SAME served
model under different runtime/capture config — i.e. one of them is a live
config-override VARIANT (``cfg-<sig>``) of the other, so a bare model-name
request should serve in place rather than switch.

Each config-variant record carries the operator-declared resident key it
derives from. Two named residents remain distinct even when they share the
same checkpoint storage; variants are transparent only within their own
named resident.

Request-driven stable-VA model auto-switch.

When multi-model stable-VA residency is engaged (``eng.stable_va`` is not
None and holds a registered resident set), resolve the request's ``model``
field to a registered resident and, if it is NOT the currently-active one,
switch the resident model BEFORE serving via the recapture-free
``aswitch_model`` park-then-wake.

Returns the RESOLVED resident key (or ``None`` when residency is
disabled / no model named): callers pass it to
``eng.asubmit(target_resident=...)`` so the engine-side factory can
re-validate against the ACTUAL active resident — in process mode
this routing decision ran against the API child's CACHED view, and
the factory's check is what makes a stale cache unable to mis-serve
(it performs the missed switch itself).

Model-name contract — a well-behaved OpenAI client must work WITHOUT
knowing the resident's filesystem path:

  * residency disabled (``eng.stable_va is None``) → no-op (the default
    single-model path is byte-identical — this is the FAST RETURN before
    anything else).
  * ``model`` empty / None / whitespace → the ACTIVE resident (serve
    whatever is resident — no switch).
  * ``model`` strips a ``:lora`` suffix (the LoRA resolver owns that), then
    resolves via :func:`_resolve_resident_key`:
      - a conventional default alias (``"default"``) → the ACTIVE resident;
      - an exact resident key (incl. a configured ``served_name``) → it;
      - an unambiguous path-basename of a resident key → that resident.
    A request resolving to the CURRENT model is the fast path (no switch —
    ``aswitch_model`` is itself a cheap no-op, but we short-circuit here to
    avoid even taking the swap lock).
  * a DIFFERENT registered resident → ``await eng.aswitch_model(target)``,
    which serialises on the engine critical section (drain + pause the
    scheduler) and re-checks the active model under the lock so concurrent
    requests for the same target coalesce into ONE switch.
  * an UNKNOWN / ambiguous model (resolves to no resident) →
    ``HTTPException(404)`` with the clear "Available: [...]" list. We never
    silently mis-route an unrecognised name to the active resident.

Single-model residency (one record) never switches: the lone key is always
active, so a matching request takes the fast path and a genuinely unknown
name 404s.

Gate admission via the watermark + queue-depth check.

Raises ``HTTPException(429)`` with ``Retry-After`` when the engine
refuses a fresh request; ``HTTPException(503)`` when the server is
draining. ``priority`` is forwarded to the engine's priority-aware
backpressure check so a batch backlog never 429s interactive
traffic and batch yields at the door.

Every refusal writes a server-side receipt first
(:func:`~arbi_serve.server.shed.record_shed`) — log line + counter,
naming the reason and the threshold that governed it.

Resolve the request's auth / tenant identity BEFORE submit (P3).

Resolving before publish avoids a race where the scheduler could read
``req.tenant`` for fairness / quota / the prefix-cache namespace before
identity was set. The identity folds into the ``SubmitMsg``
construction: callers splat the returned dict into
``eng.asubmit(**_auth_fields(request), ...)`` and the engine-side
RequestFactory binds it onto the ``Request`` before the scheduler can
see it.

Derivation is unchanged: the bearer token IS the tenant id (an
explicit ``X-Cache-Tenant`` header overrides); auth-off / anonymous
falls back to the empty string — the global cache namespace, so
single-tenant deployments keep cross-request sharing for free. The
rich :class:`~arbi_serve.tenant.context.TenantContext` the auth
middleware stamped on ``fast_request.state.tenant`` rides along BY
REFERENCE in-proc (process mode reconstructs from the tenant-id
string — P4b).

``fast_request=None`` (batch-worker path, no HTTP request object)
returns the anonymous defaults.

Resolve the per-request per-step-timing opt-in.

Three flips that all enable timing:
  1. ``cfg.timing_debug = True`` (server-wide opt-in).
  2. ``return_timing: true`` body field.
  3. ``X-Arbi-Return-Timing: 1`` request header (case-insensitive
     truthy values: ``1``, ``true``, ``yes``, ``on``).

Any one is sufficient. Returns False when none match.

Resolve the per-request scheduling priority class.

Two surfaces, body wins over the default but the header can opt a
request DOWN to batch even when the body omitted the field:

  1. ``priority: "batch"`` body field (validated to the
     ``interactive`` | ``batch`` enum by the request schema).
  2. ``X-Arbi-Priority: batch`` request header (case-insensitive).

Resolution: if the body explicitly set ``batch``, honour it. Else
consult the header — ``batch`` (case-insensitive) downgrades;
anything else (including absent) stays ``interactive``. The header
can only LOWER priority, never raise it, so it is safe to expose on
an untrusted edge: a client cannot use it to jump the interactive
queue. Any unrecognised value falls back to ``interactive`` — the
safe default that preserves today's latency contract.

Map a ``finish_reason == "error"`` terminal state to HTTP 500.

An ``"error"`` finish is an unhandled ENGINE failure (client-input
validation already rejects at submit time with 400) — returning it as
HTTP 200 with zero tokens makes a total outage look like success to
load balancers and load clients. Blocking endpoints raise 500 instead;
when the engine's infrastructure-failure latch is set
(``engine/infra_health.py``: CUDA OOM / cublas alloc / illegal
address), the latched reason is included so the response says WHY.
Streaming responses cannot change status mid-stream — they keep the
``finish_reason="error"`` frame; ``/health/ready`` (503 on the same
latch) covers outage detection there.

A backpressure DROP (``finish_reason in {"overloaded",
"admission_failed"}``) — the scheduler refused/abandoned the admission
(e.g. a recurrent-row pool exhausted past the deferred-queue bound, or
an unexpected admission failure) — maps to HTTP 503 with a
``Retry-After`` so it reads as transient backpressure, NOT a successful
200-with-zero-tokens (which would mask the refusal from load
balancers / clients). Like the engine-error case, a stream that has
already begun keeps its terminal frame.

Block until an embed/rerank request finishes (HTTP 408 on timeout,
HTTP 500 on an engine-error finish). Reads ClientRequest state only —
the OutputApplier fires ``new_token_event`` after applying the
terminal ``FinishOut`` (finish_reason + embedding / rerank_score).

Encode a whole request's texts in ONE batched, off-loop call.

Embedding / rerank requests fan out to one engine request per input,
so encoding each input separately would fire one tokenizer-pool
thread dispatch per input (queue-lock contention at high batch
sizes). A single ``encode_batch`` collapses that to one dispatch and
lets the Rust tokenizer parallelise internally.

``add_special_tokens`` lets the model's own tokenizer post-processor
append its trained terminal token. For Qwen3-Embedding that is
``<|endoftext|>`` (151643) — which is what last-token pooling reads —
NOT the ``eos_token`` (``<|im_end|>``, 151645); appending the latter
by hand reads the wrong token's hidden state and silently corrupts
every embedding. Reranking passes ``False`` (its chat-templated
prompt already ends at the position whose logits encode the
judgement; no terminal token is appended).

Serialise ``body`` with orjson directly.

Embedding / rerank responses are large float arrays (e.g. 256×1024
floats); FastAPI's default ``jsonable_encoder`` + stdlib ``json``
path is slower than ``orjson.dumps`` on that payload and shows
up as a top wall-time cost under load. orjson serialises Python
float lists natively.

Tear the watch down on every exit path, including the error ones.

A task left running past the response outlives the request it was
watching and calls ``cleanup_sync`` against an id that may since have
been reused — the watchdog would then cancel somebody else's request.

Admin surface: attention-backend hot-swap, profiling, distortion analysis,
build-batch stats, request timeline, cache flush, calibration, model reload,
config tuning, and memory release/resume/snapshot.

Split into one sub-module per concern; this package re-assembles them onto the
single shared :class:`~fastapi.APIRouter` so the public import path is
unchanged — ``from arbi_serve.server.routes.admin import router`` (and the
other module-level names) keeps working. The sub-modules are imported in a
fixed decoration order so the shared router registers the admin routes in a
stable order.

Every route body is a thin API-side validator plus ONE
:func:`~arbi_serve.server.routes.admin._dispatch.admin_call` into the
``ADMIN_METHODS`` registry (P5b, docs/engine_core_process.md §6) — the
same engine-side code path in-proc (marshaled onto the engine loop) and
in process mode (utility RPC).

All heavy / torch imports stay function-local so this package participates in
the engine-free ``dump_openapi`` path. The ``ARG001`` ruff ignore (``request``
args admin handlers keep for shape but don't read) is carried per-handler with
inline ``# noqa: ARG001``.

Repoint the app-state ``TokenizerPool`` at the engine's current
tokenizer after a model swap/reload.

The pool is built ONCE at startup around the boot model's tokenizer
(``app.py`` lifespan). A cross-architecture swap replaces
``eng.tokenizer`` with a different-vocabulary tokenizer; without this
repoint, the chat/completion routes keep tokenizing with the OLD
tokenizer, whose special-token ids can exceed the new model's vocab
(Qwen3.5-0.8B ``<|im_start|>``=248045 vs Qwen3-0.6B vocab 151936) and
index the embedding out of range → CUDA device-side assert. The engine
swap has already drained in-flight requests, so the rebind is safe.
No-op when no pool is installed (e.g. embeddings-only deployments).

Raise 409 if the engine's drain/swap lock is currently held.

The single guard behind every admin mutation that drains in-flight
work before reconfiguring the engine (backend swap, calibration
reload, model reload, config override, group swap). The lock is read
defensively so an engine without it is treated as "not locked".

ONE admin dispatch path: route → ``ADMIN_METHODS`` (P5b).

Design (``docs/engine_core_process.md`` §6)
-------------------------------------------
Every admin route body reduces to API-side validation (auth already ran
in the middleware; path confinement via ``resolve_admin_path``; enum /
string parsing) plus one :func:`admin_call`, which routes by mode:

* **process mode** (API child, ``app.state.engine_client`` set) —
  ``await engine_client.call(method, args)``: the RPC crosses as a
  ``UtilityCallMsg`` and executes on the engine loop via
  ``run_utility_call``.
* **in-proc modes** (thread / inline, ``app.state.engine`` set) — the
  SAME registry coroutine ``dispatch_utility(eng, method, args)``,
  marshaled ONTO THE ENGINE LOOP via
  ``LoopBridge.call_on_engine`` (own-thread mode) or awaited inline
  (single-loop mode). ``CriticalSection`` acquire/drain/mutate therefore
  runs on the engine loop in every mode — process-mode semantics, one
  code path.

HTTP error mapping: the boundary erases exception TYPES (an engine-side
failure crosses as ``UtilityResultMsg(error=..., error_types=[...])``),
so each route passes its ``except SomeError`` clauses as an
``http_errors`` name→status mapping. Matching walks the exception's MRO
names (concrete first) so subclass semantics survive the wire; the
``detail`` is the engine-side ``str(exc)`` verbatim — byte-identical to
the in-proc responses. Registry-level failures map uniformly:
``AdminArgError`` → 400, unknown method / unsupported-by-config → 501.
Unmapped failures re-raise (→ 500), exactly as an uncaught exception did
in the route bodies.

Heavy imports (``engine.proc.*`` pulls torch via ``messages.py``) stay
function-local so this package keeps participating in the engine-free
``dump_openapi`` path.

Normalize an :func:`admin_call` failure to ``(type_names, message)``.

For the rare route that needs a message-conditional status (the
memory release/resume "phase2 only" → 409 special case) on top of
the declarative ``http_errors`` mapping. ``type_names`` is the
exception's class-name MRO (concrete first); ``message`` is the
engine-side ``str(exc)`` — identical for both dispatch modes.
Returns ``None`` for non-engine failures (e.g. an ``HTTPException``
already mapped by :func:`admin_call`).

Execute one registry admin method for this request's app, by mode.

Returns the JSON-safe result (a dict for every catalog method except
``aunload_lora`` → ``None``). Raises ``HTTPException`` per
``http_errors`` + the builtin registry mapping; 503 ``engine not
ready`` when neither an engine nor an engine client is attached.

Dual-source 409 pre-check for drain-gated admin mutations.

In-proc: the direct ``eng._swap_lock`` read. Process
mode: the ``swap_in_progress`` READ RPC (the API child cannot see
the lock). Advisory either way — the lock itself is the
authoritative serializer; this keeps the operator-facing 409
behavior identical across modes.

Repoint the tokenizer pool after a model swap/reload — in-proc only.

Process mode deliberately no-ops here: the engine emits a
``model_switched`` ControlEventMsg and the API child's sink performs
the refresh (P4b wiring; see ``engine/proc/client.py``).

Flush radix prefix-cache pages — admin / privacy hammer.

``tenant=None`` (default) flushes EVERY tenant's cache tree.
Pass ``?tenant=<id>`` to scope the flush to one tenant. Pages
currently held by in-flight requests stay live until those
requests finish; the tree is detached so future matches miss.
Pages at refcount 0 (already evictable) are returned to the pool
immediately.

Returns ``{"freed_pages": N, "tenants_flushed": [...]}``.
Increments
``arbi_serve.prefix_cache.eviction_total{reason="admin_flush"}``.

Useful for: tenant-initiated cache wipe (privacy / right-to-erase),
operational reset before a benchmark, post-incident purge after
a config drift.

501 when the active page-table backend has no cross-request cache
to flush (the registry's ``AdminMethodUnsupportedError`` mapping).

List the calibration bundles in this deployment's configured dirs.

Scans exactly what auto-discovery scans -- the curated mount
(``ARBI_SERVE_CALIBRATION_DIR``, ``/cal`` under the shipped compose)
and the persistent autotune cache -- in the same precedence order,
one non-recursive ``*.json`` glob per directory. The route takes no
path input, so there is no directory outside that configuration a
caller can reach.

Read-only and best-effort: an unreadable directory is reported in
``missing_dirs`` rather than failing the call, because "the /cal
mount is not there" and "the mount is there and empty" are different
facts an operator needs to tell apart.

Drain, replace the calibration JSON, re-apply, resume.

Same async-lock + drain shape as ``/v1/admin/attention_backend``;
re-uses the engine's existing ``_apply_calibration_to_tkv_ops``
bit-width fail-fast guard.

Parameters
----------
path
    Absolute path to a Lloyd-Max calibration bundle.
backend
    When set, the JSON is registered as the per-backend override
    for that name. When omitted, replaces the global default.
min_improvement_pct
    Quality gate; 0 (default) skips it. Higher values run the same
    gate as ``--autotune``.

Which KV basis this engine is serving under.

Read-only and cheap. Answers two questions an operator needs separately:
whether a rotation is installed, and whether a bundle fit IN it exists --
a rotation without its bundle does not serve, because the codec refuses a
bundle in the wrong basis in both directions.

Per-cache REUSE-vs-COMPILE status captured at boot — every boot cache the
startup panel colours a phase by.

``present=True`` ⇒ reused (a **hit**, cheap); ``False`` ⇒ (re)built/JIT'd
this boot (a **miss**, the expensive phases). Merges the compile-cache
report (inductor/triton/tkv/…) with the flat-weights + activation-profile
caches recorded separately at boot. Each entry carries the ``phase`` it maps
to. Pure module reads — no engine RPC, never on the hot path.

Admin: live runtime-knob tuning + bespoke config override.

Covers ``/v1/admin/config`` (single dynamic knob), ``/v1/admin/config_override``
(sparse delta vs the active config, prepared+parked via the stable-VA pool)
and ``/v1/admin/pool`` (prepared-variant catalogue). These are ADMIN endpoints:
they change the whole engine → affect EVERY request.

``POST /v1/admin/config_override`` body — a sparse delta by param name.

Every key is a canonical param name (a registry entry in
:mod:`arbi_serve.config_overrides` — itself derived 1:1 from the boot CLI
args). ``{"max_batched_tokens": 4096}`` flips one thing;
``{"max_batched_tokens": 4096, "split_mixed_decode_prefill": true}`` flips
two; everything else inherits from the currently-loaded config.

Current active full config + known-param catalogue.

Returns the resolved value of every overridable param, the active variant
key + its delta vs boot, and the param catalogue (target / runtime_flag /
capture_affecting / doc per param).

Apply a sparse config delta vs the active config; return the now-active config.

Capture-affecting deltas (MBT / max_batch / max_context / GMU / prefill
buckets / capture flags) trigger a one-time build+capture; runtime-only
deltas matching a prepared variant are instant; a fresh-read-only flag delta
is applied live with no re-capture. 409 when a swap is already in progress.

Resume serving the last config that actually served.

The affirmative way back from an override whose rollback could not put a
serving member back — the only case that latches the sticky fault
(``/health/ready`` 503). An override that failed while the previous member
kept serving never latches and needs no restore. 400 when no config has
ever successfully served.

Ephemerally drop parked config variant(s), reclaiming their VRAM.

``?key=<variant>`` drops one parked variant; no query drops EVERY non-active
variant, resetting residency to the single active baseline member. This is
the ephemeral-A/B teardown: after a run the residency holds exactly the
active member — no permanent KV/memory tax. Refuses to drop the active member
(400). 409 when a swap is already in progress.

One read for one tick (``/v1/admin/console_snapshot``).

WHY THIS EXISTS, AND THE COST ARGUMENT IS THE WEAKER HALF.

Rendering one swimlane tick took THREE requests of the engine: the console's
``serving_state`` is itself ``/health/ready`` plus ``/v1/admin/server_info``,
and the timeline is a third. Three round trips to draw one panel is the wrong
number whatever each one costs — but the reason it is wrong is not the cost.

THE THREE READS WERE TAKEN AT THREE DIFFERENT INSTANTS. A panel could render a
queue depth from one moment against rows from another, and a readiness verdict
from a third; the picture it drew was not a state the server was ever in. A
single snapshot is not merely cheaper, it is COHERENT, and that is a
correctness property rather than a performance one.

WHAT IT IS NOT. Not a fat endpoint that recomputes everything for a caller who
wanted one field: ``include`` selects the parts, and a caller asking for one
part pays for one part. The DEFAULT is what a tick needs, because the tick is
the thing that repeats.

WHAT IS DELIBERATELY STILL SEPARATE. Panels on genuinely different cadences —
server logs at 3 s, swim metrics at 2 s, the memory card at 10 s — keep their
own reads. Folding a 10-second panel into a 1-second tick would make the tick
carry nine readings nobody looked at, which is the fat-endpoint failure with
extra steps.

Serve ``GET /v1/admin/console_snapshot``: everything one console tick needs, once.

``include`` is a comma-separated subset of ``serving`` / ``timeline`` /
``stats``; unknown names are ignored rather than refused, so a newer
console asking an older server for a part it does not have degrades to the
parts it does. Every part is optional and none is computed unless asked
for.

``serving`` is the LOAD BALANCER'S OWN VERDICT, not a second opinion: it
calls the same :func:`~arbi_serve.server.health.health_ready` an LB polls,
so the strip can never disagree with what routing would do. It carries
``memory_released`` explicitly because a console reading a 200 cannot
otherwise tell a serving engine from one that is merely ready-shaped.

Readiness as the load balancer sees it, plus the one thing 200 hides.

``health_ready`` returns a plain dict on 200 and a ``JSONResponse`` on
503 — both shapes are normalized here rather than in the console, so the
console never has to know how the probe reports.

``memory_released`` is read from the SAME engine view the gate used, not
from a follow-up ``server_info`` call. Two calls answered it at two
instants, and the second one was mostly redundant: the gate already
refuses an asleep engine with ``reason=engine_asleep``.

Admin: capture / build-batch debug counters.

Covers ``/v1/admin/build_batch_stats``, ``/v1/admin/capture_hist``,
``/v1/admin/prefill_coverage`` and ``/v1/admin/decode_pad_cudagraph`` —
the bench-probe instrumentation toggles + read-outs. The counter reads
and resets live in the ENGINE process (runner attributes + the
prefill-coverage module-global), so every route body is one
``admin_call`` against the registry's ``debug_stats_read`` /
``debug_stats_reset`` / ``set_decode_pad_cudagraph`` helpers.

Return ``_build_batch`` per-call timing stats.

Active when ``ARBI_TIME_BUILD_BATCH=1`` is set on the server
process. Otherwise returns zero counters. Used by the perf bench
harness for the CPU-prep parallelism probe.

Returns ``{calls, ns_total, ns_per_call_avg}``. Idempotent reads.

Reset build_batch timing counters.

Used by the bench probe to bracket a measurement window — the
bench script POSTs to reset before warmup, runs a fixed number
of requests, then GETs the accumulated counters.

Return the per-batch-size decode captured-graph hit/miss histogram.

Active when ``ARBI_DEBUG_CAPTURE_LOOKUP=1``. Maps each observed
decode batch size B to ``{hit, miss}`` counts so the bench probe
can measure cudagraph coverage under real concurrency. When MTP is
on, the verify-forward + drafter-chain histograms ride alongside
(``verify_by_shape`` / ``verify_pad_by_shape`` /
``drafter_by_shape``).

Return the piecewise PREFILL cudagraph coverage distribution.

Three outcomes per ``(num_tokens, num_seqs)`` shape — hits /
too_large_eager / unexpected_misses (the cap-and-eager design). The
acceptance gate is ``total_unexpected_misses == 0`` (the
too_large_eager bucket is the deliberate above-cap eager path,
expected non-zero at high concurrency / long context).
Lets the bench bracket a clean measurement window via the reset POST
below — no SIGUSR1 / file round-trip needed.

Flag-truth report: declared contract vs ACTUAL fire count per flag.

Returns ``{counters, rows, must_fire_failures}`` — the per-path
:mod:`arbi_serve.flag_truth` counter snapshot, the per-flag contract
evaluation (:func:`arbi_serve.flag_contracts.evaluate_contracts`),
and one formatted line per MUST_FIRE flag whose path shows zero
fires. The GPU gate (``tests/test_flag_truth_live_gpu.py``) drives
the canonical smoke and asserts ``must_fire_failures == []``.

Flip the decode-cudagraph pad-up path on/off at runtime.

Same-process toggle so a bit-exactness harness can compare ON vs
OFF without a cross-process bf16 GEMV-vs-GEMM noise confound. Only
intended for validation / benches.

Per-layer per-bit dequant relative error for every TKV backend.

Runs an in-process analysis (no scheduler hop) and returns a
per-(backend, layer) rel-L2 read-out — the
"show me the calibration's quality on this model" one-call API.

Query params:
  - ``max_tokens_per_prompt``: cap per prompt (default 64).
  - ``num_prompts``: how many of the engine's eval-set prompts to
    use (default 8). The eval set is taken from the active
    recipe's ``eval_set`` field if present, else a small built-in
    fallback — resolved ENGINE-side (``eng.eval_prompts`` is engine
    state; see the registry's ``distortion_analysis``).

When the model does not expose a ``capture_kv()`` hook the report
is computed against synthetic Gaussian K/V; this measures the
codec's intrinsic quantization error per (layer, bit width) — a
useful relative ranking but not a model-conditioned error.

Admin: memory release/resume + flat-weight dump + memory snapshots.

Pairs the stable-VA (cuMem) sleep/wake release/resume endpoints with the
flat-weights dump, the Python-heap (``tracemalloc``) snapshot, and the
CUDA-side memory-snapshot tooling (live snapshot + allocation-history
record start/stop/dump). All ``/v1/admin/*`` so they sit behind
:class:`AuthMiddleware` like every other admin route.

Everything CUDA-allocator-adjacent (release/resume, live snapshot,
record start/stop/dump) executes engine-side through the registry — the
allocator and the named pools live in the ENGINE process.
The one exception is ``memory/python_heap``: ``tracemalloc`` is
per-process and this route reports THIS (API) process's heap, so it
stays API-side by design.

Shared release/resume dispatch with the "phase2 only" 409 shim.

``release/resume_memory_occupation`` on a pool running the CPU/no-GPU
stub raises ``RuntimeError("... phase2 only ...")``; surface that as
a clean 409 (unsupported) rather than an uncaught 500 so callers can
detect it gracefully — any other engine failure re-raises (→ 500).

Drain in-flight requests, then exit this process.

Named for what the server actually does. Whether the exit becomes a RESTART
is the supervisor's policy, which no process can observe about itself, so
the response reports what the deployment declared rather than promising an
outcome. A caller that wants to offer a "restart" reads ``expect_restart``.

Not refused when unsupervised: stopping the server is a legitimate action
on its own. It is the CALLER's job to say which one it is offering.

Drain in-flight, drop every GPU allocation, mark released.

``release_memory_occupation`` is the stable-VA (cuMem) sleep path. On a
pool running the CPU/no-GPU stub it raises; surface that as a clean 409
(unsupported) rather than an uncaught 500 so callers can detect it
gracefully.

Rebuild every GPU allocation released by ``release_memory``.

``flat_dump_dir`` (optional query param) is ACCEPTED AND IGNORED on this
route. The stable-VA backend remaps the released physical pages at their
original VAs, so there is no weight reload for a flat dump to accelerate;
``resume_memory_occupation`` deletes the argument. Flat dumps are still
consumed by the stable-VA residency register
(``ResidentModelRecord.flat_dump_dir``), a different entry point.
The param stays on the signature so already-generated clients keep working.

``flat_dump_dir`` is a QUERY param, not a request body, on purpose: an
optional request body makes ``openapi-python-client`` serialize an omitted
body as the literal ``UNSET`` sentinel (``TypeError``), so a client calling
``resume_memory()`` with no args would crash. A query param has no such
footgun — the generated client omits it cleanly.

Dump every model parameter as a single safetensors file.

Run after warmup. The dump is what a stable-VA resident registered with
``flat_dump_dir`` reloads from on a cross-model wake; ``resume_memory``
ignores the path (see that route).

Top-N Python-heap allocators by size, via ``tracemalloc``.

Returns 409 when tracemalloc isn't running. Each entry has
``size_bytes``, ``count`` and a formatted ``traceback`` list.

Deliberately NOT routed through the admin registry: ``tracemalloc``
is per-process state, and this surface reports the process serving
the HTTP request (the API child in process mode). An engine-process
heap RPC is a separate, not-yet-needed method.

Live snapshot of GPU memory grouped by named MemPool.

Returns:

``pools`` — per-pool ``{allocated_bytes, reserved_bytes,
num_blocks, num_segments}`` from
:func:`arbi_serve.runtime.memory_snapshot.live_snapshot`,
including a ``default`` bucket for segments not attributed to any
named pool.

``device`` — ground-truth totals from ``torch.cuda.memory_allocated``
/ ``memory_reserved`` AND ``torch.cuda.mem_get_info`` (driver-
visible used / free / total). The latter is hardware truth and
matches ``nvidia-smi``.

``reconciliation`` — ``{accounted, residual_alloc, residual_reserved}``.
``accounted`` sums the named pool entries. ``residual_alloc =
device.allocated - accounted.allocated``; non-zero means torch
holds bytes not attributed to any named pool we know about
(NCCL workspace, sleepable buffers, raw default-pool allocations,
etc.). Tracking this delta is how we keep the budget honest.

``recording`` — whether allocation-history recording is active.
Optional ``top_allocations`` (present when ``top>0``) — the top-N
live allocations by size. The CUDA allocator + named pools live in
the ENGINE process, hence the registry dispatch.

``card`` — THE CLOSED VIEW, and the block a renderer should use:
``rows`` (resident physical, ``driver.residual`` included and computed
server-side from the card down), ``free_rows`` (the unallocated remainder
named by what will allocate it — the serving-step reserve itemised term by
term, and ``unclaimed.*`` for the bytes nothing declares) and ``ghost_rows``
(address space holding no physical). ``rows_bytes + free_rows_bytes ==
total_bytes`` exactly when ``closes`` is true; when it is false the
identity has a hole and ``unaccounted_bytes`` names its size. A client does
no arithmetic: every row the client would have to derive is a row the server
can name and the client cannot.

``taxonomy`` — the pool taxonomy (category headline, per-pool
caption, per-row PROVENANCE and the leak-signal flag) from
:mod:`arbi_serve.runtime.pool_taxonomy`, so a UI names, captions and
qualifies every row from the source of truth rather than a copy.

``driver_measured`` — the ``driver.*`` rows (CUDA context, kernel-stack
local memory, loaded cubins, cudagraph exec, serving-time cubin loads,
foreign processes). They are outside every allocator, so a view of the card
that omits them is short by their sum. The terms that can move while
serving are RE-MEASURED on each read and the post-boot growth is bracketed
as ``driver.modules_serving``. ``driver.residual`` is excluded from THIS
block because it is the remainder of the driver sub-identity alone; the
whole-card residual is computed for the ``card`` block above and named
there.

``unpooled_owners`` — the tensors holding
``unpooled.torch_default_pool``, named by the boot heap walk.

``state_arena_mapped`` — the ``state.gdn_recurrent.mapped`` row (the state
pools' sentinel-alias arenas), same four states as ``kv_mapped``. The other
cuMem region no allocator snapshot can see; on a hybrid model it is the
whole recurrent slab.

``kv_mapped`` — the ``state.attn_kv.mapped`` row (the physically-mapped
growable KV slab) with the STATE of the reading: ``mapped`` / ``pooled``
(its bytes are on ``state.attn_kv``) / ``absent`` (no KV pool) /
``unresolved`` (the read failed, the bytes are unknown). It is the largest
row of a serving card and invisible to the allocator snapshot, so it is
served from the same in-process source as the boot VRAM ledger rather than
left to a second system.

``unresolved`` — ``[{name, reason}]`` for every contribution this snapshot
could not establish. Non-empty means the rows do NOT sum to the card, and a
consumer must say so rather than book the difference on a residual row.

Dump the recorded snapshot to ``path`` (default
``/cache/memory-snapshots/<timestamp>.pkl``).

The pickled file format is what
https://pytorch.org/memory_viz consumes — drop it into the viewer
to see the alloc / free timeline + per-call-site flame graph.

Admin: switch-to-or-load a model (``/v1/admin/model``).

One coherent runtime-model surface. ``mode="swap"`` (default) routes through the
fast stable-VA park/wake — switch to the model, building it into the residency
if absent — so a subsequent swap-back is fast. ``mode="reload"`` is the explicit
destructive fallback (drain → free → rebuild), the same one-notch-larger drain
as a backend hot-swap.

``admin_call`` whose UNMAPPED failures still name their cause.

``http_errors`` above enumerates four types; anything else reached
FastAPI's default renderer, so a refused hot-swap told the operator only
``HTTP 500 Internal Server Error``. One live-arch run reported exactly
that sentence for two unrelated engine faults. The status stays 500 — it
IS a server fault — but the body carries ``<ExceptionType>: <message>``.
Scoped to this route: other admin routes re-map the original exception
themselves and need ``admin_call`` to re-raise it untouched.

Admin: list the checkpoints available in the model store for the reload picker.

The engine already holds the resident model's full path (``eng.cfg.model.path``);
its parent directory IS the store (e.g. ``/models/Qwen3.8-27B-exl3-4.0bpw`` →
``/models``). Scan that parent for sibling checkpoint directories (a dir holding
a ``config.json``) so the admin console can offer a model DROPDOWN instead of a
free-text path. Read-only, best-effort: no engine loop, no new env flag; if the
store can't be resolved or read, return an empty list rather than failing.

Sibling checkpoints of the resident model, for the reload dropdown.

Returns ``{"root": <store dir or None>, "models": [{"path", "name",
"resident"}]}``. The resident model is always included and flagged, even if
the scan finds nothing else, so the dropdown never comes up empty on a live
engine.

Admin: live GPU profiler (``/v1/admin/start_profile`` + ``/v1/admin/stop_profile``).

These sit under ``/v1/admin`` (not the bare root paths) so they are
gated by the admin-authorization check in :class:`AuthMiddleware` like
every other operator endpoint, and inherit the ``/v1/admin``
concurrency-exemption so a profiler toggle never trips backpressure.

Start a live ``torch.profiler`` on the running server.

Profile the REAL server (incl. TP2) without a restart. The
trace config is read from ``ARBI_PROFILE_DIR`` /
``ARBI_PROFILE_WITH_STACK`` / ``ARBI_PROFILE_RECORD_SHAPES``. Under
TP>1 the driver broadcasts the start to every rank so each dumps its
own per-rank Kineto trace. Fails 400 if a profiler is already running
or the dir is not writable.

Profile a small steady-state window: POST /v1/admin/start_profile,
drive a few seconds of decode traffic, then POST
/v1/admin/stop_profile to flush the trace(s) under ARBI_PROFILE_DIR.

Stop the live ``torch.profiler`` + flush its trace(s).

Under TP>1 the driver broadcasts the stop so every rank flushes its
own trace. Returns the output dir. Fails 400 if no profiler is
running.

Admin: small "what am I looking at" header strip for the admin UI.

No existing route returns model + uptime + version + TP size together
(see the arbi-serve admin-UI scouting notes) — this composes them from
scattered existing sources. Read-only, best-effort: any field this
process can't resolve comes back ``None`` rather than failing the
whole call.

The turbo-attn (``tkv``) backend version the running server is built
against, read at RUNTIME from the actually-loaded ``tkv`` — the pin baked
into the serving image, or a bind-mounted dev copy.

arbi-serve carries no release version of its own: a turbo-attn release
carries arbi-serve forward automatically, and two arbi-serve builds on the
same turbo-attn differ only by git sha (see :func:`_read_git_sha`). This is
deliberately SEPARATE from the API-surface version in the OpenAPI document
(``FastAPI(version=...)``), which bumps only when the HTTP contract changes,
never when the backend does.

The arbi-serve IMAGE's build commit — an ``ARBI_GIT_SHA`` launcher stamp
if present, else the ``/etc/arbi-build-sha`` baked at image build. This
identifies WHICH image; whether the running code is that image or a
bind-mounted checkout is a SEPARATE fact (see :func:`_running_from_bind_mount`
/ the ``dev_mount`` field), never conflated into this sha.

True when the running ``arbi_serve`` is a source checkout (a ``.git`` at
its repo root) rather than the image's installed package — i.e. bind-mounted
dev code overriding the image. The image sha then labels the image, not the
running code, so the header shows a distinct 'bind-mount' marker.

Admin: which ``arbi_serve`` source this engine is running.

Read by the admin console beside it, which is the same repository and
must not serve a different tree than the engine it drives. The comparison
and the reason it exists live in :mod:`arbi_serve.source_identity`.

Content identity of the ``arbi_serve`` package this process imported.

A pure filesystem read of the running package -- no engine RPC, no
torch, and answerable while the model is still building, which is when
the console asks. The image label (``git_sha``) and whether the code
is a bind-mounted checkout are the SEPARATE facts ``server_info``
already carries; this route answers what those two cannot, which is
what the bytes are.

Serve ``GET /v1/admin/request_timeline``: flat per-stage timeline of recent + in-flight requests.

Powers the Grafana multi-request *swimlane*: one row per request
lifecycle stage (``queued`` / ``prefill`` / ``prefill_cache_hit`` /
``decode``) with absolute epoch-ms ``start_ms`` / ``end_ms``, grouped
by request id so overlapping bars show concurrency. Rows come from
two sources merged together: the timeline ring (finished requests,
written by the engine finish path) and a live read of the scheduler's
``waiting``/``running`` state for anything still in flight (``live:
true``, ``end_ms`` pinned to now so the bar visibly grows on each
poll) — engine-process state either way, hence the RPC. This is an
endpoint, NOT an OTLP read. 503 ``engine not ready`` before an engine
(or engine client) is attached.

SERIALISED ONCE, NOT PER POLL. A finished request's rows are built on
the finish path and concatenated here; the response is encoded straight
from the engine's own dicts with ``orjson``. It is deliberately NOT run
back through the ``response_model``: that validated every row into
``TimelineRow`` models and dumped them again, which is two full passes
over data this process constructed itself and already trusts, on a
payload that is thousands of rows. ``response_model`` stays declared so
the schema, the OpenAPI document and the generated client are unchanged;
``tests/test_timeline_serialised_once.py`` holds the body to it.

Fill in the requests that are ADMITTED but not yet running.

The two halves of "in flight" live in different processes and neither
could see the whole. The engine reports what its scheduler is RUNNING;
the HTTP layer's ``RequestSemaphore`` holds a permit for every request
admitted past backpressure — running AND waiting — and a request that has
been submitted but not yet pulled into the scheduler is in neither the
engine's ``running`` nor its ``waiting``. It sat in the intake, invisible,
which is why a 24-request burst against an 8-wide batch reported eight in
flight and nothing queued.

So the depth is the difference, taken from the one process that holds both
numbers at once. Reported as a COUNT rather than as rows: the engine has no
per-request handle on a request it has not admitted, and inventing rows for
them would put bars on a timeline for work with no measured stamps.

``None`` when the semaphore is not mounted (a test app, an embed-only
build) — null, never zero and never absent, because "nobody is waiting"
and "nobody counted" are different answers and a missing key collapses
them. The response used to be rebuilt through ``RequestTimelineResponse``,
whose field defaults supplied the two keys when this function bailed out;
the body is now the wire form, so it carries them itself.

Serve ``GET /v1/admin/request_stats``: per-(model, backend, config)
throughput aggregates over the finished-request ring.

One group per distinct *model load* — model + backend + config
fingerprint — each with average prefill/decode length (tokens) and rate
(tok/s), so an A/B (hotswap, backend change, or a settings flip) is
directly comparable. Backs the swimlane's completed-request summary. 503
before an engine is attached.

Serve ``POST /v1/admin/request/{request_id}/answer_now``: stop this request reasoning.

The transition the thinking budget performs, decided by an operator rather
than by a token count. The request is armed to emit its closing marker on
the next forward pass, so the model TRANSITIONS to answering instead of
being truncated: the response stays well-formed and the reasoning->answer
boundary is stamped exactly as a natural close stamps it.

This CHANGES WHAT THE CALLER RECEIVES — the completion is answered from a
shorter reasoning trace than the model would have produced. It is counted
under its own path counter (``thinking_closed_by_operator``, distinct from
the budget's) so an intervention is never read back as the configured cap
engaging.

``request_id`` is either the caller's own ``X-Request-ID`` (matched first)
or the engine's per-process counter, which is what the swimlane shows. A
caller who sent no ``X-Request-ID`` gets a generated one back on the
RESPONSE header — after the request is over — so in practice this endpoint
is addressable only by a caller who chose to supply an id up front, or by
an operator reading the id off the timeline.

Refuses, with the reason, on a model with no reasoning surface, a request
that is not running, and a request not inside an open thinking block.

Serve ``POST /v1/admin/request/{request_id}/stop_now``: stop this request generating.

Forces a stop token on the next sampling step, so the request ENDS the way
it would have ended anyway -- a real ``finish_reason``, usage reported, KV
released down the ordinary path, the trace written. A dropped connection
gives none of that, and on this stack it is not even reliably noticed.

Honoured under ``ignore_eos``: that flag declares the MODEL's stop
uninteresting, not the request unstoppable, and a decode-rate bench is
exactly the kind of run an operator needs to be able to stop.

Refuses, with the reason, when the tokenizer publishes no stop id, when the
request is not running, and when it has produced no token yet -- a
mid-prefill chunk samples nothing, so there is no step to carry the forced
id and reporting success would promise a stop that lands only once prefill
ends.

Admin: watermark detection — score text or tokens for the keyed watermark.

Model-free math (:func:`arbi_serve.sampler.watermark.detect_tokens`) plus a
tokenizer round-trip for raw text; no engine dispatch. Admin-scoped on
purpose: a public detector is an oracle an attacker can iterate against to
learn how to strip the mark. Key material is accepted (as an override for
rotated-out historical keys) but never logged and never echoed back.

``POST /v1/audio/transcriptions/stream`` — continuous streaming ASR / S2TT.

The OpenAI-shaped batch ``/v1/audio/transcriptions`` (upload a file, get one
transcript back) is the wrong primitive for the real product need: fluid
transcription of a source that runs for minutes-to-hours, with text yielded as
it is recognised. Step-Audio-2-mini is generative with a bounded context, so it
cannot ingest an hour in one forward.

This endpoint streams BOTH ways. The request body is a live PCM16 feed (mono,
little-endian, ``sample_rate`` query param); as it arrives it is driven through
:class:`~arbi_serve.realtime.streaming_transcriber.StreamingTranscriber`, which
VAD-segments it and transcribes each utterance as one bounded call. Finalized
segments stream back as SSE frames the instant an utterance ends — so a client
sees transcript within a sentence of the speaker, for as long as the source
runs, with flat memory (``O(one utterance)``; see the transcriber's docstring).

Frames (``data: {json}\n\n``):

  * ``{"type":"transcript.segment", "text","start_ms","end_ms","language",
     "is_final":true}`` — one finalized segment. Timestamps are absolute
     (from stream start) at VAD-boundary granularity, NOT model word timings.
  * ``{"type":"transcript.completed", "text": <full>}`` — end of stream.
  * ``{"type":"error", "error":{...}}`` — a mid-stream failure (then ``[DONE]``);
    the stream never just stops, so a client can tell truncation from a clean end.
  * ``data: [DONE]`` — always the last frame.

Input is raw PCM16 only: a live feed can't carry a decodable container header,
and streaming-decode of compressed audio is out of scope. A finite short clip
that just wants one transcript should POST it to the (batch) endpoint instead.

Query params:

  * ``sample_rate`` (default 24000) — the PCM16 body's rate.
  * ``target_language`` — the OUTPUT language knob. Omit it (or pass
    ``original`` / ``none`` / ``auto``) to transcribe VERBATIM in whatever was
    spoken; pass ``en`` / ``zh`` / ``ja`` to FORCE every utterance into that
    language regardless of input. Only en/zh/ja are real targets — the model
    silently returns Chinese for anything else, so an unsupported target is a
    clean 400 (``translate_to`` is a back-compat alias for this).
  * ``vad_backend`` (auto|silero|energy|ten), ``vad_threshold``, ``silence_ms``
    — turn-detection knobs. ``max_tokens``, ``voice`` as usual.
  * ``diarize`` (bool) — label each utterance with a speaker (``"S1"``/``"S2"``/…)
    by CAM++ x-vector clustering; ``max_speakers`` caps the cast (e.g. 2 for a
    witness + examiner). Needs ``onnxruntime`` + ``torchaudio`` + the checkpoint's
    ``token2wav/campplus.onnx``; absent any of these, diarization is silently off.
  * ``prosody`` (bool) — attach intonation markers (pitch trend / loudness) to each
    utterance. Deterministic acoustics, no extra latency; intonation, not sentiment.
  * ``emotion`` (bool) — attach a model-inferred emotion word (Step-Audio-2's
    paralinguistic understanding) to each finalized utterance. This is SEMANTIC
    sentiment and costs a SECOND bounded model pass per utterance, so it is off by
    default and adds latency; ``prosody`` is the zero-cost acoustic alternative.

Beyond finalized segments the stream also emits interim ``is_final:false`` frames —
the growing hypothesis of the utterance the model is still decoding (the live
reveal). A client replaces the in-progress line on each partial, commits on the final.

The wire shape of one segment — identical for SSE and WS.

``is_final`` is ``False`` for an interim (still-decoding) hypothesis and
``True`` for the committed segment; a client replaces the current in-progress
line on each partial and commits it on the final (same ``start_ms``).

The forced output language, or ``None`` for verbatim ("original").

``target_language`` is the primary knob and matches how an operator thinks
about it: OMIT it (or pass ``original`` / ``none`` / ``auto``) to transcribe
in whatever language was actually spoken; pass a code (``en`` / ``zh`` /
``ja``) to force EVERY utterance into that language no matter the input.
``translate_to`` is accepted as a back-compat alias. An unsupported forced
target is rejected downstream (the model silently returns Chinese otherwise).

Continuous transcription over a WebSocket — the browser-mic transport.

A browser cannot stream a ``fetch`` request body over HTTP/1.1 (Chrome
requires HTTP/2), so a live microphone feed can't drive the SSE POST route.
WebSocket is the correct browser primitive: bidirectional, no duplex/HTTP2
caveat. Protocol:

  * client -> server: FIRST a JSON config frame
    ``{"sample_rate":16000, "target_language":"en"|null, "vad_backend":...,
       "max_tokens":...}``; THEN binary PCM16 frames (mono, little-endian).
    A JSON ``{"type":"flush"}`` finalizes the utterance in progress without
    closing; ``{"type":"stop"}`` finalizes, emits ``completed``, and ends;
    closing the socket does the same.
  * server -> client: ``{"type":"ready"}`` once configured, then
    :func:`_segment_payload` frames as utterances finalize, a
    ``{"type":"transcript.completed","text":...}`` at end, or
    ``{"type":"error", "error":{...}}`` for a bad config / mid-stream fault.

Same VAD-segmented core as the SSE route (``target_language`` semantics
identical): omit/original -> verbatim, en/zh/ja -> forced.

``/v1/embeddings`` + ``/v1/rerank`` (+ no-prefix aliases).

``_run_embeddings`` is the embedding core shared by the HTTP route and
the Batch API worker; ``batch_api`` imports it via the ``server.api``
aggregator. Heavy imports (engine embed ops, SamplingParams) stay
function-local to keep this module engine-free at import.

One access-log IN line for a fan-out API call.

Embed / rerank / score turn one API call into one engine request per
input, so the engine's own per-request pair logs at DEBUG and the call
is announced here instead — once, with the count that says how much it
contains. Its OUT line is the ``done`` summary at the end of the same
handler; the ``request_id`` both lines carry in their log context pairs
them when calls overlap.

Embedding core, shared by the HTTP route and the Batch API worker.

Fans out one engine request per input (all routed through
``eng.asubmit`` for continuous batching) and assembles the OpenAI
``list`` response dict. ``fast_request`` is optional: the HTTP route
passes it so auth/tenant binding happens; the batch worker passes
``None``. ``priority`` is stamped on each request's SamplingParams so
batch-origin embeddings (worker passes ``"batch"``) yield to
interactive ones. Raises ``HTTPException`` on validation / timeout
failure.

Rerank core, shared by the HTTP route and the Batch API worker.

Mirrors :func:`_run_embeddings`: ``fast_request=None`` on the batch
path (anonymous auth), and ``priority`` is stamped on each engine
request so batch-origin rerank yields to interactive traffic.

Score core (vLLM score-API), shared by the HTTP route and the
Batch API worker.

Pairs ``text_1 × text_2`` and scores each pair with the reranker:

  * ``text_1`` str (or 1-list) × ``text_2`` list → one query against
    N documents (the common retrieval shape);
  * the symmetric N × 1 case;
  * equal-length lists → element-wise pairs.

The texts are concatenated VERBATIM — the client applies the model's
judgement template itself (cf. ``/v1/rerank`` which templates
server-side). ``data`` preserves the input pair order (no relevance
sort — index-addressable).

``/v1/loras`` admin surface: load / list / unload LoRA adapters.

Every body dispatches through the admin registry (``aload_lora`` /
``aunload_lora`` / ``list_loras`` — the ``lora_store`` is engine state);
the engine-side LoRA exception types cross as ``error_types`` names and
map back to the historical statuses here.

Admin: load a LoRA adapter from a PEFT directory.

Body: ``{"name": "<id>", "path": "<dir>"}``. Returns the loaded
adapter's metadata (extracted engine-side — the live adapter object
never crosses). 400 on validation failure (mismatched ranks,
missing target, unknown module). 409 if a LoRA with the same name
is already loaded.

Prometheus exposition endpoint, co-mounted alongside the OTLP
push pipeline.

The OTEL `PrometheusMetricReader` (see server/metrics.py) registers
every instrument with the global Prometheus client registry; this
route serves `generate_latest()` from that registry. Same metrics,
two consumers:

  - OTLP push to the OTEL collector for the ARBI/Grafana pipeline
  - Direct scrape here for legacy Prometheus deployments

Both stay in sync because the same `MeterProvider` feeds both
readers. Disable Prometheus at runtime with
`ARBI_SERVE_ENABLE_PROMETHEUS=0`; the route still registers but
will return an empty exposition.

AUTH-EXEMPT, AND THE EXEMPTION IS THE POINT (server/auth.py
``_AUTH_EXEMPT``). A Prometheus scrape carries no bearer token, and
this repo's own ``observability/prometheus.yml`` scrapes without one;
requiring auth here does not make deployments authenticate, it makes
their dashboards go blank. What made the exemption dangerous was not
the exemption, it was the AMPLIFICATION: one unauthenticated request
bought 202 ms of work from the process running the forward pass,
which is 500x what an equally-exempt ``/health`` costs. An attacker
who can reach the port can already flood ``/health``; what they must
not be able to do is buy a multiplier.

So the endpoint is bounded twice, and neither bound is a rate limit:

* per scrape, by the collection scope — one allocator walk and no
  ``smaps_rollup`` walk, rather than 129 and one
  (:mod:`arbi_serve.server.metrics._collection_scope`);
* across CONCURRENT scrapes, by :func:`_exposition` below — building
  the exposition off the HTTP loop gives the coalescing an await
  point to happen at, so a hundred simultaneous scrapers cost one
  build, not a hundred.

Rate is what is left, and rate alone no longer buys more engine time
than the same rate against any other exempt endpoint.

Build the exposition, off the HTTP loop, at most once per wave.

``generate_latest`` drives every observable-gauge callback in the
process — it is the collection, not a formatting step — so running it
inline held the HTTP loop for its whole duration, stalling streamed
responses to real clients behind a scrape.

Off the loop it also becomes COALESCIBLE, which is the part that
bounds the endpoint. A scrape notes the completed-build counter, then
waits for the lock; if a build finished while it waited, that build
started after this request arrived, so its result is not stale for
this caller and is returned instead of collecting again. N concurrent
scrapes therefore cost one collection. Nothing here is a cache: a
scrape that arrives when no build is running always gets its own.

Serve ``GET /v1/models`` (and the no-prefix ``/models`` alias).

Includes a ``capabilities`` list keyed off the instance's serving
task. Downstream consumers (e.g. ARBI's ``fetch_embedder_model_name``
/ ``fetch_reranker_model_name``) discover which served model to call
by scanning ``data[].capabilities`` for ``"embed"`` / ``"rerank"``
and reading that entry's ``id``.

``/v1/chat/completions`` (OpenAI-compatible chat, tool calling, SSE).

``_run_chat`` is the batch-worker core re-exported through the
``server.api`` aggregator (``batch_api`` imports it from there). Tool
parsing + multimodal preprocessing imports stay at module top / call-time
respectively; ``Engine`` is annotation-only so this module is engine-free
at import (``dump_openapi`` path).

Non-streaming chat-completion core, shared by the HTTP route and the
Batch API worker.

Submits one request through ``eng.asubmit`` (so the engine continuous-
batches it alongside everything else), awaits the result, and returns
the plain ``chat.completion`` dict. No ``FastRequest`` dependency — the
batch worker has no HTTP request object, so admission gating / auth
binding / per-request backend swap (all request-scoped concerns) stay in
the route wrapper below. Raises ``HTTPException`` on validation / timeout
failures; the batch worker turns those into per-line error rows.

``priority`` is stamped onto the SamplingParams (the Batch API worker
passes ``"batch"`` so its rows yield to interactive traffic; HTTP
callers keep the default ``"interactive"``).

Attach the trace under every wire name, or none if there is no trace.

Absent (not ``null``) when the response carries no reasoning at all —
thinking off, or a model with no reasoning markers — so a non-thinking
response never carries a reasoning field.

Did the model's own chat template leave generation inside a think block?

The single source of truth for whether a response carries a reasoning
trace. Read off the RENDERED prompt rather than ``enable_thinking``,
because the flag's meaning is per-checkpoint: Qwen3.6-27B treats
undefined as thinking-ON, Qwen3.5-0.8B treats undefined as
thinking-OFF, and Qwen3-4B-Thinking-2507 pre-fills ``<think>``
unconditionally and ignores the flag entirely. The rendered prompt is
the one artifact that already reflects whichever rule this model
follows.

Await chat generation to finish, extract any tool calls, and build the non-streaming ``chat.completion`` response (HTTP 408 on timeout).

``req_handle`` is the request's ClientRequest — the applier applies
every text delta BEFORE ``finish_reason`` becomes visible (in every
detok mode), so observing the finish means ``output_text`` is
complete; no detok drain exists or is needed.

Prebuild the static bytes around a content delta's JSON string.

The per-token content chunk is the SSE hot path (one frame per token
per stream). Its payload is STATIC for the stream's lifetime except
the content string itself, so instead of re-serializing the whole
``{**base, "choices": [...]}`` dict per token, splice
``orjson.dumps(text)`` between two prebuilt byte halves. orjson
escapes a top-level string exactly as it does
a nested value, and the template reproduces the dict's insertion-order
serialization — the wire bytes are IDENTICAL (asserted by
``tests/test_sse_frame_template.py`` over unicode / quotes / control
chars). Tool-call deltas keep the generic ``_sse_frame`` path.

Prebuild the static bytes around a reasoning delta's JSON string.

Same splice trick as :func:`_content_frame_template`, extended to the
two wire names in ``_REASONING_FIELDS``: the trace is serialized ONCE
and spliced into both slots, so carrying both names costs one extra
memcpy per frame rather than a second orjson pass.

Frame tail that closes a content/reasoning chunk with ``token_count``.

Spliced in place of the templates' static tail so the counted frames
stay on the same prebuilt-bytes path as the uncounted ones: a client
measuring per-token latency must not be paying a serialization cost
the served stream does not, or it measures its own instrumentation.

Frame a batch of staged tool-stream events into SSE byte frames.

``events`` are the kind-tagged dicts produced by
:meth:`arbi_serve.engine.tool_stream.ChatToolStream.drain` — a
``"content"`` event carries the cleaned content delta, a ``"tool"``
event carries an OpenAI tool-call delta. Framing here is identical to
the old on-loop path (which yielded a content chunk then one chunk per
tool delta, in FIFO order), so the wire bytes are unchanged.
``content_tpl`` is the per-stream prebuilt byte template from
:func:`_content_frame_template`; when absent (non-stream callers /
older call sites) the generic dict path runs — same bytes either way.

``token_count`` is the completion tokens the engine committed since
the previous batch (``stream_options.include_token_count``). One
batch is one engine wakeup, so the whole count belongs to the batch's
FIRST frame and the rest carry zero: a step's tokens are committed
together, and splitting them across the kinds the step's text happens
to span would invent an ordering the engine never had.

Yield ``chat.completion.chunk`` SSE deltas (including streamed tool calls) as the engine emits tokens, ending with ``[DONE]``.

``req_handle`` is the request's ClientRequest, fed by the
OutputApplier — which also DRIVES tool-call extraction, staging
cleaned content-deltas + tool-call deltas on
``req_handle.tool_stream`` (see :mod:`arbi_serve.engine.tool_stream`).
This generator only POPS the pre-computed deltas and frames them as
SSE. A request without a staged ``tool_stream`` (e.g. a test fixture
that bypassed the admission attach) falls back to an on-loop
extractor so the wire output is unchanged.

``/v1/completions`` (OpenAI-compatible text completions).

``_run_completion`` (batch-worker core) and ``_completion_stream`` are
re-exported through the ``server.api`` aggregator; ``batch_api`` and
``test_timeout`` import them from there. ``Engine`` is annotation-only so
this module stays engine-free at import (``dump_openapi`` path).

Every accepted ``prompt`` form, as a list of individual prompts.

OpenAI's completions ``prompt`` is a union: one string, N strings, one
pre-tokenized id list, or N pre-tokenized id lists. The last form is
what OpenAI-compatible scoring clients send — lm-eval posts
``[[id, ...]]`` even at batch size 1 — so rejecting it rejects the
whole loglikelihood surface.

Gate the logprobs surface: armed, within cap, and not streaming.

Rejected rather than clamped or ignored. ``extra="allow"`` on the
request model means an unserved knob would otherwise be absorbed
silently, and a client that asked for logprobs and got a ``null``
back cannot tell that from a model that had nothing to say.

Both scored lanes must cover every position they claim to.

The engine scores prompt positions on the eager prefill path and
sampled tokens on the synchronous sampler. A rung that returns
last-token logits only, or a deferred sample that materializes a tick
late, yields a SHORT lane — which would serialize as a logprobs array
that is well-formed and quietly missing positions. Clients slice these
arrays positionally (lm-eval sums ``token_logprobs[ctxlen:-1]``), so a
short lane silently shifts the window onto the wrong tokens. That is a
500, never a 200.

One position's ``{token: logprob}`` alternatives map.

The chosen token goes in FIRST and the top-k follow, so a chosen
token that is also the argmax collapses to one key — which is what
makes ``token_logprob == max(top.values())`` the standard greedy
test on the client side.

Assemble the OpenAI ``logprobs`` block from per-position entries.

``entries[i]`` is ``None`` where the position has no logprob — only
the first echoed prompt token, which nothing predicts. ``text_offset``
is the running character offset of each token within the returned
``text``.

``choices[].prompt_logprobs`` — one map per prompt position.

Keyed by token id (as a JSON object key) and valued by
``{logprob, rank, decoded_token}``. Element 0 is ``null``: prompt
position 0 has no preceding context, so no distribution predicts it.

Resolve ``echo`` / ``logprobs`` into ``(text_override, block)``.

``text_override`` is ``None`` when the response text is unchanged.
Returns ``(None, None)`` for the overwhelmingly common request that
asked for neither, before touching the tokenizer.

Non-streaming text-completion core, shared by the HTTP route and the
Batch API worker. Mirrors :func:`_run_chat`. Routes through
``eng.asubmit`` for continuous batching; returns the ``text_completion``
dict. Raises ``HTTPException`` on validation / timeout failure.

``priority`` is stamped onto the SamplingParams so the scheduler can
treat batch-origin work as opportunistic (the Batch API worker passes
``"batch"``); HTTP callers keep the default ``"interactive"``.

Multi-prompt ``/v1/completions`` (standard OpenAI: ``prompt: list[str]``).

Each prompt becomes its own engine request — so the scheduler
continuous-batches them and the radix prefix cache lands on any shared
preamble — with all prompts tokenized in ONE batched dispatch. Shared
sampling params across the batch (OpenAI semantics); one ``choices[]``
entry per prompt carrying its input index. Non-streaming only.

Await completion to finish and build the non-streaming ``text_completion`` response body (HTTP 408 on timeout).

``req_handle`` is the request's ClientRequest — the applier applies
every text delta BEFORE ``finish_reason`` becomes visible (in every
detok mode), so observing the finish means ``output_text`` is
complete; no detok drain exists or is needed.

``/v1/realtime`` — OpenAI Realtime API-compatible WebSocket endpoint.

The FIRST WebSocket route in the codebase. It decorates the shared
:class:`~fastapi.APIRouter` (same side-effect registration as the HTTP
routes; see ``routes/__init__.py``) so ``app.include_router(api_router)``
mounts it with everything else — no separate wiring in ``server/api.py``.

Torch-free at import (the ``dump_openapi`` discipline, landmine #5 in
``REFACTORING_1000LOC.md``): the realtime session / engine turn generator
are imported inside the handler, not at module top.

The handler is a thin adapter — accept the socket, build a
:class:`~arbi_serve.realtime.session.RealtimeSession` bound to
``websocket.send_json`` + the engine turn generator, emit
``session.created``, then pump client JSON events into the session until
the peer disconnects. All protocol logic lives in
:mod:`arbi_serve.realtime.session`.

Barge-in note: Step-Audio-2-mini is turn-based, so "barge-in" is
VAD-orchestrated (new user speech cancels the active response); the model
never attends to overlapping audio. See the session module docstring.

Publish the ``/v1/realtime`` WebSocket message shapes into OpenAPI.

OpenAPI cannot describe a WebSocket, so this dummy HTTP endpoint exists
only to drag every realtime message model into ``components.schemas`` —
the returned body is all-``null`` and carries no information. Clients read
the referenced schemas to learn what to *send* (``*_append``, ``commit``,
``response.create/cancel``, ``session.update``) and what format to
*expect* (``session.created``, ``response.audio.delta``, ``error``, …).

``/tokenize`` — the RESOLVED-PROMPT surface (vLLM-compatible).

Returns the token ids the engine would ACTUALLY generate from for a
given chat request — i.e. the fully-expanded chat template, thinking
prefix and all. This is the observability surface that makes
template/thinking drift a RED CELL instead of a silent accept deficit:
a bench harness posts the same payload to every engine, hashes
``tokens``, and asserts equality across same-basket cells.

Shape matches vLLM's ``POST /tokenize`` (``{count, max_model_len,
tokens}``) so ONE harness client covers both.

Render the chat template + tokenize, returning the resolved prompt token ids.

Thinking resolution is IDENTICAL to ``/v1/chat/completions`` — the
same :meth:`resolved_enable_thinking` merge over the top-level field
and ``chat_template_kwargs`` — so the hash of ``tokens`` is a true
fingerprint of what the chat route would have generated from, not a
parallel re-implementation that could drift from it.

Server-side receipt for every refused / shed / backpressured request.

A refusal that leaves no trace is unauditable: a run that quietly turned
away a fifth of its offered load looks identical, in the server log, to
one that served everything. Capacity claims are then unfalsifiable in the
direction that matters — you cannot tell "it coped" from "it refused".

So every refusal path calls :func:`record_shed` exactly once, and gets
both channels:

* a **log line** at WARNING carrying the reason verbatim (so
  ``grep -c <reason> <server log>`` is a real count) plus the governing
  threshold and the value that crossed it;
* a **counter** increment, :data:`SHED_COUNTER`, labelled by layer /
  reason / status.

Labels are the closed, low-cardinality vocabulary below. The numbers —
threshold and observed — go in the LOG, never in a label: they are
unbounded and would shard the time series.

The counter goes through :mod:`arbi_serve.cache._metrics`, whose
in-process registry works with no OTEL configured and no engine handle.
That matters here: refusals happen at the HTTP layer, which reaches the
OTEL bundle differently in in-process mode (``engine.metrics``) than in
process mode (``app.state.metrics``, and ``ProcEngine`` has no
``.metrics`` at all). One emit path avoids that divergence entirely.

Emit the server-side receipt for one refused request.

Call once per refusal, on the refusing path only — never on the
served path.

``log=False`` counts without logging. Reserved for a gate that
re-evaluates the SAME queued request every step (the per-tenant
scheduler filter): logging each re-evaluation would bury the log,
which is its own observability failure. Those call sites log the
edge — the first refusal of a (tenant, reason) — and count every
one, so the rate is still exact.

Parameters
----------
layer
    Which gate refused (:data:`LAYERS`).
reason
    The refusal reason (:data:`REASONS`). Logged verbatim so it is
    greppable, and used as a metric label.
status
    The HTTP status the client will see (429 / 503 / 403).
limit
    The governing threshold — the number the operator can change to
    stop this refusal. ``None`` when the refusal is not threshold-driven
    (a drain, a latched fault).
observed
    The value that crossed ``limit``.
unit
    Suffix for both numbers in the log line (e.g. ``"%"``, ``" req"``).
detail
    Extra context for the log line only. Never a label.

Lock-free ASGI SSE streaming response.

A minimal drop-in replacement for Starlette's ``StreamingResponse`` on
the server-sent-events hot path. It pumps the response body
async-generator **directly** into the ASGI ``send`` callable — no
``anyio`` memory-object-stream, no task group, no threadpool.

Why
---
Under uvicorn, the ASGI HTTP protocols advertise ``spec_version
"2.3"`` (``uvicorn/protocols/http/{httptools,h11}_impl.py``), which is
``< (2, 4)``. Starlette's ``StreamingResponse.__call__`` therefore
takes the *legacy* branch: it spins up an ``anyio.create_task_group()``
and races ``stream_response`` against a ``listen_for_disconnect``
poller. Under concurrent SSE streaming, that per-call task-group +
per-chunk anyio plumbing holds the GIL on the HTTP thread and starves
the engine thread.

This class skips all of it. ``__call__`` writes exactly the ASGI
messages Starlette's ``stream_response`` would have written, in the
same order, with the same bytes — only without the surrounding
task-group machinery.

Wire-format equivalence
------------------------
For ``media_type="text/event-stream"`` with no explicit headers,
Starlette emits::

    http.response.start  status=200  headers=[(b"content-type",
                                                b"text/event-stream; charset=utf-8")]
    http.response.body   body=<chunk>  more_body=True   (per chunk)
    http.response.body   body=b""      more_body=False  (terminator)

We reproduce that byte-for-byte. ``str`` chunks are encoded with the
same ``charset`` ("utf-8") Starlette uses, so a mixed str/bytes
generator yields identical wire bytes.

Disconnect / cancellation
--------------------------
This response does **not** own cancellation. The route handlers wrap
the token generator in
:func:`arbi_serve.server.routes._request_ctx._stream_with_cancellation`,
which itself nests :func:`arbi_serve.server.cancellation.safe_run`.
``safe_run`` runs ``Engine.cancel`` **only** on
:class:`asyncio.CancelledError` — the exact signal Starlette's legacy
task-group path delivers when it cancels ``stream_response`` on a
disconnect. Two disconnect shapes both reach that cleanup:

  * **task cancellation** (server shutdown / Starlette tears the
    request task down): the ``CancelledError`` is delivered into the
    producer's suspended ``await`` *inside* ``safe_run``. ``safe_run``
    runs ``Engine.cancel`` under ``asyncio.shield``, re-raises
    :class:`SafeRunException`, and the wrapper swallows it — the
    generator finishes, the ``async for`` here ends cleanly, and the
    stream terminates with the engine request released. No chunk is
    dropped; no request leaks.
  * **``send`` raises** (the transport is gone before the producer
    notices — uvicorn surfaces a dead socket as ``OSError``): the
    producer is parked at its own ``await``, so we must inject the
    signal. ``_cancel_body`` calls ``athrow(CancelledError)`` (NOT
    ``aclose()`` — that throws ``GeneratorExit``, which ``safe_run``
    does not catch, leaking the request) to run the same cleanup, then
    the original transport error is re-raised to the ASGI server.

No disconnect poller is used here (intentionally): ``safe_run``
alone is sufficient for correctness on a live generation.

Drive a suspended SSE producer through its cancellation path.

The producers fed to :class:`SSEResponse` are wrapped in
:func:`arbi_serve.server.routes._request_ctx._stream_with_cancellation`,
whose :func:`~arbi_serve.server.cancellation.safe_run` guard runs
``Engine.cancel`` **only** on :class:`asyncio.CancelledError` — the
exact signal Starlette's legacy task-group path delivers when it
cancels ``stream_response`` on a client disconnect.

A bare ``aclose()`` would throw ``GeneratorExit`` instead, which
``safe_run`` does not catch → the engine request would leak. So we
``athrow(CancelledError)`` to reproduce the legacy signal. The
wrapper catches it (cleanup runs), swallows :class:`SafeRunException`,
and returns — which surfaces here as ``StopAsyncIteration`` (the
generator finished). We suppress every outcome: this runs while an
original exception is already unwinding, and cleanup is best-effort.

Lock-free ASGI streaming response for server-sent events.

Constructed with an async iterable of ``bytes`` (or ``str``) SSE
frames and an optional status / header set. It is itself an ASGI
application (``async def __call__(self, scope, receive, send)``),
so route handlers ``return SSEResponse(generator)`` as a drop-in
for ``StreamingResponse(generator, media_type="text/event-stream")``.

Parameters
----------
content:
    An async iterable yielding ``bytes`` (preferred; already-encoded
    SSE frames) or ``str`` frames. Sync iterables are **not**
    accepted — every SSE producer in this codebase is an async
    generator, and accepting a sync one would re-introduce the
    threadpool hop this class exists to remove.
status_code:
    HTTP status for the response start message. Defaults to 200,
    matching ``StreamingResponse``.
headers:
    Raw ASGI header pairs ``(name, value)`` as ``bytes``. Defaults
    to the single ``text/event-stream; charset=utf-8`` content-type
    header Starlette would set for this media type.

Stream the body generator straight into ``send``.

Writes the ASGI ``http.response.start`` message, then one
``http.response.body`` message per generator chunk
(``more_body=True``), and a final empty body
(``more_body=False``) — the exact sequence Starlette's
``stream_response`` produces, minus the anyio task group.

On a ``send`` failure (client transport gone) or a cancellation
of this task mid-stream, the body generator is driven through
its cancellation path so the ``safe_run``-wrapped producer runs
its ``Engine.cancel`` cleanup (see ``_cancel_body`` for why this
throws ``CancelledError`` rather than ``aclose``-ing).

OTEL telemetry install + resource-attribute helpers for the app factory.

Covers the OTEL ``MeterProvider`` / ``TracerProvider`` configuration,
FastAPI auto-instrumentation, and the resource-attribute bundle the
engine's ``Metrics()`` resolves instruments against. Imported back
into ``server/app.py`` so the public surface is unchanged.

Extract a stable model identifier from a checkpoint path.

Falls back to the last non-empty path component so resource
attributes always carry SOMETHING readable (rather than the full
absolute path) when ``--served-name`` isn't set.

Stamp the inbound HTTP span with the *true* client address.

Behind haproxy (``option forwardfor``, see ``swarm/haproxy.cfg``)
the transport peer is the proxy, so the span's auto-derived
``client.address`` is haproxy's IP, not the caller's. The original
client rides in ``X-Forwarded-For`` (left-most entry is the
furthest-upstream client). Overwrite ``client.address`` with it so
Tempo/Grafana can group and search requests by real origin.

Runs for EVERY inbound request — including unauthenticated and
auth-exempt paths (``/health``, ``/metrics``, ``/v1/models``) that
never reach :class:`AuthMiddleware` — so origin attribution exists
with or without a bearer token. Best-effort: never raises into the
request path.

Build the OTEL resource-attribute bundle from the server config.

Surfaces model identity (name / path / dtype / backend) plus the
engine size knobs as resource attributes so every metric / log /
span carries them as labels. ``max_batch`` / ``max_context`` may
still be the ``"auto"`` sentinel here — they emit ``-1`` and get
re-emitted as metrics once the engine resolves them at build().

Install OTEL providers + auto-instrumentation BEFORE the engine builds.

The Metrics() bundle the engine constructs resolves instruments
against the active MeterProvider at construction time, and
FastAPIInstrumentor needs a TracerProvider in place when it wires the
ASGI middleware — so this must run first on the lifespan startup path.

ThreadPoolExecutor-backed tokenizer pool.

Tokenization is CPU work that, in the Python case, holds the GIL for
milliseconds at a time on long prompts. Doing it on the asyncio event
loop blocks every other coroutine (including the engine's run_loop
tick, the disconnect poller, and other in-flight SSE handlers).

The fix: hand tokenization to a small dedicated thread pool. Every
worker thread releases the GIL while inside the Rust tokenizer's
``encode``/``apply_chat_template`` (the ``tokenizers`` library is Rust
with explicit GIL release on the hot path), so the event loop stays
responsive.

We round-robin across N workers so a slow tokenize doesn't block a
fast one queued behind it; queue depth is controlled by the upstream
:class:`RequestSemaphore`.

The module exposes a small async-friendly facade
(:class:`TokenizerPool`) so handlers don't have to touch
:func:`asyncio.get_event_loop().run_in_executor` directly.

Digest of the last :data:`_PFX_TAIL_CHARS` characters of ``text[:n]``.

O(1) in the length of ``text``, which is what makes scanning every
candidate affordable: the full-prefix digest is 0.5 ms over a 708 KB
prompt, and paying that per candidate would cost more than the encode it
is trying to avoid.

``f(texts, add_special_tokens) -> list[list[int]]`` for either batch surface, or ``None``.

``tokenizers.Tokenizer`` (the Rust binding) has ``encode_batch``; a
``transformers`` fast tokenizer batches through ``__call__`` and hands back
``{"input_ids": [...]}``. Both release the GIL for the bulk work; neither is
reachable through the other's name.

``(k', tail)`` where ``decode(ids[-k':])`` is a genuine SUFFIX of ``text``.

Widens ``k`` one token at a time (a codepoint spans at most four bytes, so
a clean boundary is a few tokens away at most) and returns ``None`` when no
window in the ladder round-trips — the caller then declines to a full
encode rather than splicing text it cannot prove it holds.

Async facade over a :class:`ThreadPoolExecutor` for tokenization.

Construct once per server (typically inside
:func:`build_app`'s lifespan startup); call :meth:`tokenize`,
:meth:`encode`, or :meth:`apply_chat_template` from the request
handlers.

The wrapped :class:`Tokenizer` instance is shared across workers.
The Rust side is internally thread-safe for ``encode`` / ``decode``
(each call constructs its own state); chat-template rendering goes
through Jinja which is also thread-safe at the rendering level.
Two callers can encode in parallel without coordinating.

The pool is closed via :meth:`shutdown` (or async-with). Pending
futures are NOT cancelled — Python doesn't support thread
interrupt — but new submissions raise :class:`RuntimeError`.

Wrap ``tokenizer`` in a thread-pool executor with ``num_workers`` threads for off-loop tokenization.

``prefix_cache_entries`` is how many CONVERSATIONS stay warm at once,
and ``prefix_cache_max_ids`` the total token ids the cache may hold
across all of them. Both are the caller's to resolve, because both are
facts about the deployment rather than about this class: the first is
how many requests the server admits concurrently, the second follows
from that and the context it admits them at
(:func:`arbi_serve.server.engine_boot.resolve_prefix_cache_bounds`).
The defaults are the single-entry cache this class shipped with, so a
caller that resolves neither is byte-identical to that.

``prefix_cache_max_ids`` of ``0`` bounds the cache by entry count
alone. That is finite but not small — one entry can hold a whole
max-context prompt — so a caller serving long contexts should pass it.

``(key, prev_len, prev_ids)`` for the longest cached prefix of ``text``.

Candidates are examined longest-first, so the reuse is always the
biggest one available. Each is rejected on its tail signature — O(1) in
the prompt's length — and only the survivor pays the full-prefix
digest, which is the proof. ``None`` when nothing matches.

The LRU order is mutated on a hit, so it is shared state that a READER
writes — but only the snapshot and the reorder are taken under the
lock, never the digests.

Store ``text``'s ids, dropping the entry it extends, then evict to the bounds.

``supersedes`` is the key the shortcut just reused: a shorter prefix of
the SAME conversation, made redundant by this longer one. Dropping it
is what keeps one conversation from filling the cache with its own
history one turn at a time.

An entry larger than the whole id budget is stored and then immediately
evicted by the loop, leaving the cache empty rather than over budget:
the budget is the promise, and a single prompt that cannot fit under it
is one this cache cannot serve.

Swap the wrapped tokenizer in place (e.g. after a model hot-swap).

The pool is built ONCE at app startup around ``eng.tokenizer``; a
cross-architecture model swap replaces ``eng.tokenizer`` with a
DIFFERENT vocabulary, so the pool must be repointed or the chat /
completion routes keep tokenizing with the old model's tokenizer.
That is not cosmetic: the old tokenizer's special-token ids can sit
ABOVE the new model's vocab size, so the chat-template tokens index
the new model's embedding out of range → CUDA device-side assert.

Workers read ``self._tokenizer`` per call and the underlying Rust
tokenizer's ``encode`` is thread-safe, so a plain rebind is safe.
Callers invoke this AFTER the engine swap has drained in-flight
requests, so no encode is racing the rebind.

Install the resolved-key → tokenizer provider for a residency pool.

Wired at app startup when stable-VA residency may route to a non-active
resident. Clears the per-key cache so a re-wire never serves a stale
member tokenizer.

The tokenizer to encode a request routed to ``resident_key`` with.

The active ``self._tokenizer`` when no key is pinned or no provider is
installed (the single-model path). Otherwise the target resident's own
tokenizer (cached per key), falling back to the active one when the
provider cannot resolve the key — a fixed key always maps to the same
member, so the cache never goes stale.

Tokenize ``text`` off the event loop.

Returns the token-id list. Same shape as
:meth:`Tokenizer.encode` directly. ``resident_key`` pins the encode to
a specific stable-VA resident's tokenizer (a deferred cross-arch swap);
``None`` uses the active tokenizer.

Blocking encode body run on a worker thread.

Routes through the Rust ``encode_batch`` even for a single text:
the single-text ``encode`` binding does NOT release the GIL,
while ``encode_batch`` releases it for the Rust bulk. Ids are
identical (asserted in ``tests/test_server_tokenizer_pool.py``).
Falls back to single ``encode`` for tokenizer stubs without the
batch surface.

Encode ``text`` by splitting it, encoding the parts together, and repairing the seams.

``encode_batch`` releases the GIL and parallelises across the parts, but
a naive concatenation is WRONG: a BPE merge can span a split, so the
joined ids differ from a single encode. Every seam is therefore
re-encoded from a window of tokens either side of it and spliced back,
which restores the exact single-encode ids — verified to the id in
``tests/test_server_tokenizer_pool.py``.

The seams are what keeps the parallelism worth having: repairing a
small window per split leaves the bulk work in the batch, where a first
attempt that re-encoded a whole chunk per seam serialised the work back
and measured no faster than a single encode at all. For the same
reason every seam is repaired in ONE batch and spliced once: a dispatch
per seam puts the repair back in series with the batch it exists to
feed, and it is that serialisation — not the seam count — that caps how
far the text can usefully be split.

Returns ``None`` when the shortcut cannot be taken (no batch surface, no
decode, too short to split) so the caller falls through to one encode.

``_raw(text)``, but reusing the cached prefix ``hit`` found for ``text``.

``hit`` has already been PROVEN to be a prefix of ``text`` — equal
length plus equal digest is the proof ``startswith`` used to give, at
0.5 ms over 708 KB against a 33-159 ms encode. What is left here is
whether the shortcut can be taken exactly: the entry must be strictly
shorter, and a backoff window must round-trip. Anything else falls
through to a full encode, and the result is always the same id list a
full encode returns — the shortcut changes only how long it takes to
get there, never what comes back.

Tokenize a whole batch off the event loop in ONE dispatch.

Embedding / rerank requests fan out to one engine request per
input; tokenizing each via :meth:`encode` would fire one
thread-pool submit per input (lock contention on the work queue
at high batch sizes). This routes the entire batch through a
single worker thread, where the Rust ``encode_batch`` parallelises
internally. Returns one id list per input, in order. ``resident_key``
pins the encode to a specific resident's tokenizer.

Render the model's chat template off the event loop.

Tools / thinking flags and any further ``template_vars`` pass
through to :class:`Tokenizer` — including its refusal of a variable
the resident's template cannot honor, which surfaces as the
``ValueError`` the routes map to 400. ``resident_key`` renders with a
specific resident's template + tokenizer (a deferred cross-arch swap
ships a different chat template too).

WHICH ``arbi_serve`` source tree a process is actually running.

Two processes in one deployment run this package: the engine, and the
admin console beside it. They are one repository and the console renders
the engine's own vocabulary -- route names, flag names, field names -- so
a console built from a different tree than the engine does not degrade,
it misreports: a panel reads a field the engine no longer sends and
renders the absence as a value.

Nothing compared the two before, and the failure is silent in both
directions. A console whose source is behind the engine's shows an older
UI against a current server; one ahead shows controls the server has no
route for. Neither logs anything.

A git sha cannot be the identity on its own. The console's container gets
the package directory and not the ``.git`` beside it, so ``git rev-parse``
answers nothing there -- the very deployment that needs the check is the
one a sha cannot describe. What both processes can always compute is the
CONTENT of the tree they imported, so that is the identity: a digest over
every source file under the package root, by path and by bytes.

The digest is over what a reader would call source -- Python, and the
admin UI's own served assets, since a stale ``app.js`` is exactly the
shape of tonight's failure. Byte-caches are excluded: they are derived,
and two processes can hold different ``__pycache__`` for identical
sources.

The directory of the ``arbi_serve`` package THIS process imported.

Derived from this module's own file rather than from a configured path:
the question is which code is running, and the only witness to that is
the code that is running.

The content identity of one ``arbi_serve`` tree.

``files`` maps each source file's path, relative to the package root,
to a short digest of its bytes. Per FILE and not one digest over the
whole tree, because the two processes being compared do not hold
identical file SETS even when they hold identical code: one runs an
installed package and the other a source tree, and packaging decides
which non-code files come along. A single tree-wide hash cannot tell
"different code" from "different packaging", and refusing a boot over
the second would be a guard that cries wolf until it is turned off.

``digest`` folds the same map, so an equality check is one comparison
when the two are identical -- which is the common case, and the one
worth being cheap.

A file that cannot be read is recorded with its error rather than
skipped: skipping lets an unreadable tree compare equal to a readable
one, which is the check reporting a pass it did not earn.

How two identities differ, split by the only distinction that matters.

``differing`` holds files BOTH trees have and whose bytes disagree --
the same code, two versions of it, which is a genuine mismatch. The
``only_*`` lists hold files one tree has and the other does not, which
an installed-vs-source packaging difference produces routinely and
which is therefore reported but never a refusal.

One line an operator can act on: the digest, the tree, and its age.

The mtime is rendered as a local timestamp because "is this the tree I
edited an hour ago" is the question being asked of it, and an epoch
float does not answer that at a glance.

Speculative decoding — MTP draft heads + verify+accept loop.

The sampler / engine surfaces use a constant K-dim shape so they stay
exercised even when no :class:`MtpDriver` is attached.

:mod:`arbi_serve.spec_decode.mtp` carries the bundled-head MTP driver,
verify-batch builder, verify-pass driver, and the
:class:`MtpStrategy` engine binding all in one place.

Shared left-to-right biased block-realization walk for DFlash re-rank heads.

Both DFlash re-rank heads — the DSpark ``VanillaMarkov`` head and the DFlash2
``CandidateSelector`` — realize a drafted block LEFT-TO-RIGHT, conditioning
each position on its realized predecessor. They differ only in how a single
position's base logit row is biased; the loop around that step — the seeded,
graph-safe categorical draw, the keyed-watermark coupling, and the exact
per-position ``q`` accumulation that keeps the draft lossless — is identical.
:func:`biased_sample_block` is that shared loop; each head supplies its own
one-position ``step_logits`` closure.

Land position ``j``'s per-row drafter ``q`` in the block destination.

The one place a block realization writes its ``q``, shared by the
left-to-right walk below and the plain parallel block. ``rows`` addresses
the destination's columns when the block realizes only a SUBSET of the
slate's rows (the inactive rows keep whatever the caller pre-filled);
``None`` means the destination is exactly this block's rows.

LEFT-TO-RIGHT biased SAMPLED block → ``(tokens (Bd, S), q_out)``.

``q_out`` is the caller's destination (see :func:`write_block_q`), not a
tensor this walk allocates: the block's ``q`` is vocab-scale fp32 on the
per-token draft path, so the caller sources it from
:func:`~arbi_serve.spec_decode.drafter.drafter_q_out` — the engine's
persistent buffer where there is one.

``step_logits(j, base_j, prev)`` returns position ``j``'s biased
full-vocab logit row given the realized predecessor ``prev``
(``first_prev`` at ``j == 0``). Each row is drawn through the drafter's
per-row chain — temperature → top_k → softmax → top_p → min_p — so the
returned ``q_j`` is EXACTLY the distribution position ``j`` sampled from,
the identity standard rejection sampling needs for losslessness. ``slot_offset=j`` decorrelates
the S positions' draws from the one shared step seed; the sampled token is
the feedback for position ``j+1``. When ``wm_ctx`` is armed each position
draws with the per-row keyed seeds the watermark verify uses at depth
``j``, coupling the realized block to the target's keyed choices — the
left-to-right realization is what makes a re-rank head couplable at all.

``tensors`` supplies the slate's drafter sampling tensors already encoded
(the captured realization keeps them in persistent buffers and refreshes
them by ``copy_()``); the walk then reads no Python ``SamplingParams`` at
all. ``toks_out`` is the token destination for a caller that owns one.

Independent (parallel) SAMPLED block → ``(tokens (Bd, S), q_out)``.

The plain-DFlash sibling of :func:`biased_sample_block`: every position
draws from its own base row through the same per-row drafter chain
(temperature → top_k → softmax → top_p → min_p) under a distinct
counter-keyed draw (``slot_offset=j``), and ``q_j`` is exactly the
distribution position ``j`` sampled from. There is no realized prefix, so
no position conditions on another and no keyed-watermark coupling
applies.

The slate's drafter sampling tensors are built once for the block (or
supplied pre-encoded as ``tensors``) and reused across the ``S``
positions, on the same routing conditions as the biased walk; the
fallback is the per-position ``sample_drafter_token`` chain.

Fixed-shape context-assembly cluster for :class:`DFlashDrafter`.

:class:`_DFlashContextMixin` holds the capture-ready draft-forward core:
the host + device fixed-shape context gather and the eager/replay draft
forward dispatch. The methods reference ``self.<attr>`` bound in
:meth:`DFlashDrafter.__init__`.

``active`` as maximal ``(dest row, slab row, count)`` runs.

A run is a stretch over which the destination row and the source slot both
advance by one, so its rows are one contiguous block of the slab copied
into one contiguous block of the destination — a single strided ``copy_``
rather than one per row. Rows that break the pattern simply start a new
run, so the worst case is the per-row blit and the steady all-active slate
is a single copy.

``ARBI_DFLASH_GRAPH_ASSEMBLE`` — assemble into the graph's own buffers.

Reads the value :meth:`freeze_graph_assemble` pinned, so the route a
step takes is the route the serving floor was sized against — by
construction, not by agreeing with a second reader. Before the freeze
(and wherever no freeze happens) it falls back to the live flag: both
settings assemble the same context and replay the same graph, they
differ only in whether the context lands in the graph's buffers
directly or is copied into them.

Pin the assemble route, and report whether the fallback is DEAD.

Called once, at the post-capture seam that sizes the serving floor —
after the draft-graph pre-sweep has run and sealed the pool, so the
pool's coverage is a fact rather than a forecast.

Returns True only when every one of :meth:`_select_draft_graph`'s
``None`` conditions is refuted for every reachable step:

* *the lever is off* — refuted by the pin: the value read here is the
  value every later step reads, and
  ``ARBI_DFLASH_GRAPH_ASSEMBLE`` is capture-affecting, so an override
  of it builds a new member (a new drafter, re-frozen) rather than
  swapping the route under a floor sized for the other one;
* *capture is not armed* — refuted by ``_draft_graphs is not None``,
  and ``ARBI_DFLASH_CAPTURE`` is capture-affecting for the same
  reason;
* *no graph covers the shape* — refuted by
  :meth:`~arbi_serve.spec_decode.dflash_capture.DFlashDraftGraphPool.covers_every_shape`
  over the pre-sweep's own shape list, which is coverage at the SLAB
  caps for every batch width ``1..max_batch`` and therefore coverage
  at every narrower live shape;
* *the step has no active row* — refuted at both call sites: the
  driver enters the draft only under a non-empty ``done``, which is
  ``active``'s own row list, and the TP-shard worker bridge guards on
  ``if active:``. The empty slate never reaches here.

A False leaves the floor reserving the per-step destination, which is
the shipped behaviour.

Both assemblies address ``(B, ...)`` buffers by the ``active``
tuple's row field, so it must be a batch row in ``[0, B)``.

Rank 0 derives it from ``enumerate(requests)``; the TP worker
shard-draft rebuilds it from the broadcast ``active_b`` against its
own ``blk_ids`` width. An out-of-range row is an out-of-bounds device
scatter — an illegal memory access, not an exception — so it is
rejected here, on the host, for every caller.

Every active row's slab layout, per draft layer.

``[layer][row] -> (dest row, slot, start, m, runs)`` from
:meth:`DraftKVSlots.context_layout` — pure host arithmetic off
``lengths``, so the whole assembly's geometry (widths included) is
known before a single slab byte is read.

The per-layer destination width the host assembly fills: the widest
held context over the active rows, floored at 1.

THE definition — :meth:`_assemble_fixed_context` narrows its
destinations with this exact function, so a caller that needs the
widths BEFORE the assembly (to pick the capture graph it will assemble
into) gets the widths the assembly will actually produce, by
construction rather than by agreement.

Copy active slots' per-layer context into FIXED-SHAPE padded
buffers for :meth:`DFlashDraftModel.forward_block_fixed`.

Returns ``(start_pos, ctx_k, ctx_v, ctx_kpos)`` where ``start_pos``
is ``(B,)`` (0 for inactive rows), and each ``ctx_*[i]`` is padded
over B rows to the layer's longest held context ``Ccap_i``. Pad
columns carry ``kpos = -1`` so ``forward_block_fixed`` masks them
to zero softmax weight — inactive rows and short rows produce a
result the caller discards / that is context-free.

Each row is blitted STRAIGHT from the slab into its destination
slice: :meth:`DraftKVSlots.context_layout` resolves the row's held
context into the one or two contiguous slab runs that back it (a
ring window of ``m <= ring_len`` positions wraps at most once), and
the widths it reports size the destinations up front. No gather
temporary is ever co-resident with the buffers it feeds.

The destinations are allocated FRESH per step and die with the
caller's references. That lifetime is load-bearing, not incidental:
the post-capture grow folds
:func:`~arbi_serve.engine.memory_budget.dflash_draft_transient_peak_bytes`
against the verify tail with a MAX on this synchronous path, which is
sound ONLY because the draft runs after the verify sync and its
destinations do not outlive the step. Retaining them across steps
makes them co-resident with every other serving transient, and the
MAX then under-reserves by exactly the retained bytes.

``into`` redirects the assembly at a caller-owned per-layer
``(k, v, kpos)`` destination set instead of allocating one — the
covering capture graph's own input buffers, so the context is written
once, where the replay reads it, rather than here and again into the
graph. That does NOT retain anything the step would otherwise free:
those buffers are the capture pool's, allocated at capture and already
resident when the grow measures free VRAM, and the replay already
rewrote them every step. The step allocates strictly LESS, so the
transient bound above still covers it.

The buffers are at least this step's width, so a row's runs land at
the same offsets either way and the pad beyond the live width — and
any row no active slot maps to — keeps whatever an earlier step left
there, masked out by the ``-1`` key positions refreshed over the whole
width, exactly as :meth:`DFlashDraftGraph.replay` leaves it.

``layouts`` accepts the row layouts a caller already resolved (to pick
that graph); omitted, they are resolved here.

A sync-free upper bound on the batch's max committed context length.

Reads ``engine.page_table.length(req)`` (host) for every active row and
adds a ``block_size`` guard — provably ≥ the device ``lengths_dev`` for
the row (the page table lags the device commit by at most one pending
stash of ≤ block_size tokens). Returns ``None`` when the engine has no
page table (CPU unit tests) so the caller falls back to the full slab
width (correctness-safe, just wider). No ``.item()`` / no D2H — this is
what keeps the device-slots draft off the per-step host-sync.

DEVICE-resident context assembly: same output contract as
:meth:`_assemble_fixed_context` but the committed lengths and the key
positions come from ``lengths_dev`` (no host ``.item()`` / no host ring
loop), so the draft forward runs off the async verify outputs. Context
width is the full per-layer slab (``layer_len``);
``forward_block_fixed`` masks the unfilled columns via ``kpos``.
Full-attention layers are WIDTH-BOUNDED to a sync-free host upper
bound on the live committed length (``_host_width_bound`` —
page-table length + block guard, no ``.item()``; sliding layers ring
to their small window regardless), so the device path stays off the
per-step CUDA host-sync instead of attending over the whole
``effective_context`` slab.

Each row's slab slice lands STRAIGHT in its destination row, with no
gathered ``(A, nkv, cap, hd)`` copy co-resident with the destination
it feeds: each byte moves once, not twice. Rows whose destination and
slab index advance together are blitted as ONE strided ``copy_``
(:func:`_row_slot_runs`), which the steady all-active slate collapses
to a single copy per layer per tensor. The run boundaries are host
ints carried by ``active``, so the assembly needs no ``.item()`` and
issues no device read — the sync ban is on device reads, not on host
iteration. The destinations too are allocated FRESH per step — the
lifetime :meth:`_assemble_fixed_context` documents.

``into`` redirects the assembly at a caller-owned per-layer
``(k, v, kpos)`` destination set instead of allocating one — the
covering capture graph's own input buffers, already resident before
the grow and already rewritten by every replay, so the step allocates
strictly less and retains nothing new. Each layer takes the leading
``[:, :, :cap]`` slice; the tail beyond it, and any row no active slot
maps to, keep whatever an earlier step left there and are masked out
by the ``-1`` key positions refreshed over the whole width — exactly
the pad :meth:`DFlashDraftGraph.replay` leaves behind.

Denoise all rows → ``(B, block, H)`` post-norm hidden.

Eager by default. When ``ARBI_DFLASH_CAPTURE`` is set and a
cudagraph exists for this (B, Ccap-bucket) shape, the captured
:class:`DFlashDraftGraph` is replayed instead of re-tracing the
5-layer forward every step. The graph is greedy-bit-identical:
same ``forward_block_fixed`` math over the same persistent
buffers. A miss (new shape) falls through to eager
— never a deadlock (the draft forward carries no cross-rank
collective on the replicated build; the sharded build's all_reduce
is captured symmetrically on every rank).

The capture graph this step will assemble INTO, picked from slot
state alone; ``(None, layouts)`` ⇒ allocate a fresh per-step
destination and dispatch as usual.

The graph pool keys on ``(B, caps)``, and ``caps`` is a pure function
of the per-layer context WIDTHS
(:meth:`DFlashDraftGraphPool.caps_for_widths`). Both assemblies derive
those widths from :class:`DraftKVSlots` host arithmetic —
:meth:`_host_context_widths` off the row layouts, or
:meth:`DraftKVSlots.context_width` on the device path — so the widths
are known, and the graph is therefore known, before any slab byte is
read. The host layouts are resolved once here and handed to the
assembly rather than recomputed.

``None`` when the lever is off, when capture is not armed, when no
graph covers the shape, or when the step has no active row (the
assemblies' empty-batch shapes are their own; there is nothing to save
by routing them through a graph).

Assemble this step's fixed-shape context and denoise it → ``(B,
block, H)`` post-norm hidden.

``borrow_output`` returns the captured graph's persistent output
buffer instead of a per-step clone of it. The borrower owes
consumption before the next draft forward — assert
:attr:`_draft_forwards` is unchanged while it holds the result. A step
that runs eager, or that falls back to a fresh destination, returns a
tensor of its own either way, so the flag is a permission and never a
promise: the result is always at least as long-lived as a borrow.

The single draft-forward entry point. When a captured graph covers the
step's shape and ``ARBI_DFLASH_GRAPH_ASSEMBLE`` is armed, the context
is assembled STRAIGHT into that graph's own input buffers and replayed
in place: those buffers are already exact-width, already contiguous
and already the addresses the recorded kernels read, so the copy
:meth:`DFlashDraftGraph.replay` would otherwise make from a separate
destination is pure duplicate traffic. Otherwise the step allocates a
fresh destination and dispatches through :meth:`_run_draft_forward`
(replay-with-copy, or eager).

This SIZES the draft-step reserve. The direct route ALLOCATES LESS —
it skips the per-step destination entirely — and RETAINS nothing new:
the graph's buffers belong to the capture pool, exist before the
post-capture grow measures free VRAM, and were already rewritten by
every replay. When :meth:`freeze_graph_assemble` proved the direct
route is the ONLY route this member can take,
:func:`~arbi_serve.engine.memory_budget.dflash_draft_transient_peak_bytes`
stops bounding a destination — so the fallback below is refused rather
than run, because there is no longer free VRAM standing behind it.
Without that proof the bound covers the fallback and over-covers the
direct route, exactly as before.

Both routes write the same context and replay the same graph, so the
drafted hidden is bit-identical: the columns and rows they leave
untouched differ only in which earlier step's draft K/V they hold, and
every one of those carries a refreshed ``kpos = -1`` that
``forward_block_fixed`` masks to exactly zero softmax weight.

DFlash2 candidate-selector block realization.

:class:`_DFlashDflash2Mixin` holds the DFlash2-checkpoint method of
:class:`DFlashDrafter`: the candidate-path GREEDY block realization that
re-ranks the parallel block through the checkpoint's
:class:`~arbi_serve.spec_decode.dflash2_layers.CandidateSelector`. The method
references ``self.model`` bound in :meth:`DFlashDrafter.__init__`.

The selector is the DFlash2 sibling of the DSpark semi-AR markov re-ranker
(:mod:`arbi_serve.spec_decode._dflash_driver_dspark`): both walk the block
LEFT-TO-RIGHT, conditioning each position on its realized predecessor. Where
the markov head adds a low-rank bigram logit bias, the selector re-scores the
per-position top-k candidate pool with a bilinear predecessor→successor term
over low-rank codebooks and a hidden projection, then takes the argmax
candidate. A checkpoint carries one head or the other, never both.

The block's per-position top-k pool: ``(unary, candidates)``.

``base`` is ``(rows, S, V)``, so the whole block is selected in ONE call.
The pool depends only on the base logits, never on the realized
predecessor, so neither realizer may sink this into its walk.

Runs on ``base``'s OWN dtype, before any fp32 cast a caller needs
downstream: CUDA top-k is a radix select and runs one digit pass per key
byte, so a fp32 row pays 4 passes where bf16 pays 2 — 209 us against 153 us
at ``V = 248320``, ``rows = 64`` on a 4090. Casting before the selection
buys nothing and costs that 36%.

``sorted=False`` because every consumer (argmax + gather on the greedy
walk, ``scatter_add_`` on the sampled one) is order-free. It is worth about
4.5 us of the above, not more — the cost is the select, not the sort.

LEFT-TO-RIGHT candidate-path GREEDY block: ``(Bd, S)`` tokens.

For each block position, score the ``selector_top_k`` argmax
candidates against the realized predecessor — the candidate's own
logit (``unary``) plus the bilinear term
``(pred_codebook[prev] * hidden_projection(hidden)) · succ_codebook[cand]``
— take the argmax candidate, and feed it forward as the next
position's predecessor (``first_prev`` at position 0). Greedy, so
lossless under greedy verify: the target re-checks every drafted
token against its own argmax regardless of how the draft was formed.

Bit-identical to the greedy branch of
:meth:`~arbi_serve.spec_decode.dflash2_layers.CandidateSelector.select`
(temperature 0) over the same logits + hidden — the same modules and
ops, just written device-only and fixed-shape (no python-list stack,
no ``.item()``, no host reads, no data-dependent control flow), the
same capture-eligible shape the DSpark semi-AR greedy block holds.

Served by the captured realization
(:mod:`arbi_serve.spec_decode.dflash_realize_capture`) when one covers
the shape; :meth:`_dflash2_greedy_walk` is both the recorded core and
the eager fallback, so the two can never drift.

LEFT-TO-RIGHT candidate-selector SAMPLED block.

The stochastic sibling of :meth:`_dflash2_greedy_block` and the exact
DFlash2 analog of the DSpark markov semi-AR sampled block
(:meth:`~arbi_serve.spec_decode._dflash_driver_dspark._DFlashDsparkMixin._semiar_sample_block`).
Where the markov path adds a low-rank bigram bias to the base row, the
selector adds its predecessor-conditioned bilinear term to exactly the
``top_k`` candidate positions (their base logit is the ``unary`` the
greedy realizer scores), leaving the vocab tail at its base logit so
the proposal keeps full support. The biased row is then drawn through
the SAME :meth:`GraphSafeRejectionSampler.sample_drafter_token` the
bundled MTP head and the DSpark semi-AR chain use — per-row
temperature → top_k → softmax → top_p → min_p — so the returned
``q_j`` lives on the same processed-logits manifold as the verify
pass's ``p_target``, the identity standard rejection sampling needs
for losslessness. (The selector's own ``select(temperature>0)`` branch
softmaxes the top_k pool at raw temperature and is NOT on that
manifold — it is the eager/parity path, not the served one.)

The sampled token is the feedback: position ``j+1``'s bilinear term
conditions on it (``first_prev`` anchors position 0), and — when
``wm_ctx`` is armed — position ``j`` draws with the same per-row keyed
seeds the watermark verify uses at depth ``j``, coupling the realized
block to the target's keyed choices exactly as the DSpark path does.

Returns ``(tokens (Bd, S) int64, q_out)``; ``q_out`` / ``q_rows``
are the caller's block-q destination
(:func:`~arbi_serve.spec_decode._dflash_block_walk.write_block_q`).

The selector walk's one-position biasing step over ``base`` /
``hidden``: ``step(j, base_j, prev)`` returns position ``j``'s biased
full-vocab fp32 row given the realized predecessor.

The candidate pool and the hidden projection are taken once for the
block; each step adds the bilinear predecessor→successor term at the
candidate columns only — a candidate's base logit already is its
``unary`` — so the vocab tail stays at its base logit and the proposal
keeps full support. The upcast and the sparse add run as one pass over
the row (:func:`upcast_with_sparse_bias_`) into a row buffer reused
across the walk: every step overwrites it completely before the draw,
and the sampler chain masks it in place.

DSpark confidence dynamic-K + semi-AR block-realization cluster.

:class:`_DFlashDsparkMixin` holds the DSpark-checkpoint methods of
:class:`DFlashDrafter`: the confidence-head dynamic-K stash/decision and the
markov-biased semi-AR block realization (greedy + sampled). The methods
reference ``self.model`` bound in :meth:`DFlashDrafter.__init__`.

Score each drafted position with the DSpark confidence head and
stash the per-position accept probs on each done request, for the
next step's verify-plan build to truncate the draft length.

No-op unless a knob asked for confidence (``dflash_dynamic_k`` or
``dflash_conf_dump``), the checkpoint shipped a confidence head, and
we are at TP1 (the SPMD worker mirror carries ``mtp_next_drafts`` but
not this per-position confidence — truncating rank 0 alone would
desync). One small D2H on the already-eager draft path; skipped
entirely on the default path.

Confidence-head dynamic draft length for the slate's uniform K.

Returns a truncated ``dyn_k`` in ``[1, step_k]`` when
``ARBI_DFLASH_DYNAMIC_K`` is armed, the checkpoint shipped a
confidence head, and every row carries a per-position confidence
prefix (written by the previous step's :meth:`draft`) at least
``step_k`` long. ``dyn_k`` is the MAX over rows of each row's
confident prefix (the leading run whose predicted accept prob is
``>= threshold``) — the verify pass is uniform-K, so this never
truncates a row below its own confident length. Any missing / short
confidence prefix returns ``step_k`` unchanged (static K — safe).

Runs at any TP: DFlash forces SPMD OFF (the rank-0-driver topology —
:func:`resolve_dflash_spmd_routing`), so ``build_verify_plan`` — and
thus this method — executes on rank 0 ALONE, where the real drafter
owns the confidence. Worker ranks never build a plan (they only join
the target-forward collectives), so there is no per-rank K to
diverge. (A bundled-MTP SPMD path builds per-rank plans, but that head
has no confidence-K hook — this method is DFlash-only.)

True iff this slate's semi-AR block SAMPLES (and reports its real
per-position ``q``) instead of the greedy biased argmax + one-hot.

Rides the same ``ARBI_TRUE_STOCHASTIC_DRAFT`` routing the bundled
MTP head's true-stochastic chain uses, so one knob governs sampled
drafting engine-wide. Pure function of env flags + the rows'
SamplingParams — rank-symmetric with zero exchange.

LEFT-TO-RIGHT markov-biased GREEDY block: ``(Bd, S)`` tokens.

Position ``j``'s token is ``argmax(base_j + B(prev))`` where ``prev``
is the token realized at ``j-1`` (the anchor for ``j=0``). Lossless
under greedy verify — the target re-checks every drafted token
against its own argmax regardless of how the draft was formed.

The chain itself is
:func:`~arbi_serve.spec_decode.markov_head.semiar_greedy_tokens`:
device-only, fixed-shape, no ``.item()``, so the device-slots path is
unaffected and the block is cudagraph-capturable as one unit.

Served by the captured realization
(:mod:`arbi_serve.spec_decode.dflash_realize_capture`) when one covers
the shape; :meth:`_semiar_greedy_walk` is both the recorded core and
the eager fallback, so the two can never drift.

LEFT-TO-RIGHT markov-biased SAMPLED block.

Each position draws from its own per-row filtered distribution via
:meth:`GraphSafeRejectionSampler.sample_drafter_token` — the SAME
seeded, graph-safe categorical the bundled MTP head's
true-stochastic chain uses (per-row temperature → top_k → softmax →
top_p → min_p, so the reported ``q`` lives on the same
processed-logits manifold as the verify pass's ``p_target``; greedy
rows inside a mixed slate collapse to argmax + point mass).
``slot_offset=j`` decorrelates the S positions' draws from the one
shared step seed. The sampled token is the feedback: position
``j+1``'s bias conditions on it, and the returned ``q_j`` is
EXACTLY the distribution position ``j`` sampled from — the identity
standard rejection sampling needs for losslessness.

``wm_ctx`` (keyed watermark active) switches each position's draw
to the keyed-coupled sampler: position ``j`` draws with the same
per-row context seeds the keyed verify uses at depth ``j``
(``mtp_draft_seed_row`` over ring ++ realized prefix), so the
realized block tracks the target's keyed choices. The left-to-right
realization is what makes DSpark couplable at all — the plain
parallel DFlash block has no realized prefix to hash.

Returns ``(tokens (Bd, S) int64, q_out)``; ``q_out`` / ``q_rows``
are the caller's block-q destination
(:func:`~arbi_serve.spec_decode._dflash_block_walk.write_block_q`).

Mid-collective fail-loud cluster for :class:`DFlashDrafter`.

:class:`_DFlashFaultMixin` holds the two methods that fail loud on an OOM
inside a half-issued TP collective and the test fault-injection hook. The
methods reference ``self.engine`` bound in :meth:`DFlashDrafter.__init__`.

A draft-step fault fired INSIDE a mid-flight TP collective — fail loud.

The peer ranks are still blocked in the same ``all_gather`` / ``all_reduce``
on this NCCL communicator, so cold-pathing this step and continuing would
wedge NCCL silently (both GPUs 100%, ``/health`` = 000). Instead: latch the
engine's STICKY fatal fault (``/health/ready`` → 503, admission →
``engine_context_poisoned``) and raise :class:`DrafterCollectiveFault`,
which the engine loop re-raises out of ``run_forever`` (engine thread dies
→ orchestration restarts). Fail loud + tear down, never a silent wedge.

Test hook: force a simulated OOM at a named collective for validation.

``ARBI_DFLASH_FAULT_INJECT_COLLECTIVE={embed|lmhead|forward}`` raises a
fake CUDA OOM at that collective so the fail-loud path
(:meth:`_raise_collective_fault`) can be exercised on real hardware WITHOUT
having to engineer a genuine mid-collective allocation failure. Off (empty)
by default — never trips in production.

Context-stash method cluster for :class:`DFlashDrafter`.

:class:`_DFlashStashMixin` holds the engine forward-observer methods that
project tapped context features to per-request draft K/V and disable
speculation on a position gap / capacity exhaustion. The methods reference
``self.<attr>`` bound in :meth:`DFlashDrafter.__init__`.

Count — and RATE-report — one lost drafter observe pass.

Called from the two engine seams that catch this drafter's observe hooks
(``run_step._run_mtp_seed_forward`` and
:func:`~arbi_serve.spec_decode.mtp_verify_accept._observe_verify_forward`).
They are the same event seen after two different forwards, so they share
one counter and one escalation ladder rather than two that a reader would
have to add up.

NOTHING IS DROPPED: every failure still emits an ERROR line, so a log's
failure COUNT is unchanged. What changes is that only the first failure of
each ``(phase, exception type)`` carries a traceback — the 2nd through
535th copies of one stack are not evidence, they are what buried the rate —
and that the line at every doubling states the standing rate against the
passes that succeeded, which is the quantity "speculation is degraded"
actually means.

The contiguous tap slab backing ``target._dflash_tap_bufs``, or
``None`` when the target was armed without one.

The slab's column blocks are ordered by the ARMED ``_dflash_tap_ids``;
a drafter reading it whole gets its layers concatenated in that order,
so a drafter whose own ``target_layer_ids`` disagree would silently
receive mis-ordered context. Refuse instead.

``batch.<name>`` read from the pinned host mirror when there is one.

``batch.host_mirror`` is the pinned ``PiecewiseBuffers.h_*`` ring the
persistent build path H2D-copied FROM, so its tensors are bit-identical
to the device twins by construction, not a re-derivation. Reading the
DEVICE twin instead (``.tolist()`` / ``int(t[i])``) blocks the host until
the stream drains every queued kernel — on the observe path that is the
whole forward that was just enqueued. Falls back to the device tensor on
the fresh-alloc / CPU-stub path, where no mirror exists.

A verify row's absolute anchor position, read host-side.

``StepPlan.flat_positions`` is the host list the verify batch was
assembled from (``_assemble_verify_flat_batch``), so ``[lo]`` is the
row's anchor with no device read. Plans built outside
:func:`~arbi_serve.spec_decode.mtp_verify_plan.build_verify_plan` (CPU
stubs, ad-hoc test plans) carry none; those fall back to the slot's own
covered-to position, which keeps the old no-op behaviour rather than
manufacturing a spurious gap.

The first ``n`` tokens of each tap layer in tap-id order →
``(n, len(tap_ids) * hidden)``. Reads the persistent
capture/compile-safe slab when armed (production) — the tap-id
column order of :meth:`setup_dflash_tap` makes that read a VIEW —
else the eager dict (CPU tests). ``None`` when the dict path
captured nothing (misrouted forward).

The slab view aliases live tap memory: the next tapped forward
overwrites it, so callers must consume the result before the next
forward (every caller collects and projects in the same block).

PER-FORWARD contract: the tap holds the MOST RECENT tapped
forward's tokens only (each prefill chunk overwrites the slab
from row 0), so ``n`` must be that forward's ``cu_seqlens_q[-1]``
— every caller here collects immediately after its forward. A
read past the last forward's span (the chunked-prefill trap)
RAISES via :func:`ensure_tap_read_within_last_forward` instead
of silently returning another forward's rows.

``forward_tokens`` hands the guard the count the caller already
holds. Every production caller here passes its own forward's
``cu_seqlens_q[-1]``, read from the host twin precisely so the
observe pass never parks the host on the just-enqueued forward —
so the guard was reading back from the device, once per prefill
chunk, the very number the caller had. A caller that cannot name
its forward omits it and pays the device read.

``first_pos`` continues.

``lengths[slot]`` counts the committed positions the slot holds and
``_slot_base[rid]`` is the absolute position its index 0 maps to, so
``_slot_base + lengths`` is the absolute position the slot is covered
TO. A span that starts anywhere else means the drafter missed a
commit (a depth-0 step, a prefix-cache hit) or jumped.

Two independent triggers, because neither covers the other:

* ``Request.mtp_drafter_context_stale`` — set by whatever ran the
  row without feeding the drafter (the load valve's batch-wide
  ``mtp_k=0``). Exact, and the only one available where the host
  ``lengths`` mirror is not authoritative.
* the position compare — catches every OTHER way a commit can go
  unobserved. ``compare=False`` for the device-resident stash,
  whose appends advance ``lengths_dev`` only, leaving the host
  ``lengths`` stale and the compare meaningless there.

RE-SEED rather than disable: the drafter's RoPE positions and context
tags are all slot-relative and it attends only over its own KV, so a
slot may start at any absolute position. The cost is a shorter
context that refills, not a request that never speculates again.
``None`` when the slot pool is exhausted (the request is disabled).

Project a contiguous feature span to draft K/V and append it to
the request's slot. Position-gap or slot/capacity exhaustion
disables speculation for the request (output stays correct via
the cold verify path).

Returns the slot id on a successful append (so the sharded-drafter
observe broadcast can ship the resolved slot to workers), else
``None`` (row skipped / disabled).

Record a new high-water in the REALIZED context-assemble width.

The serving floor reserves this path's peak at the widest span the
step limits allow (``min(chunk_prefill, max_batched_tokens)``), and
that reserve comes out of the KV pool. Whether serving ever BUILDS a
span that wide is a different question from what the config asked
for — admission narrows the prefill chunk to what fits and says so
only in a boot warning, so a run can spend its whole life at half the
configured width with nothing but the request rate to show for it.

This is the census that answers it, taken where the bytes are actually
spent rather than where the knob is read. One line per new high-water,
so a boot emits at most a handful and a reader can compare the widest
span served against the row bound the reserve was sized at.

SHARDED build: ship the resolved per-row stash spec to workers so
each rank populates its OWN shard KV from its local tap capture.

No-op on the replicated build (rank 0 owns the only draft KV) and on
every worker rank / TP=1. ``obs`` rows are ``(slot, tap_lo, n_keep,
first_pos)`` collected from the successful :meth:`_stash` calls.

Stash per-request context K/V after a legacy-path forward
(prefill chunk or K=1 seed step). Called from
``run_step._run_mtp_seed_forward`` right after the model forward.

Row boundaries and absolute positions are read HOST-side (see
:func:`_host_batch_twin`): this runs after the seed forward is
enqueued, so a device read here would park the host for the whole
forward every step.

Stash the accepted-prefix context K/V after a verify pass.

Per row, the verify forward processed ``K_eff + 1`` tokens
``[anchor, d_1 .. d_K]``; the accepted run keeps slots
``0 .. n_accepted`` — exactly the positions that extend the
request's context to the new committed length minus one.
``results`` is aligned with ``mtp_row_indices_host`` (rows with
K_eff > 0); cold rows (K_eff == 0) contribute their anchor slot.

DEVICE-resident verify-pass stash (async path).

The async analogue of :meth:`observe_verify_forward`: consumes the
verify accept's DEVICE ``n_accepted`` ``(len(mtp_rows),)`` tensor
(no host pull) and scatters each row's accepted-prefix context K/V
into its slot with :meth:`DraftKVSlots.append_device`. Called from
the async verify path right after ``_async_verify_accept`` produces
``n_accepted_gpu`` and BEFORE the drafter seed, so the next draft
reads fresh context without the sync-verify host wait.

Only the K_eff>0 verify rows are stashed here (they own the uniform
``K+1`` span); cold rows never reach the async path. A row already
disabled / holding no slot is skipped. The host path's
position-continuity guard runs here too (:meth:`slot_for_span` off
the plan's host ``flat_positions``) — verify extends inductively
only while every commit feeds the drafter, which a depth-0 step
does not.

Numerically identical to the host stash: the same
:meth:`DFlashDraftModel.project_context_stacked` projection over the
same tap features + absolute positions, the same committed-prefix
length (``n_accepted + 1``), the same slab layout — just driven by
device tensors end to end (parity-tested).

True iff this row's context K/V must be stashed this step.

``sampling.mtp_k`` is the depth for THIS step and is transiently
zeroed by the scheduler's load-adaptive spec disable and by any
depth policy. A row that already holds a slot keeps being fed at
depth 0: ``lengths[slot]`` IS the drafter's absolute position, so
a skipped commit leaves the slot one token behind the request
forever, and the next stash at a real position trips the
:meth:`_stash` gap check and disables the request for its whole
lifetime.

Resolve ``ARBI_TRUE_STOCHASTIC_DRAFT`` into the per-row routing threshold.

``None`` = stochastic drafting disabled (mode ``"0"`` — every slate
keeps the greedy-argmax + one-hot proposal). Otherwise a slate row
drafts stochastically iff ``temperature > 0 and temperature >=
<returned value>``: mode ``"1"`` returns ``0.0`` (every stochastic
row), mode ``"auto"`` returns ``ARBI_TRUE_STOCHASTIC_DRAFT_TEMP``
(per-row temperature routing; ``inf`` ≡ mode ``"0"``, ``0`` ≡ mode
``"1"``).

This is a PURE function of the env-derived runtime flags, and the
per-row decision (:func:`stochastic_draft_row`) is a pure function of
the request's SamplingParams — both are identical on every SPMD/TP
rank, so the routing is rank-symmetric with zero exchange (the params
themselves ride the existing DrafterOp broadcast on the rank-0-driver
path and the request mirror on the SPMD path).

The resolved thinking regime the drafter sees for this row.

Reads the per-request signal
(:attr:`arbi_serve.engine.request.SamplingParams.resolved_thinking`)
that rode the admission->``SamplingDigest``->every-rank rail to this
routing site, so it is identical on every SPMD/TP rank (both the real
``SamplingParams`` on the rank-0 driver path and the
``_DigestSamplingView`` mirror on the worker path expose it). ``True`` /
``False`` = the concrete thinking fact; ``None`` = "no regime info"
(non-chat / raw completions).

PLUMBING WITNESS ONLY — the signal is now READABLE here, but
:func:`stochastic_draft_row` deliberately does NOT consult it, so this
accessor's only current consumers are tests witnessing distinct
per-request regimes arrive here without perturbing today's
byte-neutral routing.

True iff this row drafts by sampling the drafter's own distribution.

``threshold`` is :func:`stochastic_draft_temp_threshold`'s non-``None``
value. Greedy rows (``temperature <= 0``) never qualify regardless of
the threshold (a ``threshold == 0.0`` must reproduce mode ``"1"``,
which only routes ``temperature > 0`` rows).

The per-request thinking regime is available alongside this decision via
:func:`draft_row_regime` (``sp.resolved_thinking``) but is intentionally
NOT read here — routing stays a pure function of ``temperature`` +
``threshold``.

Whether ``ARBI_TRUE_STOCHASTIC_DRAFT`` routes ANY row to sampled drafting.

The mode knob alone, with no request in hand: ``"1"`` / ``"auto"`` can
route, ``"0"`` never can. This is the BOOT half of the sampled-drafting
decision — the flag is not in the live config-override registry
(:mod:`arbi_serve.config_overrides`), so an installed flag overlay cannot
carry it and the answer is fixed for the life of the process. Boot-time
consumers (the stochastic drafter twin's capture, the depth gate's boot
refusal, the verify pool's dense-``q`` scratch) read it here rather than
each spelling the ``!= "0"`` comparison, so there is ONE answer to "can
this boot draft a distribution".

The per-ROW half is :func:`stochastic_draft_row`; the two together are
:func:`true_stochastic_draft_active`. No counter fires here: this is a
boot-shape question, not the routing read that
:func:`stochastic_draft_temp_threshold` certifies.

Whether ANY verify slate on ``eng`` can carry a DENSE drafter ``q``.

The reachability predicate for the vocab-scale dense-``q`` working set:
:attr:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers.sc_draft_probs`
(the densified ``(K, B, V)`` proposal the rejection sampler reads) and the
per-row ``(K, V)`` carries the next slate is densified from. Both exist iff
some row's ``mtp_next_draft_probs`` is a dense tensor rather than
:data:`~arbi_serve.spec_decode.drafter.POINT_MASS_Q`, and that is decided
by the DRAFTER, not by the verify step.

True on either of two boot-fixed grounds:

  * :func:`true_stochastic_draft_enabled` — the bundled MTP head, DFlash
    and DSpark all realize a sampled draft (and report its real ``q``)
    exactly when this knob routes;
  * an external draft MODEL is configured or attached. Its chain reports a
    dense per-position ``q`` for EVERY stochastic slate, on no knob of its
    own, so the flag does not describe it.

Both grounds survive for the life of a ``VerifyBuffers`` generation: the
flag is boot-only, and ``mtp_draft_model_path`` is ``scope="backend"`` in
:mod:`arbi_serve.config_overrides`, whose rebuild allocates a fresh pool.
Read the CONFIG as well as the attached driver so the answer is the same
before and after attach — the budget resolves it at profile time, when no
drafter exists yet.

Encode below-threshold auto-mode rows as greedy point masses.

A sampled slate uses one dense proposal tensor for every row. Rows that
do not qualify for sampled drafting retain greedy-drafter semantics by
using a clone with ``temperature=0`` and ``min_p=0``. The input sequence
and qualifying rows are left unchanged.

Per-row effective speculation depth — FAIL LOUD, never clamp.

A request's ``sampling.mtp_k`` is bounded to ``driver.max_k`` at
ADMISSION (``engine.submission._build_request`` raises HTTP 400 for a
larger K), because the drafter+verify cudagraphs and the verify
snapshot pool are sized for ``max_k`` at boot. So by the time the
verify pass runs, every row already satisfies ``mtp_k <= max_k``.

A row whose ``mtp_k`` exceeds ``max_k`` at this point is an admission
bug (or a peer mutating sampling.mtp_k post-admit), so we RAISE rather
than silently serving a different depth. ``mtp_k == 0`` (MTP off / cold)
floors to 1 at the call site (decode-via-verify).

DIAGNOSTIC: unified accept-accounting dump.

``ARBI_MTP_ACCEPT_DUMP=<path>`` makes every token-emitting decode path
append one JSONL record, so client-visible output tokens can be
reconciled 1:1 against the engine's step accounting:

  * ``{"reqs": [...], "acc": [...], "k": [...]}`` — a verify commit
    (:meth:`MtpDriver.commit_results`), one committed token per row plus
    that row's accepted drafts.
  * ``{"cold": [...], "dl": [...], "ol": [...]}`` — a whole-slate K=0
    cold tick (:func:`_commit_cold_and_accepted` k_eff==0 rows): one
    frontier-argmax token per row, ZERO speculation. ``dl`` is each
    row's cached-draft length at commit time, ``ol`` its output length
    (age proxy). These ticks bypass ``commit_results`` entirely.
  * ``{"seed": [...], "pf": [...], "emit": [...]}`` — an MTP seed /
    legacy step (:func:`_run_mtp_seed_forward`): ``pf`` marks prefill
    rows, ``emit`` marks rows that actually sampled a token this step.

Every record carries ``pid`` so rank-0 / SPMD-worker duplicates can be
separated in analysis. Eager host path only; not for perf runs.

The dump path, or ``None`` when the diagnostic is off.

Reads the RuntimeFlags registry (a cached snapshot), NOT ``os.environ``
directly — the single-source invariant (every ARBI_* flag is a RuntimeFlags
field) is enforced by ``tests/test_runtime_flags_single_source.py``, and the
snapshot is also cheaper than an env lookup per call.

``ARBI_MTP_PLAN_DUMP`` — the per-rank verify-plan dump path.

Records the per-tick verify-plan DECISION on every rank so the
rank-symmetry of the TP>1 MTP prefill seed
(:meth:`DistributedEngineDriver._spmd_seed_mtp`) can be PROVEN rather than
argued. EVERY rank resolves drafts in the SAME place —
:func:`mtp_verify_plan._verify_step_k_and_drafts`, reached by rank 0
through :func:`build_verify_plan` and by each SPMD worker through
:func:`run_verify_step_spmd_worker` -> ``build_verify_plan`` — and each
appends ``{"slate": [...], "effk": [...], "dh": <drafts digest>, "pid":…}``.
If two ranks ever disagreed on which rows are cold, or on the draft token
values, their per-pid record sequences would differ — diff the per-pid
dump files to compare.

(``distributed.spmd.derive_verify_tensors`` carries the same collapse
logic but is reached only from tests — the served SPMD verify path goes
through ``build_verify_plan``.)

Append one per-rank verify-plan record (no-op when the flag is off).

``eff_drafts`` rows are ``None`` on a cold (collapsed) step; they
digest as the literal string ``"cold"`` so a rank that collapsed and
a rank that did not produce visibly different records.

Per-slot draft-DEPTH gate for the bundled-head chain.

A static ``K`` buys the same depth on every step. The drafter's own
logit row says, per slot, how sure it is: the gap between the top-1 and
top-2 probability of ``softmax`` over that row. This reads that gap at
slot ``d - lag`` and decides whether to run depth ``d + 1`` at all, so a
step whose run is about to be long keeps its depth and one whose drafter
has already hedged stops.

WHY PER SLOT AND NOT PER STEP. A per-step rule sets the NEXT step's
uniform K from this step's outcome, which ratchets: a confident prefix at
step ``t`` caps step ``t+1``, and the cap has almost no way back up. The
decision has to be taken inside the chain, against the slot the drafter
just produced.

WHY THE MARGIN AND NOT TOP-1. A slot at ``p1=0.50, p2=0.45`` is ambiguous
in a way one at ``p1=0.50, p2=0.05`` is not, and top-1 cannot tell them
apart. Both signals separate acceptance, but the margin does it at a much
lower stop rate — it needs a threshold near the bottom of the range where
top-1 needs one near the middle.

WHY THE LOGITS AND NOT THE PROPOSAL ``q``. The gap is a property of the
drafter's LOGIT row, which every drafting mode computes — the greedy
chain has to form it to take its argmax. Reading it off ``q`` instead
made the gate a function of two things it has no business depending on:
whether the slate happened to draft stochastically at all (a greedy
chain carries ``POINT_MASS_Q``, a marker with no margin in it), and the
REQUEST's own sampling knobs (a ``top_k=1`` row's filtered ``q`` has
margin 1.0 however unsure the model was, so the gate went inert exactly
where it was cheapest to help). ``softmax`` over the raw row is the
model's own confidence and nothing else's, it is the SAME number ``q``
carried wherever the filtering was inert, and it exists in every mode —
which is what lets one threshold mean one thing.

WHY LAG. The decision needs a host-visible number, so the device→host copy
of slot ``d``'s margin has to land before depth ``d + 1`` is launched. At
``lag = 1`` the copy issued after slot ``d`` is not read until the launch
of depth ``d + 2``, so a whole head forward sits between the copy and its
sync and the host never waits on the device.

WHAT THE DEPTHS RUN ON. A decision between depths needs a seam the host
can act in, which a fused K-step cudagraph does not have — it replays every
depth it recorded in one launch. A drafter chain captured as SEGMENTS (one
graph per step, ``--mtp-capture-chain-segments``) has that seam and still
replays recorded kernels, so the gate rides it; on a boot without the
segments the only way to serve the gate is the eager chain, which costs
several times what the gate can save.

WHY IT NEEDS ``ARBI_ACCEPT_INVARIANT``. A depth gate varies ``T = D + 1``,
the verify block's width. On a recurrent (GDN) target under default
numerics the verify forward's recurrent half is a chunked scan over ``T``,
and its output seeds the next step's drafter — so how much of the text the
gate changes is itself a function of the gate. ``ARBI_ACCEPT_INVARIANT=1``
pins the accept-critical GEMM row count and routes GDN verify to a
per-token recurrence, both width-free, which makes a varying depth a
throughput decision instead of a trajectory one. That is a correctness
property of the gate, not a benchmarking convenience, and
:mod:`arbi_serve.spec_decode.depth_gate_boot` refuses the combination
rather than serve it.

CHAIN ONLY. The decision needs a per-slot seam, which only the bundled-head
chain has: DFlash realizes every position from one forward and a tree takes
its step size from the node count. An armed gate on a slate with no depth
to decide runs the full depth and counts itself inert rather than reading
as a gate that chose ``K``.

OFF costs one attribute load per drafter chain.

The resolved per-slot depth policy for this process.

``margin`` is the top-1/top-2 softmax gap at or below which a slot
stops the chain. ``lag`` moves the predictor strictly further back: the
decision to run depth ``d + 1`` reads slot ``d - lag``.

``null_rate`` replaces the margin test with a coin at that per-slot
rate while keeping every other cost — the same forwards, the same
device reduction, the same device→host copy and the same sync. It is
the control an effect this size needs: an arm that cuts at the same
rate through a signal that carries nothing. A same-arm repeat cannot
play that role, because the mechanism under test is a deterministic
function of the thing being varied.

``ARBI_MTP_DEPTH_GATE`` resolved to a policy, or ``None`` for static K.

OFF costs one attribute load: nothing below the flag check runs, so an
unarmed boot never builds a policy and never touches a margin.

Per-row top-1/top-2 margin for one drafted slot, from its LOGITS.

THE definition of the gate's signal, and the only one: every producer
calls this and none re-derives it. ``logits`` is the drafter's ``(B,
V)`` row at the slot — the same row the pick reduces over, penalties
included — and the result is ``(B,)`` fp32.

Taken over ``softmax`` of the raw row, not over the request's filtered
proposal. The filtered ``q``'s margin answers "how peaked is the
distribution this row will be SAMPLED from", which a ``top_k=1`` row
pins at 1.0 whatever the model thought; ``softmax`` of the row answers
"how sure is the drafter", which is the quantity a depth decision
wants and the one that exists identically under greedy and stochastic
drafting. Where the filtering is inert the two are the same number.

One slot's per-row margins reduced to the slate's decision.

MAX across rows, not min or mean: the slate proposes a UNIFORM depth,
so a row that is still confident would pay for its neighbour's doubt.
Taking the most confident row lets the gate stop only where the whole
slate has hedged, and collapses to the row's own margin at ``B = 1``,
which is where the threshold was calibrated.

An armed gate met a slate with no depth to decide.

A ``K = 1`` slate drafts the one slot every chain drafts before a
predictor exists, so the gate is inert on it however its flag reads.
Counted so a boot serving such traffic reads as "nothing to decide"
and never as "the gate chose K".

The PER-SLOT rate at which a chosen-depth histogram stopped chains.

Per SLOT INSPECTED, not per chain: a chain that stopped early was never
asked about the slots it did not reach, so dividing stops by chains
would report a rate no coin can be set to. This is the number the
rate-matched null takes, and matching the wrong one makes the null cut
harder than the gate and flatters the gate by the difference.

``(slots, rows)`` fp32 device scratch for one eager chain's margins.

Held rather than allocated per chain for the same reason the pinned
landing buffers are: the contents are dead the moment the slot that
owns them has been decided, and a fresh allocation per chain would
cost more than the reduction it lands.

One drafter chain's depth decisions.

Two calls, in the order the chain makes them: :meth:`observe` hands it
the per-row margins of the slot the head just produced, and
:meth:`stops_before` asks whether the next depth should run at all.

The two are separated by ``lag`` slots on purpose. :meth:`observe`
issues an ASYNC copy of the slot's margin and records an event;
:meth:`stops_before` waits on that event. With ``lag = 1`` a whole head
forward is launched between them, so the wait is already satisfied when
it is reached and the host never blocks the chain.

Queue slot ``slot``'s margin for the host, without waiting for it.

``row_margins`` is the ``(B,)`` device tensor the head filled at
this slot (:func:`slot_row_margins`). The slate reduction, the
threshold test and the device→host copy all stay HERE, host-issued
in the seam between depths, so a captured chain hands this the
buffer its recorded kernels wrote and pays nothing extra for the
decision.

Whether depth ``slot`` should be skipped, from slot ``slot-1-lag``.

The first ``lag + 1`` slots have no predictor behind them and are
always drafted; the coin in the null is flipped exactly where the
real gate would have consulted a margin, so the two arms inspect
the same slots and cut at the same per-slot rate.

What the per-slot depth gate needs from the rest of the stack.

Checked ONCE at boot, by name, with the flag that caused each one. Every
condition here fails the same way if it is not checked: the gate serves,
the chosen-depth histogram looks spread, tok/s looks plausible, and the
number the A/B reports is not the number the gate is worth. None of them
raises on its own.

The load-bearing one is ``ARBI_ACCEPT_INVARIANT``. A depth gate varies the
verify block's width, and on a recurrent target under default numerics
that width IS the chunk boundary of the verify scan whose output seeds the
next draft. The gate would then be scored on a trajectory it moved. There
is no fallback that fixes this — a gate that varies depth on unpinned
numerics is measuring itself — so the combination is refused rather than
served.

Refuse a gate on a drafter whose chain depth is not a decision.

Read off the DECLARED capability — a direct attribute read, not a
probe with a default. A probe answers "no" identically for a drafter
that cannot gate and for one that simply never declared, which is how
a capability hook comes to read zero for a structural reason while
looking like a tuning one. A drafter missing the declaration fails
here, at boot, by name. Not an ``isinstance`` either: the capability
is the thing, and a class check would admit a future subclass that
dropped it and refuse a future drafter that grew it.

Only the bundled-head chain takes its depth one slot at a time.
DFlash realizes every position from ONE forward, so there is no depth
to skip; an external draft model runs its own loop that has never been
handed a per-slot signal. On either, the flag would read ON while the
depth never moved — which is exactly what a null looks like.

Refuse a varying depth on unpinned accept-critical numerics.

A gated step's verify block is ``chosen_depth + 1`` rows instead of
``K + 1``, which is the same hazard every other data-dependent verify
width carries — so the condition, and the message it raises, come from
the one place that states it
(:func:`~arbi_serve.spec_decode.verify_width_boot.refuse_unpinned_verify_width`)
rather than from a copy here.

Refuse the gate and a tree together.

A tree takes its step size from the NODE COUNT, not from a chain depth
the gate could shorten, and its per-node expansion returns the top-``W``
candidates rather than the row a margin is read from. The gate would be
inert while its flag read ON.

Refuse the gate under TP>1.

The chosen depth is a data-dependent branch taken on a host-read
reduction, so a single-ULP difference between ranks sends them down
chains of different length: they then issue different numbers of
per-step collectives and desync the next step's slot allocation. That
is a deadlock, not a slowdown, and no capture form fixes it — the
divergence is in the decision, not in what the depths run on.

Refuse a policy that cannot describe a decision.

A negative lag would read a slot the chain has not drafted yet, and a
lag at or past the slate's K leaves no depth the gate can decide, so
the flag would read ON over a chain it never shortens.

Device-resident draft tokens between the drafter chain and the next verify.

The drafter chain writes its ``(K, B)`` draft block on the device. The next
verify plan needs those tokens only ON the device (as verify input ids and
as the accept test's draft tensor), so pulling them to the host after every
chain is a stream drain the step does not need: the host waits for the
chain to finish, then does the drain, the plan build and the next forward
launch with an empty GPU queue.

Deferring the pull lets the host build step ``N+1``'s plan while step
``N``'s chain is still running. Each row keeps a :class:`DeviceDrafts`
handle — a snapshot of the chain output (the captured chain's output buffer
is borrowed and is overwritten by the next replay), its column, and the
event recorded after the snapshot — and its host ``mtp_next_drafts`` list
holds :data:`PENDING` placeholders of the right length, so every consumer
that only needs the depth keeps working. A consumer that needs the token
values calls :func:`resolve_host_drafts`, which performs the pull for that
row alone.

Whether drafts may stay on the device between chain and verify.

Requires the flag, a CUDA engine and a single rank: at TP>1 the worker
mirror derives its verify plan from the host draft lists.

Snapshot the chain's ``(>=depth, B)`` block and hand each row its handle.

The snapshot is enqueued on the current stream right behind the chain,
so it is stream-ordered after the chain's last kernel and safe against
the next replay overwriting the chain's own output buffer.

Drop every row's carried drafts because the drafter that made them has.

A device handle names a snapshot allocated under the OUTGOING member's
allocator, and a host list — placeholders or values — names draft-KV
slots in a store that is being replaced, so neither survives a drafter
swap. Clearing both cold-paths each row's next step, which is the one
step the incoming drafter needs in order to draft for itself.

Copy each row's first ``k`` device drafts into ``draft_tensor[:k, :B]``.

One copy when every row shares one snapshot in column order (the steady
state), else one copy per row. Each source event is waited on by the
current stream first, so the copies order after the chain that produced
them.

Write ``draft_view`` (``(k, B)``) into ``input_ids`` at ``flat_index``.

``flat_index`` lists the flat verify positions of row 0's ``k`` drafts,
then row 1's, and so on — the transpose order of ``draft_view``.

DFlash block-diffusion draft model (eager, transformers-free).

A standalone port of z-lab's ``DFlashDraftModel`` for in-engine speculative
decoding. The draft is a small dense Qwen3
transformer (typically 5 layers) that, conditioned on hidden states
tapped from several target-model layers, denoises a masked block of
``block_size`` tokens in a single parallel forward.

It is deliberately self-contained: the slim runtime image carries no
``transformers``, so the Qwen3 attention / RoPE / SwiGLU are reimplemented
here in plain torch (reusing only :class:`arbi_serve.models.layers.RMSNorm`).
The draft shares the target's token embeddings and LM head — the caller
passes ``noise_embedding = target.embed_tokens(block_ids)`` and applies
``target.lm_head`` to the returned hidden, exactly as the reference does.

Conditioning: target hidden states from ``target_layer_ids`` are concatenated
(``len(ids) * hidden`` wide), projected by ``fc`` + ``hidden_norm`` to
``hidden``, and injected as additional K/V context at every draft layer —
each layer attends over ``[context ++ noise_block]``. Full-attention layers
are bidirectional; ``sliding_attention`` direction is selected by the
checkpoint architecture (see ``DFlashConfig.layer_types``). Runs eager;
the per-block attention is tiny (context + block_size positions).

Layout note: the config dataclass / ``config.json`` parsing lives in
:mod:`arbi_serve.spec_decode.dflash_config`; the layer primitives (cache,
SWA mask, RoPE, attention, MLP, decoder layer, ``_set_submodule``) live in
:mod:`arbi_serve.spec_decode.dflash_layers`. Both are re-exported here so
``arbi_serve.spec_decode.dflash`` remains the single import surface.

Enter :func:`declare_drafter_forward` around a method that runs one of
the drafter's own quantized linears.

On the CALLEE rather than at each call site, which is the opposite of the
choice :func:`~arbi_serve.runtime.forward_declaration.forward_declared`
makes for the target, and for a reason that only applies here: these
methods are reached from the driver, from the memory-budget probes, from
the worker mirror and from the capture pre-sweep, and a declaration owned
by the call site would have to be remembered at each. The one method that
is NOT decorated is ``forward_block_fixed``, and that is not an omission —
it is replaced by a ``torch.compile`` trampoline at boot
(``install_torch_compiled_methods``), so a context manager here would sit
INSIDE the compiled region and be traced. Its callers declare instead, and
``tests/test_forward_declaration_seam.py`` holds them to it.

Block-diffusion draft. Returns the post-norm hidden for the block.

The caller supplies ``noise_embedding`` (the target's embedding of the
masked block, with the verified anchor in slot 0) and the raw
``target_hidden`` (concatenation of the tapped target layers), and
applies the target LM head to the returned hidden. ``position_ids``
must span ``[context ++ block]`` (length ``ctx + block_size``); the
queries take the trailing ``block_size`` positions.

Eager whole-sequence denoise — the NUMERICAL REFERENCE, not served.

No production caller reaches this; the parity tests use it to pin
:meth:`forward_block_fixed` (the served forward, over
:class:`~arbi_serve.spec_decode.dflash_kv_slots.DraftKVSlots`) against
the upstream semantics. See :class:`DraftKVCache`.

Per-position accept LOGIT from the DSpark confidence head.

``draft_hidden`` ``(..., H)`` is the post-norm block hidden that
produced each drafted position's base logits; ``prev_token_ids``
``(...)`` is the REALIZED predecessor token feeding that position's
markov bias (the anchor for block slot 0). The head's single Linear
maps ``concat([draft_hidden, markov_w1[prev]])`` (``H + markov_rank``)
to one logit; ``sigmoid`` is the predicted probability the target
accepts this drafted token. ``markov_first`` flips the concat order
(calibration knob — the checkpoint's trained layout is resolved
empirically). Fails loud if the checkpoint shipped no confidence head
or no markov head (the input needs the markov embedding).

Stack per-layer K/V projection + K-norm weights for
:meth:`project_context_stacked`. Dense drafters get the one-GEMM
weight; EXL3 drafters (no dense ``.weight``) get only the stacked
norm weights (their GEMMs stay per-layer). Idempotent; call after
weights are bound.

All-layer context K/V projection in one pass.

Returns ``(k_all, v_all)`` shaped ``(L, ..., nkv, hd)`` — K
post-norm + post-RoPE, V raw — numerically equal to stacking
:meth:`project_context`'s per-layer output (bit-identical on the
EXL3 path, where the GEMMs stay per-layer; the dense one-GEMM
path is the same matmul over a concatenated weight). Fixed-shape
and loop-free over rows → cudagraph/compile-safe.

Project tapped target features straight to per-layer draft K/V.

Returns ``[(k, v), ...]`` per draft layer, each ``(1, n_kv, n,
head_dim)``, post-norm and post-RoPE — exactly the context
entries :meth:`forward` would have appended to a
:class:`DraftKVCache`. Lets the engine persist draft K/V
incrementally (per accepted token) and drop the feature buffer.

Denoise one block against precomputed context K/V.

Equivalent to :meth:`forward` with a cache holding the same
context entries (see :meth:`project_context`); queries/keys for
the block positions are computed fresh and never persisted.

``ctx_key_pos[i]`` carries the ABSOLUTE position of each stored
context key for layer ``i`` — used to build the SWA mask. When
omitted, the context is assumed to span contiguous positions
``[0..n_ctx)`` (full per-request context, no ring). A sliding
layer fed only its trailing window passes the window's true
absolute positions so the mask matches the full-context result.

``(plan, window)`` for one layer's maskless attention geometry.

``window`` is the architecture's sliding rule, or ``None`` on a
full-attention layer (which is bidirectional over whatever context
view it is fed). ``span`` bounds how far back a query can reach —
the window, or the whole stored width when there is none — and is
what sizes the packed row stride.

Plans are keyed by geometry, not by layer, so a uniform drafter
allocates ONE set of buffers for all its layers; the per-layer
:meth:`DFlashVarlenPlan.refresh` before each use is what keeps a
shared plan correct.

One layer's attention with NO mask tensor.

Q/K/V keep the head axis last throughout, so the projections need no
transpose and the varlen call takes them as produced. The window and
the per-row key count carry what the additive bias carried; see the
module docstring of
:mod:`arbi_serve.spec_decode.dflash_flash_attend` for why the two
formulations select the same keys.

Fire the per-layer attention-route counters for this build.

Which route each draft layer takes is decided by ``_varlen_provider``
and that layer's attention sink, both fixed before the first forward,
so the pair is recorded ONCE here instead of per layer per call from
inside :meth:`forward_block_fixed`.

That placement is load-bearing, not tidiness. ``PathCounter.fire``
mutates a module-global's integer attribute; traced by Dynamo it
becomes a guard on that attribute's value, which the traced body then
invalidates on its way out. Every call to a compiled
``forward_block_fixed`` therefore failed its own guard and re-traced —
the per-call recompile that made compiling the draft forward cost a
trace per call rather than a trace per shape. Out here the counters
say exactly what they said before and the compiled region reads no
global at all.

Idempotent in meaning only: it fires once per call, so callers record
a build once. The pair still partitions every attended draft layer.

Denoise B blocks in one FIXED-SHAPE forward (cudagraph-capturable).

The capture-ready sibling of :meth:`forward_block_batched`: no
python per-row loops, no data-dependent shapes — every mask is
computed by broadcasting from ``start_pos`` / ``ctx_kpos`` device
tensors, so the whole body can be recorded once and replayed with
fresh buffer contents. Numerically equivalent to running
:meth:`forward_block` per row (parity-tested): a context key
column is REAL iff ``0 <= kpos < start_pos[row]``; everything
else is pad and gets ``-inf`` bias (softmax weight exactly 0, so
over-padded buckets do not perturb the result). Sliding layers
additionally apply the architecture's sliding-window rule against
the true absolute key positions; full layers stay bidirectional
over the row's real keys. Rows with ``start_pos == 0`` (inactive /
padding) see no context and produce garbage the caller discards.

Block positions are ``start_pos[b] + [0..block)``; queries take
the whole block (slot 0 = anchor).

``ARBI_DFLASH_ATTN_BACKEND=flash_varlen`` carries the same two
constraints without a mask tensor — see
:mod:`arbi_serve.spec_decode.dflash_flash_attend` — which is the only
thing that makes a fused flash backend eligible for this attention.

Denoise B blocks in a SINGLE forward, one row per active request.

Padded-with-mask batching of :meth:`forward_block`: per layer the
rows' context K/V are right-padded to the batch's longest context
and an additive attention mask zeroes both the pad slots and the
masked-out positions, so each row computes EXACTLY what its
single-row :meth:`forward_block` would. Context lengths differ per
row (each slot has its own committed length + per-layer ring
window); ``ctx_key_pos[i][r]`` carries the true absolute key
positions of row ``r`` at layer ``i`` (the ring returns them out of
``[0..n)`` order), used to build the sliding-window mask and to mask
the pad tail. Full-attention layers stay bidirectional over each
row's real keys.

Returns the post-norm block hidden ``(B, block, H)``.

Build the optional modules the checkpoint's key set declares, so
the state-dict load binds every tensor.

``shapes`` maps canonical checkpoint keys to their LOGICAL tensor
shape. ``quant_mirror(runtime_path, in_features, out_features)``
returns a bound quantized Linear for a module the checkpoint
quantized, or ``None`` for a dense one; the dense loader passes no
mirror at all.

DSpark heads ship their OWN ``embed_tokens.weight`` — the drafter
must use that, not the target's (the head was trained against its
own copy) — and may ship ``lm_head.weight`` too. An ``lm_head``
WITHOUT an ``embed_tokens`` is a malformed checkpoint → fail loud;
an ``embed_tokens`` without an ``lm_head`` projects through the
TARGET's head. z-lab DFlash heads carry neither and borrow both.

``confidence_head.proj.*`` (DSpark accept-rate predictor) is loaded
into :attr:`confidence_head`; :meth:`confidence_accept_logits`
consumes it for confidence-head dynamic-K
(``ARBI_DFLASH_DYNAMIC_K``).

Load a DFlash checkpoint (config.json + safetensors).

Weight quantization resolves through the
:mod:`arbi_serve.weight_quant` backend registry — the same path the
main model takes: the checkpoint's TENSOR KEYS select the backend,
which supplies each projection's mirror class and binds its payload
(:mod:`arbi_serve.spec_decode.dflash_quant_load`). A checkpoint no
registered backend claims loads the dense bf16 path; the norms
(RMSNorm ``.weight`` tensors) stay dense in both. A checkpoint whose
``config.json`` declares a ``quantization_config`` that no backend
detects is refused rather than loaded dense.

``tp_shard=True`` builds the TENSOR-PARALLEL drafter: q/k/v and
gate/up projections column-shard (head-aligned), o_proj and
down_proj row-shard with an all_reduce, ``fc`` + norms stay
replicated. Requires the ambient parallel config
(``get_parallel_config``) to be initialized; at tp_size == 1 it
degrades to the replicated build.

Swap the dense per-layer projections for column/row-parallel
linears and bind this rank's shard from the full checkpoint.

``state`` is ``None`` on the warm flat-dump path: the parallel
modules are still built (they carry this rank's shard shapes) but
nothing is bound — the dump fills them.

Fused upcast + sparse candidate bias for the DFlash2 sampled walk.

One position of the DFlash2 selector walk biases exactly ``top_k`` columns of
a base logit row and leaves the rest at their base logit. Materialising that
as a full-vocab fp32 copy followed by a ``scatter_add_`` reads and writes the
row twice; :func:`upcast_with_sparse_bias_` does both in one pass over the
row: each column is upcast to fp32 and, where it is one of the position's
candidates, the candidate's bias is added — the same single fp32 add the
scatter performs, so the result is bit-identical.

``out = fp32(base_row); out[r, cands[r, k]] += bias[r, k]`` in one pass.

Falls back to the copy + ``scatter_add_`` pair off CUDA or without Triton
— the same values, two passes.

DFlash2 draft-head primitives (transformers-free).

Two novel modules a DFlash2 checkpoint adds on top of the v1 draft:

* :class:`GroupedDynamicCausalConv` — a depthwise causal conv applied
  around attention and the MLP in every layer. Each group of channels is
  convolved with a static ``base_kernel`` plus an input-dependent kernel
  produced by ``kernel_projection``. ``prepare`` runs the pre-sublayer half
  (kernel index 0) and returns the post-sublayer dynamic kernel; ``finish``
  applies it (kernel index 1).
* :class:`CandidateSelector` — re-ranks the block by walking a predecessor
  path: for each position it scores the ``selector_top_k`` argmax
  candidates against the realized predecessor via low-rank codebooks and a
  hidden projection, then takes the argmax successor.

Both are faithful ports of z-lab's reference modules; they depend only on
torch (not transformers), so the arch/module port is self-contained.

Depthwise causal conv with a static ``base`` kernel plus a per-token
``dynamic`` kernel, over ``kernel_size`` causal offsets.

``hidden`` ``(B, L, H)``; ``base`` ``(kernel_size, H)``; ``dynamic``
``(B, L, kernel_size, groups)``. Channels split into ``groups`` blocks of
``group_size``; offset 0 is the current token, offset ``k`` the token
``k`` steps back (causal left-pad).

Grouped dynamic causal conv wrapped around a sublayer.

``base_kernel[0]`` / ``base_kernel[1]`` are the pre- / post-sublayer
static kernels; ``kernel_projection`` produces both dynamic kernels from
the input in one GEMM. Parameter names match the checkpoint
(``base_kernel``, ``kernel_projection``).

Predecessor-path re-ranker over the block's top-k candidates.

For each block position it scores the ``top_k`` argmax candidates
against the realized predecessor token — ``unary`` (the candidate's own
logit) plus a bilinear term ``(pred_codebook[pred] * hidden) · succ_codebook[cand]``
— then greedily takes the successor, which becomes the next position's
predecessor. ``anchor_ids`` seeds the walk (block slot 0).

GREEDY ONLY, and the served drafter does not walk here at all: the engine
realizes a DFlash2 block through
:class:`~arbi_serve.spec_decode._dflash_driver_dflash2._DFlashDflash2Mixin`,
which holds both regimes device-only and fixed-shape. This module is the
reference port the parity test scores that realizer against.

Return ``(path, candidates, sample_probs)`` for the GREEDY walk.

``path`` ``(B, L)`` is the selected token per block position;
``candidates`` ``(B, L, top_k)`` the argmax pool. ``sample_probs`` is
the reference selector's third slot and is always ``None`` here,
because this port realizes the greedy walk only.

``temperature`` must be ``0``. A sampled DFlash2 block is drawn by
:meth:`~arbi_serve.spec_decode._dflash_driver_dflash2._DFlashDflash2Mixin._dflash2_sample_block`
instead, and the difference is not an optimization: that path biases the
FULL base row and draws it through the same processed-logits chain
(temperature → top_k → softmax → top_p → min_p) the verify pass builds
``p_target`` on, which is the identity standard rejection sampling needs
to stay lossless. A softmax over the top_k pool at raw temperature — the
obvious thing to write here — is a different proposal ``q`` on a
different manifold, so it is refused rather than offered.

Cudagraph capture for the DFlash draft forward.

The DFlash drafter proposes a whole block in ONE parallel
:meth:`DFlashDraftModel.forward_block_fixed` call — unlike the bundled
MTP head / external drafter, there is no autoregressive K-step chain, so
this does NOT use the K-step chain scaffold
(``arbi_serve.runtime.capture.drafter``). It is a single-graph capture:

  * Persistent INPUT buffers (masked-block noise embedding, per-slot
    ``start_pos``, and per-layer padded context K/V + key positions) and
    one OUTPUT buffer (post-norm block hidden) are allocated ONCE per
    per-layer-context-capacity at the widest captured batch and never
    move; each ``(batch, caps)`` graph takes leading-dimension
    ``narrow`` views of that set (:class:`_DFlashGraphBuffers`).
  * The forward is recorded once against those addresses; every later
    step of the same shape ``copy_``\s fresh contents into the input
    buffers and calls ``graph.replay()`` — no re-trace of the 5-layer
    forward, no per-step Python driver overhead.

The capture stays greedy-bit-identical: the captured body is exactly the
eager ``forward_block_fixed`` over stable buffers.

Bucketing: the per-layer context capacity ``Ccap_i`` is rounded UP to the
next power-of-two step so a handful of graphs cover the whole run instead
of one per exact length (the pad columns are masked to zero softmax
weight by ``forward_block_fixed``, so over-padding is numerically inert).
A shape with no captured graph replays nothing (returns ``None``) and the
caller falls back to eager — never a deadlock.

Round a context length up to a coarse bucket (multiple of 64, then
power-of-two beyond 256), clamped to ``hi`` when given so a bucket never
exceeds that per-layer slab width. ``hi >= n`` for any live width, so the
clamped bucket still covers ``n``.

One max-context ``(B, caps)`` per batch width ``B`` in ``1..max_batch``:
``caps`` is the per-layer slab width, the widest context that layer holds.

A live context replays into its batch's single graph by padding up to the
slab (the pad tail is masked). Bounding the pool to ``max_batch`` graphs —
rather than one per context bucket — is what keeps the captured
working-set count independent of context length.

Draft-graph I/O allocated once at the widest batch, sub-viewed per graph.

The pre-sweep captures one graph per batch width ``1..max_batch`` at the
same per-layer capacities, and exactly one of them replays per step. A
per-graph allocation therefore holds ``sum(1..max_batch)`` copies of a
buffer set whose widest member already covers every narrower one: the
context K/V pair alone is ``B x nkv x cap x head_dim`` per layer, the
dominant draft-graph term.

Allocating at ``width`` and handing each graph ``narrow(0, 0, B)`` views
keeps every graph's baked addresses stable and byte-identical to a
standalone allocation — a leading-dimension narrow of a contiguous tensor
is contiguous with the same strides — while the resident bytes are those
of the widest graph alone.

Sound because the buffers carry NO state between steps:
:meth:`DFlashDraftGraph.replay` rewrites ``start_pos`` / ``noise`` and
every layer's context rows it uses, and re-marks the key-position pad with
``-1``. Sharing across graphs is therefore invisible to any single replay,
and replays never overlap (one captured graph runs per step). ``out`` is
handed back as a clone unless the caller asks to borrow it, and a borrower
owes consumption before the next replay of ANY graph at these caps — see
:meth:`DFlashDraftGraph._returned`.

The pre-capture warmup passes for one ``(B, caps)`` draft shape.

Two forwards, so any first-call init that fires on call #1 rather than
call #2 is settled before anything records, then a device sync. Under
``ARBI_DFLASH_COMPILE`` these are also what drives Dynamo + Inductor for
this shape, which is why the warmup is a function both callers share:
:class:`DFlashDraftGraph` runs it immediately before recording, and
:meth:`DFlashDraftGraphPool.warm_shapes` runs it alone from the
compile-warmup phase. A shape warmed by one is warm for the other only
while both drive the SAME call — hence one body, not two.

Graph-pool ceiling that SCALES with the served ``max_batch``.

DFlash captures one graph per DISTINCT ``(B, caps)`` it sees, at the exact
slate row count ``B`` — so the reachable shape space grows linearly with
``max_batch`` (up to ``max_batch`` distinct widths). A fixed cap
under-covers at high concurrency: at c64 a steady shape set can exceed 32
and new shapes fall to the EAGER draft (a loud warn, then a per-step perf
loss on the drafter path). Deriving the ceiling from ``max_batch`` engages
fully at whatever concurrency is launched — the 27B/TP2 (mb16) ceiling
stays at 32 (mb16×2 floored at 32) so the tight card pays no extra
draft-graph VRAM, while the roomy 0.8B/TP1 (mb64 → 128) path widens to
cover the extra slate widths a fixed cap would drop to eager at c64. The
pool only ever ALLOCATES a graph for a shape that ACTUALLY occurs, so
VRAM tracks the real workload (a bench at one context length fills only a
few slots), not the ceiling. Floored at ``_DFLASH_MIN_GRAPHS`` so no
config pays more draft-graph VRAM than that floor. ``ARBI_DFLASH_MAX_GRAPHS``
is an explicit positive override (bound it on a tight card, or lift it
for a context-diverse workload).

Lazy cache of :class:`DFlashDraftGraph` keyed by (B, capacity bucket).

``max_graphs`` bounds the captured-shape count; derive it from the served
``max_batch`` via :func:`dflash_pool_max_graphs` so it scales with launched
concurrency (the reachable ``(B, caps)`` space is ``max_batch``-wide). The
``32`` default is only the module-test fallback — the driver always passes a
max_batch-derived value.

``graph_pool`` is the MemPool the captured graphs record into;
``buffers_pool`` is the SEPARATE MemPool holding the persistent
:class:`_DFlashGraphBuffers` the graphs bake by ``data_ptr``. They must not
be the same pool — the capture allocator reuses addresses across
non-overlapping replays, so a persistent buffer placed in it can land at an
address an earlier graph already baked into its kernel arguments (see
:func:`arbi_serve.runtime.capture.decode.capture_decode`).

Copy fresh inputs into the persistent buffers, replay, return the
output — a clone by default (a view would alias the next replay), or
the persistent buffer itself under ``borrow_output``
(:meth:`_returned`).

The replay result: the persistent :attr:`out` buffer when the
caller borrows it, else a fresh clone.

BORROWING IS A CONTRACT ON THE CALLER. :attr:`out` is a view of the
buffer set every graph at these caps shares, so the NEXT replay of
ANY graph overwrites it. A borrower must consume the tensor — and
drop every view of it — on the same stream before that next replay.
Every replay is counted (:attr:`replays`) so a borrower can assert
the count did not move while it held the result.

Replay over inputs the caller already wrote into this graph's own
persistent buffers, and return the output (a clone, or the borrowed
persistent buffer under ``borrow_output`` — see :meth:`_returned`).

The sibling of :meth:`replay` for a caller that assembled straight
into :attr:`start_pos` / :attr:`noise` / :attr:`ctx_k` / :attr:`ctx_v`
/ :attr:`ctx_kpos` — there is nothing left to copy, so the whole
second pass over the context is skipped. The caller owes exactly what
:meth:`replay` writes: every row of ``start_pos`` and ``noise``, a
``-1`` refresh of each layer's whole ``ctx_kpos``, and the live
columns of each layer's ``ctx_k`` / ``ctx_v``. Columns and rows left
untouched carry the allocation zero-fill or an earlier step's draft
K/V — finite either way, and masked to exactly zero softmax weight by
the refreshed ``-1`` key positions.

Writing these buffers outside the capture region is safe: they are
ordinary device memory whose addresses the recorded kernels baked in,
the writes and ``graph.replay()`` are enqueued on the one stream in
program order, and exactly one graph replays per synchronous
:meth:`DFlashDrafter.draft` call — no replay is ever in flight while
they are written.

The bucket tuple for the live per-layer context WIDTHS, clamped to
the slab widths so it matches the pre-sweep's keys exactly.

Takes the widths themselves rather than the tensors that carry them,
so a caller that derives them from
:class:`~arbi_serve.spec_decode.dflash_kv_slots.DraftKVSlots` layout
arithmetic can pick its graph BEFORE assembling anything into it.
:meth:`_caps_for` is this function applied to ``shape[2]``, so the two
agree by construction, not by coincidence.

The set every pre-swept graph at ``caps`` views.

``None`` outside an announced sweep, or for a width past it — those
graphs allocate their own set, exactly as an unshared pool does.

Run each shape's pre-capture warmup forward WITHOUT recording it.

The compile half of :meth:`precapture_shapes`, hoisted so it can run
in the compile-warmup phase instead of inside the capture sweep. Same
shapes, same shared buffers and same warmup call, so the sweep's own
warmup then finds every compiled callable already built and the
Inductor artifacts are in the process before the boot persists its
compile-cache blob. Returns the number of shapes warmed.

Announces the sweep width exactly as :meth:`precapture_shapes` does,
so the buffer set allocated here is the one the sweep views — warming
early moves those bytes earlier in the boot, it does not add any.

The graph that serves shape ``(B, caps)``; ``None`` ⇒ run eager.

The pool's whole shape-selection policy, over the SHAPE alone: a
caller holding no assembled tensors can pick its graph here and then
assemble straight into that graph's persistent buffers. A sealed pool
reuses a covering graph or misses (one warn per shape); an unsealed
pool captures on first sight until it is full.

The returned graph is always at batch width ``B`` exactly (both the
exact-key hit and :meth:`_covering` filter on it), so its buffers carry
``B`` rows.

True when :meth:`resolve` returns a graph for EVERY shape in ``shapes``.

Asked with the pre-sweep's own shape list
(:func:`dflash_presweep_shapes`), this is the statement *no served
draft step can miss this pool*. The sweep captures at the per-layer
SLAB widths, which are the widest context any layer can hold, so a
live step's per-layer caps are ``<=`` the swept caps componentwise and
:meth:`_covering` — which accepts any resident graph at the same batch
width whose every cap is ``>=`` the asked one — resolves to the swept
graph. Coverage at the slab caps for every batch width therefore
implies coverage at every reachable width, by construction rather than
by enumeration.

It is the same :meth:`resolve` a step calls, so a True here cannot be
true by a different rule than the one serving runs.

An UNSEALED pool answers False without asking: :meth:`resolve` would
MINT a graph for a shape it does not hold, so on an unsealed pool the
question captures rather than reports, and a pool that still captures
on demand has no fixed coverage to prove.

Return the captured hidden for this shape (``replay`` pads a live
context up to the bucket internally); ``None`` ⇒ caller runs eager.

``borrow_output`` hands back the graph's persistent output buffer
under :meth:`DFlashDraftGraph._returned`'s contract.

Every captured graph this pool holds.

The engine's sleep walk introspects DATACLASS entries of iterable
pools; this pool keeps :class:`DFlashDraftGraph` objects in a plain
dict, so it declares them instead. Without this the draft graphs'
instantiated execs — driver memory outside every torch and cuMem pool
— stay resident through every sleep.

True iff ``c`` is a DFlash2 draft checkpoint config.

Detected by the declared architecture or by the presence of any of the
conv/selector hyper-parameters under ``dflash_config`` (or top level),
so a v2 checkpoint is recognised whichever on-disk layout it uses.

Map checkpoint module prefixes to the runtime DFlash module paths.

DFlash-family repositories share one computation graph but may expose
architecture-owned names for the context encoder.  The mapping is
selected by the checkpoint's declared architecture and is intentionally
narrow so unknown layouts fail strict loading.

Return ``state`` with architecture-owned prefixes made canonical.

DFlash2 additionally stores its selector codebooks as bare tensors
(``candidate_selector.predecessor_codebook``); the runtime holds them
as ``nn.Embedding`` so their canonical key carries the ``.weight``
suffix (the reference does the same via ``from_pretrained`` key_mapping).

True iff the checkpoint config declares per-head attention sinks.

Read from the ``dflash_config`` sub-dict when it carries the key,
otherwise from the top level — the same two on-disk layouts
:func:`dflash_config_fields` handles.

Extract ``(mask_token_id, target_layer_ids)`` from a DFlash-family
``config.json`` dict.

Two on-disk layouts exist: z-lab DFlash checkpoints
(``architectures: [DFlashDraftModel]``) nest the fields under a
``dflash_config`` sub-dict; DSpark heads
(``architectures: [Qwen3DSparkModel]``) put them at the top level.
``target_layer_ids`` is mandatory (its absence means "not a DFlash
checkpoint"); ``mask_token_id`` is ``None`` when absent — the boot
asserts don't need it, the model loader (:meth:`DFlashConfig.from_dir`)
does and fails loud there. A ``causal`` / ``causal_head`` flag is NOT
read: the block forward is bidirectional for every checkpoint, so the
flag selects nothing.

Extract ``(output_multiplier, final_logit_softcapping)`` from a
DFlash-family ``config.json`` dict.

Both are epilogue terms the head was trained to emit its logits
through, and the drafter's vocab projection must reproduce them: the
sampled draft reports the post-epilogue distribution as its ``q``, and
rejection sampling is lossless only when ``q`` is the distribution the
head defines. Dropping either silently costs acceptance AND breaks
that identity, so they are parsed rather than defaulted away.

Read as one group from one source block — nested ``dflash_config``
when it declares either, else the top level — matching how
:func:`dflash_config_fields` and the DFlash2 conv/selector group pick
their source. Absent from both is the identity epilogue
``(1.0, None)``, which projects raw logits.

A cap of ``0`` means "no cap" (the transform is undefined at zero) and
normalizes to ``None``. A negative cap or a non-positive multiplier is
refused: neither has a meaning the epilogue can honour, and applying
one would silently invert or flatten the head's distribution.

Resolve the draft block width (anchor ++ ``block_size - 1`` masks).

A ``block_size`` declared by the checkpoint always wins — it is that
head's trained width. A checkpoint that declares none takes the runtime
width the server derives from ``--mtp-n-draft``, bounded by
:data:`UNDECLARED_MAX_BLOCK_SIZE`; with neither it refuses.

DFlash2 nests its trained ``block_size`` under ``dflash_config`` and is
read from there. v1 checkpoints keep the existing top-level-only lookup
(their nested ``block_size``, when present, is intentionally
runtime-owned — see ``build_dflash_drafter``).

Parse ``config.json`` from a DFlash / DSpark checkpoint directory.

``runtime_block_size`` is the block width the server derived from
``--mtp-n-draft``; it is used only by checkpoints that declare no
``block_size`` of their own (see :func:`dflash_block_size`).

``target_rope_scaled`` is whether the TARGET applies a RoPE scaling
schedule. The draft head consumes the target's hidden states and
predicts its tokens, so the two must rotate positions identically;
a schedule the target does not apply is dropped and reported.

DFlash block-diffusion drafter — engine binding (:class:`Drafter` sibling).

Wraps :class:`arbi_serve.spec_decode.dflash.DFlashDraftModel` behind the
:class:`arbi_serve.spec_decode.drafter.Drafter` Protocol so the existing
MTP verify machinery (:func:`run_verify_step`, ``MtpStrategy``) drives it
unchanged. One ``draft(...)`` call denoises a whole masked block in a
single parallel forward — no autoregressive chain — conditioned on hidden
states tapped from several target-model layers.

Draft K/V management: tapped context features are projected to per-layer
draft K/V **at observe time** (:meth:`DFlashDraftModel.project_context`)
and written straight into :class:`DraftKVSlots` — persistent per-request
K/V buffers. Only committed positions are ever written, so accept/reject
"rewind" does not
exist: the block's own K/V are scratch inside the draft forward and never
persist. No feature buffer, no growing python-list cache, fixed base
pointers (cudagraph-ready).

Invariant: a request's slot length always equals its anchor position (the
last committed token, not yet forward-processed). The seed-path position
check disables speculation for a request that violates it (e.g. a
prefix-cache hit skipped part of the prompt); verify keeps the output
correct either way.

The drafter realizes plain DFlash blocks in parallel. The shared
``ARBI_TRUE_STOCHASTIC_DRAFT`` policy selects greedy argmax with the compact
point-mass ``q`` marker or sampled realization with the exact per-position
``q`` reported to the rejection sampler.

DSpark checkpoints (``markov_rank > 0``) add a semi-autoregressive
re-ranking pass over the same single parallel forward
(``ARBI_DFLASH_MARKOV_SEMIAR``): the block's base logits are produced
once, then the block is realized LEFT-TO-RIGHT with a low-rank bigram
bias conditioned on each position's realized predecessor
(:class:`~arbi_serve.spec_decode.markov_head.VanillaMarkov`). Greedy
slates argmax the biased logits (lossless — verify re-checks every
token); under ``ARBI_TRUE_STOCHASTIC_DRAFT`` a stochastic slate SAMPLES
each position from its own per-row filtered distribution and reports
that exact per-position ``q`` (the same
``GraphSafeRejectionSampler.sample_drafter_token`` contract the bundled
MTP head's true-stochastic chain uses), which keeps standard rejection
sampling lossless by construction.

TP>1: the 5-layer draft model + per-request draft-KV state live on rank
0; worker ranks join the two vocab-parallel collectives
(``embed_tokens`` all_reduce, ``lm_head`` all_gather) in lockstep via
``DFlashDraftOp`` so rank 0's borrowed-module forwards don't deadlock.
Self-contained heads (checkpoint-owned dense embed/lm_head — the DSpark
layout) carry NO vocab-parallel collective in either leg, so those
broadcasts are skipped.

``cfg.batch.max_batch`` for the draft graph-pool ceiling.

``max_batch`` is concrete from CLI parse onward, so a missing or ``< 1``
value here is a config-plumbing bug — FAIL LOUD. Substituting a concrete 1
would silently under-cover the graph pool while reporting a max_batch the
operator never configured.

The target's input-embedding module, across architecture layouts.

Raises when no known layout matches: a borrowing head silently reading the
wrong module would draft from a different embedding than the target
verifies with, which costs acceptance without failing.

Width of the vocabulary-prefix head this drafter may project through.

0 when ``--draft-vocab-prefix`` is off. Raises when it is ON but this
drafter cannot honour it — never silently ignores it, because the operator
named a width and a drafter whose reach they cannot read off their own
config is a surprise, not a lever.

Checked HERE, at drafter construction, because it is the first point that
can see both facts: the sliced head is installed before ``_profile_and_size_kv``
(so its VRAM is charged to the KV budget) and this drafter is not built
until after it, inside the KV phase.

Bind the loaded draft model to its target engine.

At TP>1 the draft model + per-request draft-KV state live on RANK 0
ONLY; worker ranks join rank 0's ``embed_tokens`` / ``lm_head``
collectives in lockstep via ``DFlashDraftOp`` (see :meth:`draft`).

Append ``(conf_prob, accepted)`` calibration rows (JSONL) when
``ARBI_DFLASH_CONF_DUMP`` is set. Each verified row's stashed
per-position confidence (:attr:`Request.mtp_next_draft_conf`) is
joined to that step's accepted-prefix length (``res.num_accepted``):
position ``j`` was accepted iff ``j < num_accepted``. Diagnostic
only (gated OFF by default); used to measure the confidence head's
accept-prediction AUC before trusting dynamic-K.

Arm cudagraph capture of the draft forward.

Builds a :class:`DFlashDraftGraphPool` that records
:meth:`DFlashDraftModel.forward_block_fixed` once per (batch,
context-capacity bucket) shape and replays it thereafter. Called
from :func:`build_dflash_drafter` when ``ARBI_DFLASH_CAPTURE`` is
set (GPU only). The eager fixed-shape forward stays the fallback
for uncaptured shapes / capture failure. Idempotent.

``graph_pool`` holds the captured graphs; ``buffers_pool`` holds the
persistent draft-graph I/O the graphs bake by ``data_ptr``. They are
different pools by design — see :class:`DFlashDraftGraphPool`.

Arm capture of the block realization
(:mod:`arbi_serve.spec_decode.dflash_realize_capture`) when
``ARBI_DFLASH_WALK_CAPTURE`` is on. Idempotent.

The gate is the WALK, not the draft regime. Both regimes issue the
same left-to-right per-position walk from the host — a sampler chain
for a sampled slate, an argmax re-rank for a greedy one — so arming
on "sampled drafting is reachable" would leave the greedy walk eager
and charge a boot that never samples a dispatch cost that has nothing
to do with sampling. The two are then not comparable on a served A/B,
which is the whole point of running one.

How many realization graphs the boot pre-sweep records: one per
slate width per mask-plan shape, plus one per slate width for the
greedy walk where the served realizer has one. Zero when the pool is
not armed.

Boot pre-sweep of the block realization at every slate width: the
greedy walk, and the sampled walk for the top-k-capped and the
nucleus mask plans. Then seal the pool.

The greedy sweep runs first and unconditionally — a greedy walk reads
no proposal slab, so it does not share the sampled sweep's dependency
on the engine having one to bake. Without that slab the sampled sweep
is skipped and the pool stays UNSEALED, capturing on first use.

The draft-forward graphs this drafter owns, for the sleep walk.

The engine's walk introspects dataclass entries of its cudagraph
pools; the draft pool is not one of those, so the ownership is
declared here (see ``arbi_serve.engine.sleep._GRAPH_OWNER_ATTRS``).
Empty when capture is not armed.

The ``(B, caps)`` shapes the boot pre-sweep captures.

The ONE source both the capture sweep and the compile warmup read.
Two callers deriving the same list independently is how a warmed shape
set and a captured shape set drift apart, which would leave the sweep
compiling a shape the warmup never saw — the cost this hoist exists to
remove, silently back where it started.

Drive the draft forward at every pre-sweep shape, capturing none.

Called from the compile-warmup phase. Under ``ARBI_DFLASH_COMPILE``
each shape's first call is a Dynamo trace plus an Inductor build, and
running them here rather than inside the capture sweep puts them
BEFORE the boot persists its compile-cache blob, so a later boot
reloads them instead of rebuilding them. With the flag off the calls
are eager and cost one warm forward per shape.

Returns the number of shapes warmed; 0 when capture is not armed.

Capture the draft forward at every ``(B, caps)`` bucket a served
request can reach, then seal the pool. Run from the Phase-5 capture
sweep so the pool is fully grown before the KV serving ceiling is
sized; a live context replays into its bucket by inert padding. No-op
when capture is not armed.

True when the step's active rows are exactly ``range(B)``.

``active`` enumerates ``requests`` in order, so ``done`` is strictly
increasing and a full-length one can only be ``range(B)``. Under it
every per-row gather and scatter the step would issue — the drafted
hidden, the base logits, the anchor ids, the token block — is the
IDENTITY, so the step takes views instead and writes ``out`` directly.
A partial slate keeps the padded destinations, whose inactive rows
must read zero.

Active-row indices on the drafter device, staged through the
persistent pinned host mirror.

Returns a view of the persistent device buffer — valid until the
next :meth:`_stage_rows` call, so one draft step stages once and
reuses the view. The numpy assignment writes the row ids straight
into pinned storage (NumPy's C-level coercion — no intermediate
host tensor), and the H2D is a genuine async
``copy_(non_blocking=True)`` off pinned memory.

Keyed-watermark context for a left-to-right sampled block, or None.

``(cache, row-ordered requests, seed_dev)`` when watermarking is armed —
the shared skeleton both re-rank heads (:mod:`._dflash_block_walk`) pass
through to couple each position's draw to the keyed verify. None when
the watermark is not in play: non-CUDA, an SPMD worker bridge (the
mirror carries no ring, so TP1-only), or no watermark cache present.

The logit epilogue governing the drafter's vocab projection, as
``(active, apply_logit_epilogue keywords)``. Resolved once and cached.

The epilogue belongs to the OWNER of the head being projected
through, because it is part of what that head's logits mean. A
BORROWED head (``self.model.lm_head is None``) is the target's, so the
terms are the target's, asked of the model that resolved them
(:meth:`~arbi_serve.models._layer_stack_model_mixin.LayerStackModelMixin.logit_epilogue_terms`)
rather than re-derived here. A self-contained head (DSpark) is the
checkpoint's own, so its ``config.json`` governs.

Ownership of the lm_head is a SEPARATE question from ownership of the
input embedding (``draft_head_modules``' third element): a checkpoint
may ship ``embed_tokens`` and still project through the target's head.

Skipping this is not a losslessness bug — the sampled walk reports
the ``q`` it drew from either way — but it costs ACCEPTANCE: the
drafter proposes from a differently-shaped distribution than the
target scores with, and every such mismatch shows up as rejected
tokens.

The narrowed ``--draft-vocab-prefix`` head changes nothing here: all
three terms are elementwise and monotone, so applying them to the
``N``-wide row is exactly the prefix of applying them to the full one.

Resolution is lazy because a checkpoint-owned head is attached
after ``__init__``, so ownership is not yet known there. The cached
result makes the per-step guard in :meth:`draft` a host-side branch
on a constant — no device op, and nothing data-dependent.

The embed / lm_head modules ONE draft step projects through.

Embed / lm_head ownership: a self-contained head (DSpark) ships its
own dense modules and the draft must use THOSE (the head is trained
against its own copies); a borrowing head (z-lab DFlash) uses the
target's — under TP those are vocab-parallel modules whose collectives
the workers join, which is why the third element says which it is.

Resolved HERE for every caller rather than at each site: the serving
floor measures a draft step's peak by running this same pair
(:func:`~arbi_serve.engine.inprocess_capture._measure_dflash_draft_step_bytes`),
and a floor sized against a different head than the one that serves is
a floor sized against a different vocab. That is why the
``--draft-vocab-prefix`` window is resolved here too: the drafter's
projection is ``N`` wide on every path that runs it, so the floor
measures the row the step actually allocates. What stays at the SERVED
call site is :func:`~arbi_serve.draft_vocab_prefix.note_draft_step` —
a per-step attribution counter only a served step may fire.

One block-diffusion denoise pass per row → ``(k, B)`` drafts.

Rows without context (disabled / no slot) get zero-token drafts
(the verify pass rejects them; correctness is unaffected, only
accept-length).

``frontier_repair_meta`` is ignored — it overrides the bundled MTP
head's read of the MAIN pool's KV frontier under the async verify
path. DFlash owns its OWN draft-KV slots (stashed by the observers /
the device stash), so the main-pool repair meta does not apply.
Accepted for signature parity so the async seed's
``draft(..., frontier_repair_meta=...)`` call does not raise.

Sample an independent DFlash block into ``q_out`` and return it.

Every position uses the shared drafter sampling transform and a
distinct counter-keyed draw
(:func:`~arbi_serve.spec_decode._dflash_block_walk.parallel_sample_block`).
Returns tokens as ``(Bd, S)`` and the caller's ``q_out`` destination,
which the verify rejection sampler reads as the proposal. Served by
the captured realization when one covers the shape.

Realize the greedy block by captured-graph replay, or ``None`` for
eager — with the reason on ``dflash_realize_captured``.

A greedy walk carries no slate: none of the sampled route's
conditions (mask plan, ``min_p``, the baked ``q`` slab, a keyed
watermark, the fused-draw and philox gates) apply to it, so the only
thing that can refuse it is a shape the pool has no graph for.

Which block realizer a slate takes on this drafter, in either
draft regime — the greedy walk and the sampled walk differ only in
how a position is resolved, never in which head re-ranks it.

Realize the block by captured-graph replay, or ``None`` for eager.

``None`` — with the reason on ``dflash_realize_captured`` — when no
pool is armed, the slate is keyed-watermarked (its per-position seeds
are host work), the drafter sampling chain is not on its tensor-driven
routing, or the pool has no graph that can represent this slate.

The done rows of a ``(B, ...)`` tensor.

``active`` enumerates ``requests`` in order, so an all-active
slate's row indices are ``range(B)`` and the gather is the
IDENTITY — the tensor itself is the answer, and the copy the
gather would make is pure duplicate traffic (a full-vocab one
on the base logits). A partial slate gathers as before.

Maskless varlen attention for the DFlash draft forward.

The draft forward's sliding layers express two constraints — the
architecture's ``|k_pos - q_pos| <= window`` rule and the per-row validity
of the ring's pad columns — as ONE additive ``attn_mask``. Every fused
SDPA backend that would accept the shape declines on the mask's presence
(``sdp_utils_cpp.h``: flash rejects any ``attn_mask``), so the draft
attention lands on the memory-efficient kernel;
``docs/dflash-drafter-attention-decomposition.md`` reports the measured
gap and the method that produced it.

This module expresses BOTH constraints without a mask tensor:

* the window becomes the flash kernel's native ``window_size``, which is
  bottom-right aligned — key index ``j`` attends query index ``i`` iff
  ``-w <= j - (i + seqlen_k - seqlen_q) <= w``. Laying each row out as
  ``[its in-window context in ascending position] ++ [the block]`` makes
  index distance EQUAL position distance, so the native window computes
  the architecture's rule exactly;
* pad validity becomes ``cu_seqlens_k``: a row contributes only the keys
  it actually holds, so there is no pad column for a mask to zero.

The context lives in a RING (:class:`~arbi_serve.spec_decode.dflash_kv_slots.DraftKVSlots`),
whose valid span wraps, and index-distance equals position-distance only
in ascending-position order. :class:`DFlashVarlenPlan` therefore builds,
per step, the permutation that lands the ring's span contiguously —
inverted from the same ``ctx_kpos`` the masked route reads, so the two
routes attend the same key set by construction rather than by agreement.

Capture safety: every per-step quantity is derived on the device from
``start_pos`` / ``ctx_kpos`` with no ``.item()``, no host branch and no
data-dependent shape, into buffers the plan allocates once per
``(batch, width, window)`` — the same discipline
:mod:`arbi_serve.spec_decode._dflash_driver_context` holds for the
context assembly it feeds.

turbo-attn's cute-DSL varlen prefill, driven through its flat loader.

``turbo_prefill`` carries the same bottom-right window definition as FA2
(``tkv/kernels/cute/_fa/mask.py``, the Local branch:
``row + causal_row_offset - 1 - window_size_left <= col <=
row + causal_row_offset + window_size_right``) and takes
``cu_seqlens_q`` / ``cu_seqlens_k`` directly. ``BypassLoader.from_bf16``
in FLAT mode reads a ``(1, total_k, n_kv, head_dim)`` tensor at
``cu_seqlens_k[row] + n``, which is exactly the packed layout
:class:`DFlashVarlenPlan` produces.

The importable varlen+window provider, and why, memoized.

Returns ``(provider, reason)``; ``provider`` is ``None`` when nothing
importable offers varlen with a native window, and ``reason`` then names
every candidate that was tried so the refusal states what the image
lacks rather than that something is missing.

Validate ``requested`` against what this build can actually run.

``sdpa`` always resolves. ``flash_varlen`` RAISES when no provider is
importable rather than degrading silently: the flag is capture-affecting
and sizes nothing, so a silent downgrade would leave a member captured
on one route while the operator reads the other.

Per-step varlen geometry for one ``(batch, context width, window)``.

Every buffer is allocated once, at construction, and rewritten in place
by :meth:`refresh` — the draft forward is cudagraph-captured, so a
per-step allocation would be baked into the graph's private pool and a
host read would break the capture outright.

``span`` is the widest context any row can contribute: ``min(width,
window)`` for a windowed layer (a query cannot reach further back than
the window) and the full ``width`` otherwise. The packed row stride is
``span + block``, so ``batch * (span + block)`` slots hold any slate.

``window`` is carried because it decides where a row's packed span
STARTS, and the two layer kinds answer differently: a windowed layer
starts at the window's own lower bound (the ring may hold older keys the
window excludes), a full-attention layer at the oldest key it actually
holds (there is no window to bound it, and a ringed full layer's oldest
key is not position 0).

Rebuild the step's spans and permutation from the live context.

``ctx_kpos`` is the same ``(batch, width)`` absolute-key-position
tensor the masked route reads (``-1`` on a pad column), so a key this
plan packs is exactly a key that route would have unmasked.

A windowed layer's lower bound is ``max(0, committed - span)``: a
query at the anchor reaches back at most ``span`` positions, and the
ring holds at most ``span`` of them, so the two agree on the same
first position. A full-attention layer has no such bound and takes
the oldest position the row actually holds instead — which is
position 0 for a linearly stored layer and the ring start for a
capped one.

Gather ``ctx`` into packed varlen order and append the block.

``ctx`` is ``(batch, n_kv, width, head_dim)`` in ring order;
``block_kv`` is ``(batch, block, n_kv, head_dim)``. Returns
``(total, n_kv, head_dim)`` — one gather that moves the same bytes
the masked route's ``torch.cat`` moves, plus one scatter of the
block's own keys.

Persistent per-slot draft K/V storage for the DFlash drafter.

:class:`DraftKVSlots` — the fixed-base-pointer per-request draft K/V
buffers that :class:`~arbi_serve.spec_decode.dflash_driver.DFlashDrafter`
writes committed context into. Re-exported from ``dflash_driver`` so
``from arbi_serve.spec_decode.dflash_driver import DraftKVSlots`` still
resolves.

Persistent per-slot draft K/V storage, sized per layer-type.

Each draft layer owns one persistent ``(max_slots, n_kv, layer_len,
head_dim)`` K and V buffer whose base pointer never moves. A request
owns one slot for its lifetime; ``lengths[slot]`` is the number of
committed context positions written (the request's absolute context
length — it may exceed a sliding layer's ring capacity). Appends take
rope'd / normed K and raw V from :meth:`DFlashDraftModel.project_context`.

Per-layer-type sizing: a ``full_attention`` layer attends the whole
context, so its buffer holds ``effective_context + block_size``
positions. A ``sliding_attention`` layer is causal with a trailing
``sliding_window``, so the block queries (at absolute positions
``>= anchor``) can only ever attend context positions in
``(anchor - window, anchor]`` — the last ``window`` committed
positions. Its buffer therefore holds only ``min(effective_context,
sliding_window) + block_size`` positions, stored as a RING: position
``p`` lands at row ``p % ring_len``. Storing only the window is
numerically equivalent (out-of-window context is masked away — see
``_swa_mask`` and ``test_dflash_swa``), which keeps a sliding layer's
buffer far smaller than a full-attention layer's at long context.

The draft model runs EAGER (never cudagraph-captured), so the ring's
dynamic / wraparound indexing is safe.

Scatter each row's first ``n_keep`` committed K/V into its slot,
DEVICE-RESIDENT (no host ``.item()``), advancing ``lengths_dev``.

Mirrors :meth:`append` byte-for-byte (full layers store linearly at
``[base, base+n_keep)``; sliding layers ring at ``pos % ring_len``)
but takes a FIXED per-row span ``S`` and a device ``n_keep`` mask so
the whole write is one masked ``index_put_`` per layer. ``S`` is the
verify span (``K+1``); positions ``>= n_keep`` (rejected / cold pad)
are masked out and never written. Advances ``lengths_dev[slot] +=
n_keep``.

Overflow (a FULL layer's ``base+n_keep`` past its linear capacity —
real context exhaustion) is masked out rather than written, so the
slab is never corrupted; the request's context simply stops growing
(verify keeps the output correct). The host path disables such a
request; the device path lets it ride truncated (accept degrades,
not correctness) — a rare long-context edge, documented.

Cached ``(row, position-in-span)`` index pair for an ``A x S``
append, flattened.

Both are pure functions of the shape, and the shapes an append takes
are bounded — one row per slot, one column per verify-span position —
so they are built once per shape instead of once per append.

Trailing rows of an ``n``-row append that SURVIVE it.

THE SPAN THE SKIP IS ALLOWED TO PROJECT, and it is deliberately the
ring's own number rather than one recomputed from ``sliding_window``.
The justification for dropping rows is "this ring overwrites them", so
the span has to come from the ring or the justification and the code
can drift apart without either looking wrong.

Every layer here stores in a ring of ``_layer_len[i]``, and an append
of ``n >= cap`` rows writes each ring index ``ceil(n / cap)`` times —
the LAST write to every index coming from one of the final ``cap``
rows. So the final contents of an append of all ``n`` rows and of only
the final ``min(n, cap)`` are identical, BIT FOR BIT, by construction.
Note what that argument does NOT use: the attention mask. It is a
statement about storage being overwritten, not about which keys a
query can see, so it holds under a causal window and a BIDIRECTIONAL
one alike (``DFlash2DraftModel`` is bidirectional). Slicing on the
mask instead would give a different — and at the front of a
bidirectional window, wrong — span.

Returns ``n`` unchanged (nothing may be skipped) the moment ANY layer
stores linearly: a linear layer keeps every position it is given, so
no row it would receive is dead. That is the one precondition, and it
is read off :attr:`_has_linear_full` rather than assumed.

The MAX over layers, because the stacked projection produces every
layer from one pass over the same rows: a row is dead only if it is
dead for all of them. Per-layer spans still bound it individually —
the ring discards the excess for the shorter layers exactly as it
always did.

Write per-layer ``(1, n_kv, n, head_dim)`` K/V at the slot's
current absolute length. Full layers store linearly; sliding
layers ring at ``pos % ring_len`` (older entries fall out of the
trailing window and are correctly overwritten). Returns ``False``
(no write) only when a FULL-attention layer would overflow its
linear buffer (true context exhaustion).

``skipped`` is how many positions the caller did NOT project because
:meth:`retained_span` proved this append would overwrite them. Those
positions are still CONSUMED: the write starts ``skipped`` further on
and the slot's length advances by ``skipped + n``. That split is the
whole hazard of the skip — ``lengths`` tracks absolute committed
positions and :meth:`~arbi_serve.spec_decode._dflash_driver_stash.
_DFlashStashMixin.slot_for_span` compares the next span's start
against it, so coupling the advance to the number of rows WRITTEN
would silently re-seed every long request's draft context. The
advance follows the positions consumed, never the tensor's length.

Where one slot's held context for ``layer`` lives in the slab.

Returns ``(start, m, runs)``: ``start`` is the absolute position of
the oldest held entry, ``m`` the number of positions held, and
``runs`` maps them onto slab columns in ascending absolute position
as ``(dst_offset, src_offset, length)`` triples. A linearly stored
layer holds ``[0, n)`` as one run. A ring layer holds the trailing
``m = min(n, ring_len)`` positions, and since a contiguous absolute
range of ``m <= ring_len`` positions wraps the ring at most once it
is one run, or two when it wraps. Pure host arithmetic off
:attr:`lengths` — no tensor read, so a caller can size and fill the
destination without materialising a gather.

Slab columns the device context view of ``layer`` spans.

A ring layer spans its whole (small) trailing window; a full-attention
layer is bounded to ``max_full_len`` when the caller supplies that
sync-free upper bound on the live committed length.

True when ``layer`` stores its context as a trailing RING (position
``p`` at row ``p % layer_len``) rather than linearly.

With :meth:`context_width` it is the whole of a layer's key-position
layout, so two layers agreeing on both hold identical key positions.

``(R, cap)`` absolute key position of each of ``layer``'s slab
columns for rows with committed ``lengths`` ``(R,)``; ``-1`` for a
column the row has not filled.

A row with ``lengths == 0`` (an inactive destination row) yields all
``-1``, which is exactly the pad the fixed-shape forward masks out.

``out`` writes the result into a caller-owned ``(R, cap)`` int64
destination (a capture graph's persistent key-position buffer) and
returns it, so the positions are built where they will be read instead
of into a fresh tensor the caller then copies.

The destination is also the arithmetic's own scratch: the column range
is a cached ``(1, cap)`` constant and every step of the map runs in
place, so a call allocates nothing beyond the ``(R, cap)`` mask —
rather than the three full-width int64 intermediates a functional
``where`` chain builds.

``(1, cap)`` column range for :meth:`context_kpos_device`.

A leading slice of the one persistent range built with the slabs, so a
per-layer key-position build allocates no ``arange`` — at any width,
including the varying full-attention widths a live step asks for.

Per-layer context K/V plus the absolute key positions.

Full layers return the linear ``(1, n_kv, len, head_dim)`` slice
for positions ``[0, len)``. Sliding layers return only the last
``ring_len`` committed positions, REORDERED into ascending
absolute position (so the K — already RoPE'd at its absolute
position — pairs with the right key position), and the matching
absolute positions so :meth:`DFlashDraftModel.forward_block`
builds the SWA mask against true positions (not ``[0..n)``).

DEVICE-resident per-layer context for a batch of slots.

Returns ``(ks, vs, kpos, start_pos)`` where ``ks[i]`` / ``vs[i]``
are ``(A, n_kv, W_i, head_dim)`` slab slices and ``kpos[i]`` is
``(A, W_i)`` the ABSOLUTE key position of each stored column (``-1``
for a column the slot has not filled). ``start_pos`` is ``(A,)`` the
committed length gathered from ``lengths_dev``.

WIDTH BOUND (``max_full_len``): a full-attention layer's buffer is
``effective_context`` wide, but only the first ``max active length``
columns are ever populated. Passing the WHOLE slab makes every draft
step allocate + attend over ``effective_context`` keys (mostly
masked) — memory-hungry and compute-wasteful. Caller passes
the batch's max committed length so full layers slice to
``[0, min(full_len, max_full_len)]`` — the device analogue of the host
:meth:`context_views` ``[0:n]`` slice. Sliding layers already ring to
their (small) trailing window, so they are unaffected. ``None`` ⇒ full
slab (legacy). ``max_full_len`` is a SYNC-FREE host bound computed by
the caller (``_host_width_bound``: page-table length + block guard) —
no ``.item()`` / D2H, so the whole draft-context assembly stays off the
per-step host-sync that device-slots exists to remove.

Numerically identical to the host path either way:
:meth:`DFlashDraftModel.forward_block_fixed` masks by
``(kpos >= 0) & (kpos < start_pos)`` (+ the SWA window), and every
column beyond ``max active length`` is unfilled (kpos ``-1`` → masked).

DFlash draft-model layer primitives (transformers-free).

The per-block K/V cache, the sliding-window mask helper, the Qwen3 RoPE /
GQA-attention / SwiGLU / decoder-layer nn.Modules, and the
``_set_submodule`` dotted-path swap helper used by the loader. All are
re-exported from ``arbi_serve.spec_decode.dflash`` for import-path
stability.

Per-layer K/V cache for the block-diffusion draft (transformers-free).

Mirrors the reference's ``DynamicCache`` usage: ``update`` appends the
freshly-projected ``[k_ctx ++ k_noise]`` for a layer and returns the full
per-layer K/V; ``crop`` trims every layer to the accepted-sequence length
each decode step. K/V are ``(B, n_kv_heads, seq, head_dim)``. Without it
the draft would attend over only the current block + context and lose the
sequence history, which costs accept-length (not correctness).

NOT A SERVED PATH. This class is a numerical REFERENCE, reachable only
through :meth:`DFlashDraftModel.forward`, which no production caller
invokes; ``tests/test_dflash_engine_tap.py`` is the only code that
constructs one, to pin the optimized path against the upstream
semantics. Serving drafts run
:meth:`DFlashDraftModel.forward_block_fixed` over
:class:`~arbi_serve.spec_decode.dflash_kv_slots.DraftKVSlots`, where a
sliding layer holds a RING of its trailing window rather than the whole
history. So the unbounded per-step growth here, and the dense
full-history mask in :meth:`_Attention.forward` below, describe the
reference and not what the engine runs — read the ring before concluding
the served drafter scales with context.

SDPA with one extra per-head logit that enters only the softmax
denominator.

``q`` is ``(B, nh, q_len, hd)``, ``k``/``v`` are ``(B, nh, n_keys, hd)``,
``attn_mask`` is bool or additive float broadcastable to
``(B, nh, q_len, n_keys)``, ``sink`` is ``(nh,)``.

An all-zero key/value column contributes logit ``0`` and value ``0``;
the additive bias sets that column to ``sink[h]``, so the result is
``sum_j exp(l_j) v_j / (Z + exp(sink[h]))`` — the sink slot holds
probability mass without contributing a value.

GQA attention that never repeat-interleaves K/V across query-head groups.

Each KV head's ``groups`` query heads share one K/V; folding those heads
into the query sequence lets a single ``nkv``-wide SDPA cover the whole
group, so K/V stay ``(B, nkv, S, hd)`` instead of expanding to
``(B, nh, S, hd)``. The expanded K/V are the dominant intermediate a
cudagraph capture bakes at long context; folding removes that term while
computing exactly what per-head SDPA over repeat-interleaved K/V would.

``bias`` is added to every head. ``sink`` matches :func:`sdpa_with_sink`:
one extra zero-valued key column per query head whose logit is the head's
sink, entering only the softmax denominator. Returns ``(B, nh, q_len, hd)``.

The fold keeps the batch and KV-head axes SEPARATE — ``(B, nkv, g*q_len,
hd)``, not ``(B*nkv, g*q_len, hd)``. Every fused SDPA backend requires
4-dimensional inputs and silently declines a 3-D call, which drops the whole
attention onto the math decomposition: that path upcasts q/k/v and the mask
to fp32 and materialises the full ``(B*nkv, g*q_len, S)`` score matrix,
several times over, for a per-step transient orders of magnitude larger than
the fused kernel's output-sized one. Flattening the two leading axes is what
would cost it, so they stay apart.

Full-head RoPE cos/sin generator (Qwen3 rope, NEOX pairing).

``yarn`` swaps the plain ``theta ** (-2i/d)`` inverse frequencies for
the NTK-by-parts schedule and post-scales the whole table by its
``cos_sin_scale`` — the same two knobs :class:`RoPECache` applies on
the main-model path, sharing one :class:`YarnRopeScaling` definition.

Qwen3 GQA attention over ``[context ++ noise_block]`` (non-causal).

Queries come from the noise block; keys/values from the concatenation
of the projected target context and the noise block, so every block
position attends to all context and all block positions.

Reference only, like :class:`DraftKVCache` that feeds it: the served
forward inlines this body at a FIXED context width. The dense
``k_len``-wide mask below is therefore not a served cost.

Pre-norm decoder layer: attn over injected context, then SwiGLU.

``attention_conv`` / ``mlp_conv`` are ``None`` on v1 checkpoints (the
forward is then byte-identical to the plain pre-norm layer). DFlash2
attaches a :class:`GroupedDynamicCausalConv` to each: ``prepare`` runs
after the layernorm and before the sublayer, ``finish`` after the
sublayer and before the residual add, matching the reference exactly.

Registry-driven weight-quant load for the DFlash / DSpark drafter.

The drafter resolves its quantization through the same
:mod:`arbi_serve.weight_quant` backend registry the main model path uses:
the checkpoint's TENSOR KEYS select the backend(s)
(:func:`~arbi_serve.weight_quant.loader.detect_backends`), each dense
projection maps to its mirror class through
:meth:`~arbi_serve.weight_quant.base.QuantBackend.quant_class_for`, and
:meth:`~arbi_serve.weight_quant.base.QuantBackend.bind` reads the payload.
Nothing here knows a quant format — a backend that registers itself serves
drafters with no change to this module.

Ownership rules, all fail-loud:

  - A projection a detected backend owns is swapped for that backend's
    mirror and bound from the checkpoint tensors.
  - A projection the checkpoint's ``ignore`` / ``exclude_modules`` globs
    cover stays dense and must ship a float ``.weight``.
  - Anything else — a projection with neither quantized tensors nor an
    exclusion pattern, an exclusion pattern that also carries quantized
    tensors, or a quantized path with no module in the drafter graph —
    raises.

``tp_shard`` maps each projection to the head-aligned column/row split
(:data:`_COLUMN_LEAVES` / :data:`_ROW_LEAVES`) before asking the backend
for the mirror, so a backend with no parallel mirror for that class
refuses by name instead of binding a replicated weight on every rank.

Build ``cls`` from a weight-quantized drafter checkpoint.

``backends`` are the registry's detections for ``stc``, in detection
priority order. The dense graph is constructed on meta so only the
quantized representation ever occupies accelerator memory.

Logical ``(out, in)`` shapes keyed by canonical runtime name.

Dense tensors report their header shape; a quantized payload has
none (the packed weight's shape is the format's, not the layer's),
so the drafter config supplies the logical dims and the backend's
own bind-time shape validation is the check.

Captured block realization for the DFlash drafters.

A DFlash slate realizes its block LEFT TO RIGHT — the DFlash2 selector walk,
the DSpark markov walk, or the plain parallel draw — and every position of it
is a handful of small kernels issued from the host. The realization is
host-dispatch bound: its wall is several times its device time at every batch
width. That is true of the SAMPLED walk and of the GREEDY walk alike; both
run the same per-position re-rank and differ only in how the position is
resolved (a sampler chain against the slate's mask plan, or an argmax), so
the capture covers both and neither draft regime pays a dispatch cost the
other escapes.

Every input a realization reads is a device tensor and it has no
data-dependent host control flow, so it is recorded once per shape into a
cudagraph and replayed thereafter. A SAMPLED graph is keyed by the realizer,
the active row count, the block depth, the vocabulary, the slate's resolved
mask plan (a launch constant — see
:class:`~arbi_serve.spec_decode.rejection_sampler_ops.SpecMaskPlan`), whether
the slate carries ``min_p``, and the identity of the ``q`` destination it
bakes. Replay copies the fresh block logits, hidden states, anchors, the
slate's encoded sampling parameters and the step seed into persistent
buffers and reads the realized tokens back; ``q`` lands in the engine's
persistent proposal slab exactly where the eager realization writes it. A
GREEDY graph reads no sampler chain, no seed and no proposal ``q`` — its
tokens are the target's own argmax to re-check — so its key is the shape
alone and its replay copies only logits, hidden and the anchor.

The recorded kernels are the eager realization's own — the same callee with
the same tensor-driven arguments — so a replay is bit-identical to the eager
walk under the same seed. A slate the graph cannot represent (a different
mask plan, a ``min_p`` the capture has no buffer for, a ``q`` destination that
is not the baked slab, a keyed watermark) is refused loudly and takes the
eager path.

Warm ``run`` twice eagerly, then record one replay of it.

The warm-ups keep whatever the walk's kernels resolve lazily (workspace,
autotune) outside the recorded stream. ``restore`` puts back persistent
state the warm-ups and the capture wrote through on their way — the
proposal slab a sampled walk bakes can hold a live step's ``q``, and a
capture must leave it as it found it. Raises with no graph left behind
when the capture fails.

Realization I/O allocated once at the widest width, sub-viewed per graph.

Shared by every graph at the same ``(steps, vocab, hidden)`` because the
buffers carry no state between replays: every replay rewrites the rows it
reads, and exactly one realization runs per draft step.

One captured greedy realization at a fixed ``(realizer, rows, steps)``.

The greedy walk re-ranks each position against its realized predecessor
and takes the argmax candidate. It reads only the block logits, the block
hidden and the anchor, and writes only the realized tokens: no sampler
chain, no step seed, no proposal ``q``. So the shape is the whole key —
none of the mask-plan / ``min_p`` / baked-slab conditions a sampled graph
carries can make a greedy slate unrepresentable.

Lazy cache of the realization graphs, pre-swept at boot.

``cores`` maps a realizer kind to the callable that realizes a SAMPLED
block from device tensors, ``greedy_cores`` to the callable that realizes
a GREEDY one; a drafter registers whichever of the two its served
realizer has (the plain parallel block's greedy form is a single argmax
over the block, so it has no walk to record). ``max_graphs`` bounds the
resident shape count across both. ``graph_pool`` is the MemPool the
graphs record into; ``buffers_pool`` holds the persistent I/O the graphs
bake by address (distinct pools, for the reason
:class:`~arbi_serve.spec_decode.dflash_capture.DFlashDraftGraphPool`
gives).

Encode the slate and copy it, with the step seed, into the
persistent sampling buffers the recorded kernels read.

Refuses a slate whose resolved mask plan is not the recorded one (the
plan is a launch constant) or that carries a ``min_p`` the capture
has no buffer for.

Copy the fresh inputs in, replay, return ``(tokens, q_out)``.

The tokens are the persistent output buffer borrowed: the next replay
of any graph over these buffers overwrites them, so the caller
consumes them before its next draft.

Copy the fresh inputs in, replay, return the realized tokens.

The tokens are the persistent output buffer borrowed: the next replay
of any graph over these buffers overwrites them, so the caller
consumes them before its next draft.

DFlash hidden-tap collection helpers — the PER-FORWARD contract.

The tap slabs armed by ``setup_dflash_tap`` hold the residual-stream
features of the tokens of the MOST RECENT tapped forward ONLY, written
from row 0: slab row ``i`` is flat token ``i`` of that forward (multi-
sequence batches are flat-concatenated exactly like ``cu_seqlens_q``
spans). Every forward overwrites the previous forward's rows. This is
by design and cannot be "fixed" into a per-sequence accumulator:

  * the slabs are sized to ``max_batched_tokens`` (a per-STEP cap, not
    a per-sequence context), shared by all sequences of a step;
  * the feature ``copy_`` into rows ``[0, n)`` is recorded into the
    captured decode/verify target graphs at a fixed address — a
    sequence-position-dependent write offset cannot be baked into a
    captured graph.

Production consumers (``DFlashDrafter.observe_*``, the worker observe
op) honour the contract: they collect immediately after each forward
with THAT forward's ``cu_seqlens_q``, so chunked prefill is handled by
accumulating per chunk. The failure mode these helpers make LOUD is the
out-of-band consumer (trainers / tooling) that runs a request to
completion and then reads the slab once, expecting features for the
whole sequence: under chunked prefill the slab holds only the LAST
chunk, and the naive read returns silently corrupt features for any
sequence longer than ``chunk_prefill``.

Validity metadata (recorded by the model's tap dispatcher, capture-safe
device buffers on the slab path):

  * ``_dflash_tap_ntok`` — ``(1,)`` int32, tokens written by the most
    recent tapped forward (the capture-time bucket width under a
    captured graph, which upper-bounds the step's real token count);
  * ``_dflash_tap_pos`` — ``(max_num_tokens,)`` int64, the flat
    ``positions`` of that forward (``-1``-filled when positions are not
    flat, e.g. M-RoPE);
  * ``_dflash_tap_out_pos`` — eager-dict-path (CPU tests) analogue of
    ``_dflash_tap_pos``.

Number of tap rows written by the most recent tapped forward.

Slab path: reads the ``_dflash_tap_ntok`` device buffer (one host
sync). Eager dict path: the captured tensors' row count. ``0`` when
no tapped forward has run yet.

FAIL-LOUD guard for a tap read of ``n`` rows.

The per-forward contract makes rows ``[0, written)`` of the most
recent tapped forward the ONLY valid data; asking for more means the
caller is reading across forwards (the chunked-prefill trap) and
would silently get another forward's rows. Raise instead.

``forward_tokens`` is that count as the CALLER knows it —
``cu_seqlens_q[-1]``, which every production observer has already read
off the pinned host twin to slice its own row boundaries. Passing it
makes the guard free. Omitting it reads the count back from the
``_dflash_tap_ntok`` device buffer, which is a BLOCKING D2H on whatever
thread asks: the whole-sequence collector below has no forward of its
own to name and pays it, once, off the serving path.

Whole-sequence tap read for OFFLINE consumers (trainers / tools).

Returns ``(seq_len, len(tap_ids) * hidden)`` features — tap layers
concatenated in ``_dflash_tap_ids`` order, the same layout as
``DFlashDrafter._collect_tap`` — for sequence positions
``[0, seq_len)``, VALIDATING via the recorded per-forward positions
that slab row ``i`` really is sequence position ``i`` for the whole
span (i.e. the sequence was prefilled by a single forward with no
other sequences in the batch). Raises instead of ever returning the
silently-truncated last chunk of a chunked prefill.

A consumer that cannot guarantee single-forward prefill must either
size ``chunk_prefill`` / ``max_batched_tokens`` past its longest
sequence, or collect per forward like the production observers.

Presence / frequency penalties on the MTP DRAFTER's own proposal.

The verify slate already applies the three penalties per speculative
position (:mod:`arbi_serve.spec_decode.verify_penalties`), so the accepted
tokens come from the distribution the ``K=1`` path would have produced.
The drafter did not, and that asymmetry is pure throughput loss: the
proposal ``q`` is drawn from a penalty-BLIND distribution while the accept
test measures it against a penalty-BEARING ``p``, so drafts systematically
miss. Measured on one RTX 4090 / Qwen3.8-27B-exl3-4.0bpw / MTP K=4 /
non-thinking / temp 0.7 boot, ``presence_penalty`` 0.0 -> 1.5 moved
ms-per-forward-step by +0.2% (i.e. the penalty kernel itself is free) and
tokens-per-forward-step by -3.7%. All of that loss is proposal mismatch.

WHAT IT MEASURED, AND WHERE IT DID NOT
--------------------------------------
Isolated in the shipping configuration (stochastic drafter, fanned mask,
27B exl3 TP1 NOTHING at ``presence_penalty=1.5``: ms/step +0.206 (CI [+0.094, +0.318],
a significant COST) against accept +0.161 (CI [-5.240, +5.562]), and it
cut the penalty's own cost by 0.011 ms/step. Default OFF. The acceptance
interval is +/-5.2 points, so that is "no evidence of a gain", not proof
of none.

The reason is structural rather than a rig artefact, and it is worth
stating because it says where the mechanism CAN work. Softmax is invariant
to a constant logit shift, so alignment acts only on how much the bias
VARIES over the support carrying mass, relative to the drafter's own
error. ``presence_penalty`` is two-valued and loses at both ends, measured
on served text: on MIXED support (prose, 41% of draws are repeats) the
bias does vary, sd 0.74 logits, but the MTP drafter is already 0.53 total
variation from the target, so the correction is second order; on
REPEAT-DOMINATED support (arithmetic, 95% repeats) the drafter is close
(TV 0.08) but the bias is nearly the constant 1.5, sd 0.32, and the
softmax cancels most of it.

``frequency_penalty`` is different IN That is the one regime where this should be visible, and it is the
open question the flag is kept for.

WHAT THIS MODULE IS AND IS NOT
------------------------------
Speculative decoding is lossless for ANY proposal ``q`` (Leviathan-2023
Theorem 1), so nothing here can change the served output distribution --
the verify pass corrects whatever the drafter proposes. That is the whole
reason the drafter may be treated as a free parameter: the exactness bar
that binds :mod:`~arbi_serve.spec_decode.verify_penalties` (where a
broadcast committed-prefix count WOULD break distribution preservation)
does not bind here. Aligning the drafter only moves ACCEPTANCE.

SCOPE -- presence and frequency, not repetition. Both are ADDITIVE in
logit space::

    penalty_add[b, t] = presence[b] * (count[b, t] > 0) + frequency[b] * count[b, t]

so the whole correction is one ``(rows, V)`` tensor SUBTRACTED from the
drafter's logit row -- which is what makes a single chokepoint possible,
and what makes it capturable (a bias is a device tensor; the ``K=1``
penalty chain is host Python over request objects and cannot enter a
CUDA graph). The repetition penalty is a sign-dependent MULTIPLY and has
no additive form, so it cannot ride this rail; a repetition-bearing
request keeps today's repetition-blind drafter (still lossless, just
un-aligned on that one term) while its presence / frequency terms align
normally. Qwen ships ``repetition_penalty=1.0`` in every published
profile, and the knob this fix exists for -- Qwen3.8's non-thinking
``presence_penalty=1.5`` -- is squarely on the additive rail.

STATE REUSE -- zero new full-vocab memory. The committed-prefix term is
EXACTLY the persistent accumulator's own entry
(:mod:`arbi_serve.sampler.penalty_accumulator`), which the sampler already
maintains for the ``K=1`` draw, so the drafter reads it rather than
building a second one. The occurrence counts come from the same
accumulator. Only two ``(rows,)`` knob vectors are new.

DFLASH IS NOT COVERED, deliberately. It is a different drafter behind
``--dflash-draft-path`` (mutually exclusive with the bundled MTP head at
``arbi_serve/engine/mtp_attach.py``), and its structure will not carry an
exact per-position penalty: the block-diffusion forward produces the
logit rows for ALL K positions in one pass and one ``lm_head`` call
before any token is picked, so on its parallel realizers there IS no
realized prefix for positions > 0 to penalise against
(``dflash_driver.py``'s ``base[:, j, :].argmax(-1)`` over the whole block,
and the independent per-position sampled draw beside it). Its
left-to-right realizers -- the markov semi-AR walk and the DFlash2
candidate selector -- DO hold ``toks[:, :j]`` on device at pick time and
could carry one; the committed-prefix term is position-independent and
would apply to every path. None of that is done here, and none of it is
claimed: DFlash keeps today's penalty-blind proposal, which stays
lossless, so this is a missing optimisation and not a missing guarantee.

PER-POSITION EXACTNESS -- chain step ``j``'s history is the committed
prefix PLUS the ``j`` tokens this step already drafted, and the
autoregressive chain has those realized on device at pick time (unlike
the verify slate, which has to reconstruct them). The correction touches
only those ``<= K`` entries and is recomputed from the committed count,
so it reproduces the ``K=1`` formula rather than layering a delta on an
already-penalised value.

Device-resident penalty state the drafter chain reads per step.

Every tensor here is PERSISTENT and allocated once: the captured
drafter chain bakes their addresses at record time, so a replay must
pick up fresh values by having the host write INTO them, never by
swapping the objects out. :meth:`refresh` is the only writer.

``committed`` and ``counts`` are views of the sampler's own penalty
accumulator buffers -- not copies -- so the drafter and the ``K=1``
draw can never disagree about what the committed prefix contained.

Return ``logits`` carrying the K=1 presence / frequency penalty.

THE chokepoint. Every drafter pick -- greedy argmax, the live
per-request sampler, and the graph-capturable tensor-driven sampler --
reads the row this returns, so there is exactly one place where the
drafter's distribution and the verifier's can drift apart.

``draft_prefix`` is the tokens THIS chain step has already drafted,
``(j, B)`` at chain position ``j`` (``None`` / empty at position 0,
whose history IS the committed prefix). The correction it drives is
exact per position: frequency owes one unit per occurrence, presence
owes one unit at the FIRST occurrence and only when the committed
prefix did not already contain the token.

Shape-static given ``j`` and free of host reads, so the whole thing
captures into the drafter chain graph alongside the head forward.

Attach the drafter's penalty state to ``head``; idempotent.

Must run BEFORE the drafter chain is captured: the captured graph bakes
the tensor addresses it reads, so a state installed afterwards would
reach the live chain only and leave every replay penalty-blind.

``None`` (and no attribute) when the sampler's penalty accumulator is
not resident -- there is then no committed-prefix term to read and no
boot reserve that priced one, and inventing a second ``(rows, V)``
buffer here would spend VRAM the KV pool was already handed.

Fold this step's committed tokens in and re-read the knobs.

Returns whether any row asks for a presence / frequency penalty.
A slate that asks for neither leaves every buffer zero, so the
apply below is numerically inert and the captured chain needs no
second variant.

The accumulator fold is idempotent within a step -- it appends
only the tail past the length it last saw -- so calling it here
(the verify commit tail, after the step's tokens landed) and again
from the next step's ``K=1`` draw does the work exactly once, in
whichever of the two runs first.

``Drafter`` Protocol — common surface for every speculative-decoding source.

Two implementations:

  - :class:`arbi_serve.spec_decode.mtp.MtpDriver` — bundled head, shares
    the main model's KV pool, embedding, and lm_head (DeepSeek-V3 / Qwen
    3.5+ convention).
  - :class:`arbi_serve.spec_decode.external_drafter.ExternalModelDrafter`
    — a separate (smaller) draft model loaded into its own
    :class:`MultiStatePool`. K eager-mode decode iterations per
    ``draft(...)`` call (eager-only; no captured-graph chain).

The Protocol is the seam the engine + verify driver sees. Anything
beyond ``draft / commit_results / on_request_admitted /
on_request_finished / rollback_partial_accept / warmup`` is
implementation-private: each drafter manages its own state, the engine
doesn't care.

Common surface every speculative-decoding source implements.

The verify driver (`run_verify_step`) and engine boot (`build_*`)
talk only through this Protocol. The bundled-head and external-
model drafters are sibling implementations; neither inherits from
a shared base class.

Singleton marker: the drafter proposal ``q`` is a one-hot point mass
on the drafted tokens.

Carried in place of a dense ``(K, B, V)`` (or per-row ``(K, V)``)
drafter distribution wherever the greedy-draft proposal flows —
:meth:`Drafter.draft`'s stochastic return, ``Request.mtp_next_draft_
probs``, and the SPMD mirror rows. Slot ``(k, b)`` of the implied
distribution is ``q[v] = 1`` iff ``v == draft_tokens[k, b]`` — fully
determined by the draft tokens already carried alongside, so ONE
process-wide instance (:data:`POINT_MASS_Q`) represents every slate.
The rejection sampler's index-form fast path consumes it without ever
materializing the one-hot (bit-identical accept / recovery / bonus).

The ``(k, b, vocab)`` fp32 destination a drafter realizes its dense
proposal ``q`` into.

The single implementation of "where does a true-stochastic draft put its
``(K, B, V)``". Its shape is fixed at boot — ``k`` from the drafter's
block/chain width, ``b`` from ``max_batch``, ``vocab`` from the model — so
it is the engine's persistent
:attr:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers.sc_drafter_q`,
sliced, whenever the engine in hand has that pool at this width. That pool
is a budget line (``capture.io_buffers``); a fresh vocab-scale fp32 block
on the per-token draft path is not, and at
``gpu_memory_utilization = 1.0`` there is no headroom for one.

Read through the engine on EVERY call and never cached on the drafter: a
model rebuild (``reset_model_state_for_build``) sets
``eng.verify_buffers = None`` and the next build makes a new generation,
so a held reference would outlive its storage.

Falls back to a fresh ``torch.empty`` when there is no pool to slice — a
drafter with no engine handle, a pool built without vocab-scale scratch, a
slate past the worst case, a proposal width the pool was not sized for. The
contents are undefined either way; the caller fills what it returns.

Per-row drafter-q carry for ``mtp_next_draft_probs``.

The single implementation of "stash this row's slice of the drafter
distribution for the next verify step": a dense ``(K, B, V)`` q clones
its ``(K, V)`` column (the true-stochastic path); the
:data:`POINT_MASS_Q` marker carries through as-is (the greedy-draft
path — the row's point mass is fully determined by its
``mtp_next_drafts``, so there is nothing to clone).

True iff a row's ``mtp_next_draft_probs`` carry can serve a K-step
rejection accept: either the :data:`POINT_MASS_Q` marker (greedy-draft
proposal — valid at any K, the tokens ride ``draft_tokens``) or a dense
per-row ``(K', V)`` tensor with ``K' >= k_step``. ``None`` (cold seed /
first verify step after a greedy chain failure) is not a carry — that
slate accepts greedily for one step, exactly as before.

FULL (logical) vocab width the drafted argmax tokens index into.

Under TP the lm_head is VOCAB-PARALLEL, but every drafter's forward
all-gathers + trims so the returned argmax draft tokens span the FULL
vocab ``[0, dims.vocab_size)``. Sizing the one-hot ``q`` from a sharded
width would scatter a full-vocab token index into a half-width buffer →
CUDA ``index out of bounds``. Returns the engine model's logical
``dims.vocab_size`` (0 if unavailable — the caller falls back to a
driver-local width).

Width of the row the DRAFTER forms — the ``q`` a proposal buffer holds.

Normally :func:`resolve_full_vocab_size`: the drafter projects through the
same lm_head the verifier does. ``--draft-vocab-prefix`` breaks that tie by
giving the drafter its own lm_head sliced to token ids ``[0, N)``
(:mod:`arbi_serve.draft_vocab_prefix`), so every tensor the drafter builds
— its logit row, its sampled ``q`` — is ``N`` wide while verify keeps the
full vocabulary.

The two widths COEXIST by design. A ``q`` supported on ``[0, N)`` is a
valid proposal (Leviathan-2023 Thm 1) and the verify seam pads it into the
full-width residual buffer
(:func:`~arbi_serve.spec_decode.mtp_verify_accept._resolve_draft_probs`),
so a token outside the prefix is still EMITTABLE through the residual
branch — it is only never DRAFTED. What must not happen is a drafter
buffer sized from the VERIFIER's width and written from the drafter's:
that is a shape error, not a wider buffer.

Falls back to the full width whenever no sliced head is installed, so a
caller reading this before
:func:`~arbi_serve.engine.build_phases_load._apply_draft_vocab_prefix_phase`
gets the conservative answer rather than a wrong one.

Scatter one row's point-mass proposal into a zeroed dense column.

The MIXED-slate escape hatch: when a verify slate carries BOTH dense
per-row q tensors (true-stochastic rows drafted in an earlier slate)
and :data:`POINT_MASS_Q` rows, the sampler needs one dense ``(K, B,
V)`` operand — the marker rows are densified here, writing ``1.0`` at
each drafted token (byte-identical to a materialized one-hot). An
all-marker slate never calls this (the sampler's index-form fast path
skips densification entirely).

Whether this drafter's chains are actually being gated right now.

The capability AND the flag: a capable drafter on an unarmed boot
drafts the full K every step, so the verify plan must keep reading
a short chain as a cache miss.

Maximum speculative-token chain length the drafter can emit.

Per-request K is capped at this value at admission time. The
engine's verify-pass batch shape ``(K+1, B)`` is sized to
``max_k + 1`` worst-case.

Run K speculative steps and return ``(K, B)`` draft tokens.

Greedy contract (``sampling_params=None``): returns just the
``(K, B)`` int64 token tensor.

Stochastic contract (``sampling_params`` set): returns
``(tokens, q)`` where ``q`` is the per-slot drafter proposal the
verify pass's residual sampler (Leviathan-2023) consumes — the
:data:`POINT_MASS_Q` marker for the default greedy-argmax
proposal (a one-hot on ``tokens``; never materialized), or the
dense ``(K, B, V)`` distribution for true-stochastic drafting.

``prev_hidden`` is the bundled-head input (the main model's
last hidden state). External drafters ignore it (they keep
their own per-request KV history).

``requests`` is the per-row engine ``Request`` list; external
drafters use it to look up their own page-table entries.
Bundled head uses it for KV-slot allocation in the main pool.

``seed_buf`` is the engine's persistent ``int64 (1,)`` step
seed, required for stochastic sampling determinism.

LoRA: the drafter runs the BASE model (no ``lora_state`` is
threaded here) even for a request with a ``lora_id``. This is
intentional and correctness-safe — the verify forward applies the
adapter (see ``EagerModelRunner._reconstruct_verify_batch``), and
verify is the lossless authority: any draft the un-adapted drafter
proposes that the LoRA'd target would not have produced is simply
rejected. The only cost is a lower accept rate for LoRA requests,
never a wrong token.

Roll back drafter-side KV state for rejected tokens.

Per-row, ``k - n_accepted`` tokens were drafted but rejected.
Drafters that maintain their own KV cache (external) must
truncate to the accepted prefix here. Bundled head writes K/V
into the main pool's draft slots and the engine's
``free_draft_slots(k_accepted, k_rejected)`` already covers
rollback there — bundled implementations make this a no-op.

Update per-request + aggregate accept counters from verify results.

Engine-side commit (token append, page-table extend, stop
condition checks) is owned by the engine — the drafter only
updates its own metric counters here.

One-time first-call warmup (JIT compile, alloc scratch, etc.).

Called once during engine boot after the drafter is attached.
Bundled head is a no-op (warmup happens implicitly during the
main model's profile pass). External drafters run a synthetic
single-token forward here so the first real request doesn't
pay JIT-compile cost.

Shared drafter-chain circuit breaker.

The MTP verify path accounts every drafter-chain draw with an
UNCONDITIONAL pair of calls — the rank-0 sync/async seed
(:func:`arbi_serve.spec_decode.mtp_verify_accept._seed_drafter_bucket`) and
the rank-symmetric SPMD worker
(:func:`arbi_serve.spec_decode.mtp_verify_spmd`)::

    driver.note_drafter_chain_failure(exc, context=...)  # on draw failure
    driver.note_drafter_chain_success()  # on draw success

where ``driver`` is whatever :class:`~arbi_serve.spec_decode.drafter.Drafter`
is installed as ``eng.mtp_driver`` (a forwarding alias for ``eng.drafter``):
the bundled MTP head (:class:`~arbi_serve.spec_decode.mtp_driver.MtpDriver`),
the block-diffusion DFlash drafter
(:class:`~arbi_serve.spec_decode.dflash_driver.DFlashDrafter`), or a
standalone draft model
(:class:`~arbi_serve.spec_decode.external_drafter.ExternalModelDrafter`).

Mixing :class:`DrafterCircuitBreakerMixin` into EVERY drafter keeps those
two methods — and the loud-fail-on-repeated-capture-failure protection they
implement — in ONE place so the implementations cannot drift: a drafter
that omits the mixin would raise ``AttributeError`` on
``driver.note_drafter_chain_success()`` and kill every speculative-decode
step for that drafter. This module exists to make that class of bug
impossible: a new drafter that forgets the mixin still has the two methods
via the Protocol expectation, and the ones we ship all share this code.

Requires the host to (optionally) expose ``self.engine`` — an
:class:`~arbi_serve.engine.engine.Engine` or ``None``. When present the
breaker marks ``engine._serving_degraded`` (feeds the ``serving_degraded``
gauge / ``x-arbi-degraded`` header); a host without an engine reference
still gets the full counter + breaker semantics.

WHAT THE BREAKER IS FOR, AND WHAT IT IS NOT FOR. Its window counts a DEFECT —
a drafter that cannot draw — and its threshold ends the process because a
broken drafter silently degrades every step and a restart is the fix. An OOM'd
draw is a different fact: the allocation did not fit, cold-pathing it is the
correct response (it frees the draft transient), and a restart hands the
replacement process the same card. It therefore has its own window here and its
own reaction in :mod:`arbi_serve.engine.memory_pressure`, and the breaker
cannot see it. That split is what keeps this module's threshold honest: it
remains a statement about the drafter alone.

Raised when a drafter chain fails ``ARBI_DRAFTER_FAILURE_LIMIT``
consecutive times.

Each drafter-chain failure cold-paths its step (correct output, all
speculation lost — a large invisible tok/s regression). One transient is
tolerated and self-heals (any drafter success resets the count); a
PERSISTENT failure means every step is silently degraded, so the engine
loop re-raises this out of ``run_forever`` — the engine thread dies, the
``/health`` liveness probe flips 503, and orchestration restarts the
server rather than letting it serve slow forever.

A DRAFTER DEFECT ONLY. An out-of-memory draw never reaches this window:
restarting the process hands the replacement the same card, so the
threshold that is right for broken code is exactly wrong for a card short
of memory. Those go to :meth:`DrafterCircuitBreakerMixin._note_drafter_oom`
and the memory ladder in :mod:`arbi_serve.engine.memory_pressure`.

Raised when a draft-step failure fired INSIDE a half-issued TP collective.

An OOM (or any exception) thrown while a vocab-parallel collective is
mid-flight — the drafter's ``embed`` all_reduce or ``lm_head`` all_gather, or
a TP-sharded draft forward's row-parallel all_reduce — leaves the ranks
DESYNCED: rank 0 bailed but its peers are still blocked in
``all_gather_into_tensor`` / ``all_reduce`` on the same NCCL communicator. The
ordinary drafter breaker would cold-path the step and CONTINUE, which is
exactly the silent NCCL wedge we refuse (both GPUs pinned at 100%,
``/health`` = 000 — never-wedge violated).

So this is a HARD fault, not a cold-path candidate: the driver latches the
engine's sticky fatal fault (``/health/ready`` → 503, admission →
``engine_context_poisoned``) and raises THIS out of the seed-draft handler.
Because it subclasses :class:`DrafterChainBreakerError` the engine loop
re-raises it out of ``run_forever`` (engine thread dies → orchestration
restarts) — fail loud + tear down cleanly instead of wedging in place.

Drafter-chain failure accounting + circuit breaker.

Mixed into every :class:`~arbi_serve.spec_decode.drafter.Drafter`. See
the module docstring for why this lives in one place. Subclasses SHOULD
call :meth:`_init_drafter_breaker` from ``__init__`` (the class-level
defaults below make the accounting safe even if they forget) and MAY set
:attr:`_breaker_label` to name their drafter in logs / the degraded
registry.

Initialise the per-instance failure-window state. Call from ``__init__``.

``total_cold_path`` counts every drafter-chain failure that reset a
bucket to cold-path (all-time, never reset);
``_consecutive_drafter_failures`` is the breaker window, reset to 0 by
any drafter success. ``total_oom_cold_path`` /
``_consecutive_drafter_ooms`` are the same pair for the MEMORY window,
which the breaker cannot see — an OOM'd draw is counted in both totals
(the speculation is lost either way) and in only the memory window.
See :meth:`note_drafter_chain_failure`.

Account one drafter-chain failure; trip the breaker on persistence.

Called from every drafter-chain except-handler AFTER the bucket's
rows were reset to cold-path (the step itself stays correct — it just
loses speculation). Bumps ``total_cold_path``, extends the
consecutive-failure window, marks the engine's ``_serving_degraded``
registry (feeds the ``serving_degraded`` gauge / ``x-arbi-degraded``
header) when an engine reference is present, and ERROR-logs the
failure — full traceback on the first of a burst, then rate-limited to
one line per ``_DRAFTER_FAILURE_LOG_INTERVAL_S``.

Breaker (``ARBI_DRAFTER_FAILURE_LIMIT``, default 5): once the window
reaches the limit, raises :class:`DrafterChainBreakerError` — the
engine loop re-raises it out of ``run_forever`` (server death →
orchestration restart). ``0`` trips on the FIRST failure; ``-1`` never
trips (legacy log-and-continue). A single transient never trips it:
:meth:`note_drafter_chain_success` resets the window.

TWO WINDOWS, because there are two facts. The breaker window above
counts DEFECTS — a drafter that cannot draw — and its threshold ends
the process because a broken drafter degrades every step and a restart
is the fix. An OOM is not that fact and a restart is not that fix, so
:meth:`_note_drafter_oom` takes it into a separate window instead. This
is a classification, not a new threshold: there is no second limit to
set, because the memory ladder escalates on a STATE (nothing left to
narrow) and not on a count.

Account one drafter draw that OOM'd, and hand it the memory reaction.

The bucket is already cold-pathed by the caller, which is the whole of
this path's own narrowing: the draw's transient is not allocated, the
step still commits the correct token, and only the speculation is lost.
What this adds is the reaction the serving path had no way to produce —
:func:`~arbi_serve.engine.memory_pressure.note_oom` re-measures the step
budget from the card and re-arms admission with it, closes the front
door when the budget can no longer cover one row, and says so on the
``serving_degraded`` registry throughout, so a server cold-pathing under
memory pressure is never silently slow.

It raises :class:`~arbi_serve.engine.infra_health.
StepMemoryExhaustedError` on exactly one outcome: ``note_oom`` reporting
that the ladder ran out (an OOM with the door already shut and no step
succeeded since). The sticky fault is latched by then, so this raise is
only how the verdict leaves the seed-draft path — and it is raised from
HERE rather than from the caller for the same reason
``DrafterCollectiveFault`` is: every seed-draft path (sync / async /
spmd / offload / run_step) funnels through this one method, so the
behaviour cannot exist on some of them and not others.

Reset the drafter's failure windows after a successful draw.

Cheap early-out on the healthy path (one int compare against the OR of
both windows). Clears the ``drafter`` entry from the engine's
``_serving_degraded`` registry (if present) so the gauge drops back to
full-speed.

Resets BOTH windows, because a successful draw is evidence against both
readings: the drafter can draw, and the card had room for the draw.

DeepSeek-V4 DSpark drafter — the loop that drives the bundled stages.

DeepSeek's own recipe for Flash-0731 pins ``{"method": "dspark",
"num_speculative_tokens": 7, "draft_sample_method": "probabilistic"}``:
one BLOCK per draft, realized left to right against a bigram bias, then
verified in a single target forward.

This is a :class:`~arbi_serve.spec_decode.drafter.Drafter`, so everything
downstream of the proposal is the machinery that already exists — the
``SpecMode.MTP`` step, ``build_verify_plan``'s ``K + 1`` slots per row,
the rejection sampler, ``verify_and_accept``, and the engine's own token
commit. Nothing about verification or acceptance is re-implemented here.

Two seams it plugs into rather than adding:

  * ``observe_seed_forward`` / ``observe_verify_forward`` — the hooks the
    run-step and the verify driver already call on any drafter that
    defines them, used here to stash the target's tapped residual streams
    for the tokens that forward COMMITTED;
  * :class:`~arbi_serve.spec_decode._dflash_driver_dspark._DFlashDsparkMixin`
    — the DSpark confidence and semi-AR block realization, unchanged: it
    reads ``self.model.markov_head`` / ``confidence_head``, and
    :class:`~arbi_serve.models._deepseek_v4_mtp.DSparkDraftHead` presents
    both.

What the stages consume is not what the engine hands a next-token head.
A DSpark stage reads the target's residual STREAMS — ``(N, hc, D)`` per
tapped layer, concatenated — so the per-row context comes off the model's
tap, not off ``prev_hidden``.

The committed span a request's stages have not consumed yet.

``streams`` is ``(n, hc, D * len(target_layer_ids))`` off the target's
tap and ``ids`` the tokens at those positions, starting at
``first_position``. The next block's ANCHOR is the position right
after this span — the token the target emitted from it, which the
engine hands the drafter as ``last_token_ids``. That token has no
tapped row of its own (it was an OUTPUT of the forward, never an
input), which is exactly why the block feeds it through the
embedding rather than through ``main_proj``.

Check the stages draft through the TRUNK's vocab tensors.

A DSpark checkpoint ships neither an embedding nor an LM head for
its stages, so a chain bound to anything else is reading
uninitialised weights.

No-op: a DSpark block publishes nothing that outlives it.

The block runs TRANSIENT, so it writes no ring slot and no pooled
entry, and the next draft re-runs the committed span from wherever
the accept landed. The target's own state is rolled back by
``DSv4StatePool.rollback_batch``, which this drafter does not own.

Fold the verify results into the per-request and lifetime counters.

``req.mtp_proposed`` / ``req.mtp_accepted`` are what
``MtpStrategy.stats_for`` publishes per request and what /metrics
and the bench report sum — a drafter that only kept its own
aggregate would report a served deployment as speculating zero
tokens.

Stash the ACCEPTED prefix's tapped streams after a verify pass.

The forward ran ``[anchor, d_1 .. d_K_eff]`` per row and the engine
committed ``n_accepted`` drafts plus a bonus token. The rows that
extend the stages' context are the anchor and the accepted drafts —
``n_accepted + 1`` of them — because the anchor reached the stages
only as a TRANSIENT block slot, which published nothing. The bonus
token was an output, has no tapped row, and becomes the next
block's anchor through ``last_token_ids``.

Append each ``(request, first_row, count)`` span to its pending.

Reads the tap of the forward that JUST ran — it holds that
forward's tokens only, from row 0, and every forward overwrites it.

Extend the request's pending span, or start a new one.

A span that starts at or before what is held REPLACES it — the
forward re-ran those positions and its tokens are the committed
ones. A span that starts past the held one's end would leave a hole
the stages never run, so it is refused.

Propose one block per row and return ``(K, B)`` draft tokens.

``prev_hidden`` is unused: a DSpark stage reads the target's
residual STREAMS at the committed positions, which arrive through
the tap rather than as one reduced hidden state per row.

``(B, block_size - 1, ·)`` base logits and hidden per row.

The committed span runs FIRST — one flat forward across every row
— so each stage's ring carries the context before the block reads
it. The anchor's own row is dropped from both returns: it predicts
a token the target already committed.

``ARBI_DFLASH_DYNAMIC_K``'s arming predicate.

One reader of the flag, shared by the boot refusal
(:mod:`arbi_serve.spec_decode.dynamic_k_boot`), the verify plan that
consults the hook, and the hook itself — so "is dynamic-K armed" cannot be
answered one way at the gate and another at the site it gates.

OFF costs one attribute load per verify plan: nothing below the flag check
runs, so an unarmed boot never consults a drafter about it.

What ``ARBI_DFLASH_DYNAMIC_K`` needs from the rest of the stack.

Checked ONCE at boot, by name. Dynamic-K shortens the slate's UNIFORM
draft length to the confidence head's confident prefix, so the verify
block it produces is ``dyn_k + 1`` rows instead of ``K + 1`` — a
data-dependent verify WIDTH, decided from the drafter's own scores. That
is the hazard :mod:`arbi_serve.spec_decode.verify_width_boot` states, and
the reason it is refused rather than served is the reason stated there:
under unpinned numerics the truncation moves the served text, and
dynamic-K's own acceptance criterion (accepted/proposed up, ``accept_len``
not down) is measured on the trajectory it moved.

The DSpark drafter closes that loop tighter than a bare GEMM reorder does:
:meth:`~arbi_serve.spec_decode.dsv4_dspark_drafter.DSv4DsparkDrafter.observe_verify_forward`
stashes the target's tapped residual streams from the verify forward that
just ran, at the width that forward ran at, and those streams are the next
block's draft context — which is what the next confidence prefix, and so
the next ``dyn_k``, is computed from.

Neither condition here raises on its own downstream: the truncation
serves, the counters look healthy, and the number the A/B reports is not
the number dynamic-K is worth.

Refuse dynamic-K on a drafter with no confidence prefix to read.

Read off the DECLARED capability — a direct attribute read, not a probe
with a default. A probe answers "no" identically for a drafter that
K with the
flag reading ON. That is how this lever's counters came to read zero on
every chain and tree boot for a structural reason while looking like a
tuning one.

Only a DSpark checkpoint's confidence head scores positions. The
bundled MTP head and an external draft model have nothing to score
with, so on those the flag could only ever be inert.

External draft-model drafter — small model paired with a larger main model.

Eager-only K-step decode chain. The drafter holds its own
:class:`MultiStatePool`, :class:`FlatPageTable`, per-layer
:class:`AttnOp` list, and per-:class:`StateKind` metadata builders.
Each ``draft(...)`` call runs K eager decode iterations against the
drafter's own model + KV.

Supported scope:

  * single-GPU only (TP=1)
  * paged-KV only (no MLA / Mamba / GDN / ShortConv)
  * attention archs only (no recurrent state to roll back)
  * dtype assertion + tokenizer-hash assertion at boot

The 8 boot asserts in :meth:`__init__` make any violation a hard
RuntimeError with an actionable message.

Standalone draft-model drafter implementing the :class:`Drafter` Protocol.

Holds its own model, pool, page table, attn ops, and metadata
builders. The main engine threads through ``on_request_admitted`` /
``on_request_finished`` so the drafter's page table tracks the
same active set of requests; per-step ``draft(...)`` runs K eager
decode iterations.

Register a fresh request in the drafter's page table.

Idempotent — duplicate calls for the same request_id are no-ops
(admission may be retried under multi-group / preempt paths).
On admission we allocate the prompt's KV up-front by running a
single drafter forward through ``_warmup_request`` lazily on
the FIRST ``draft()`` call rather than at admit time — keeps
admit cheap and lets the drafter forward overlap with the
main model's prefill.

Free the drafter's page-table entry for this request.

Idempotent — the engine's terminal-state code may call this from
multiple finish paths (stop / length / context / timeout /
cancelled / error) per the same ``_metrics_finished`` guard.

Trim drafter-side draft slots for rejected tokens.

:meth:`draft` allocated ``k`` slots per row — slot 0 holds the
SEED token (the prior step's committed bonus, newly written into
the drafter's own KV) and slots ``1..k-1`` hold the first ``k-1``
proposals (the Kth proposal's K/V is never written; see
:meth:`draft`). A verify pass that accepts ``n_accepted`` drafts
leaves ``n_accepted + 1`` committed tokens that the drafter KV
must keep (the seed + the accepted draft prefix), so we retain
``min(n_accepted + 1, k)`` slots and free the trailing rejected
ones via :meth:`FlatPageTable.finalize_draft_slots`. Retaining
only ``n_accepted`` (ignoring the seed slot) would drop the last
accepted token's K/V every step, drifting the drafter one token
behind per step and collapsing the accept rate. A FULL accept
(``n_accepted == k``) keeps all ``k`` slots; its Kth token was
never written, so the next step re-seeds from the bonus with a
benign 1-token context lag.

The drafter is attention-only (assert-checked in __init__), so
there is no recurrent-state snapshot to roll back — page-table
finalize is the entire rollback.

Run K eager decode iterations and return ``(K, B)`` int64 tokens.

``prev_hidden`` is ignored — the external drafter maintains its
own per-request KV history and seeds each step from
``last_token_ids``.

``frontier_repair_meta`` is ignored — it overrides the
bundled-head drafter's read of the MAIN model's KV frontier (the
async verify path advances the main pool optimistically before
committing the prior step). The external drafter owns a SEPARATE
KV pool and manages its own frontier via ``allocate_draft_slots``
/ ``rollback_partial_accept``, so the main-pool repair meta does
not apply. Accepted for signature parity with
:meth:`MtpDriver.draft` so the async verify driver's
``draft(..., frontier_repair_meta=...)`` call does not raise.

``requests`` is required (the drafter needs each row's
request_id to look up its drafter-side page-table entry).

Stochastic path (``sampling_params`` set) returns
``(tokens, q_full)`` matching the bundled-head contract.

Eager, no captured-graph fast path (a future
``DrafterModelGraphPool`` could plug in here).

Replay the captured greedy K-step chain for a hit ``(B, k)``; else None.

Returns ``(k, B)`` int64 draft tokens on a captured-graph hit,
leaving the K allocated draft slots in place for the verify driver
to trim (the same post-chain contract as :meth:`draft`'s eager
path). Returns ``None`` for a stochastic draft (the captured chain
bakes the greedy argmax feedback) or an uncaptured/off-ladder
``(B, k)`` shape — the caller then runs the eager chain. A miss
warns ONCE per shape (loud fallback, never silent).

The K per-step paged-attn metas + the shared block_table are
derived by :meth:`_build_replay_step_metas` (positions / slots /
seq_lens depend only on the page-table lengths + allocated slots,
not on the drafted token values, so the whole chain is
precomputable) and copied into the captured persistent buffers by
:meth:`ExternalDrafterChainGraph.replay`.

Build the K live per-step ``AttnPagedKVMeta`` + the shared block_table.

Mirrors the per-row position / slot / seq_len derivation in
:meth:`_forward`, computed for all K steps at once. The values are
token-independent (they follow from the page-table lengths + the
allocated draft slots), so the whole chain's metas are precomputable
before replay. Each meta carries the slices
:meth:`ExternalDrafterChainGraph.replay` copies into the captured
persistent buffers (``slot_mapping`` / ``seq_lens`` /
``cu_seqlens_k`` / ``positions``); ``block_table`` is invariant
across the K steps (the draft slots live within already-allocated
pages) so it is built once and shared.

Populate the drafter's own KV with this request's prompt.

The drafter is a SEPARATE model with its OWN paged-KV pool; the
main engine's prefill writes only the main model's cache. Without
this, the decode draft chain attends over uninitialised KV (and
starts at position 0 with no context), so every proposed token is
rejected. Runs ONE multi-token ``is_prefill=True`` forward over
the full prompt — positions ``0..L-1``, one slot per token — to
write the prompt K/V, then leaves the page-table length at ``L``
so the first decode draft lands at position ``L``.

Single-request (``B=1``) — called per-row from :meth:`draft`.
Caller holds ``torch.inference_mode()``.

Run ONE drafter decode step.

Builds a single-token-per-row :class:`ScheduledBatch`, runs the
drafter model's forward, samples one token per row.

Mirrors :meth:`EagerModelRunner._build_batch` / ``forward`` but
bound to this drafter's own pool / attn_ops / metadata
builders. Everything is eager (a captured-graph replay could
replace this in the future).

Gemma-4 EAGLE3 assistant as an arbi MTP draft head.

The ``gemma-4-E2B-it-assistant`` checkpoint is a low-rank EAGLE3
drafter: a 4-layer Gemma decoder at internal width 256 wrapped by a
``pre_projection`` (2×backbone → 256) and ``post_projection`` (256 →
backbone), with Q-only attention that reads the VERIFIER's KV cache.

Forward (one draft token per call; the driver chains for K>1)::

    e = verifier_embed(token)                       # (B, 1536)  scaled sqrt(1536)
    x = pre_projection(cat([e, prev_backbone_hidden]))   # (B, 256)
    for layer in layers[0..3]:                       # Gemma 4-norm sandwich, Q-only
        x = layer(x, positions, verifier_kv[i])      # reads verifier layer 13/14 KV
    draft_hidden    = norm(x)                        # (B, 256)
    backbone_hidden = post_projection(draft_hidden)  # (B, 1536)  -> next-step feedback
    logits          = masked_embedding(draft_hidden, draft_lm_head_w)   # centroid-sparse

The head OWNS the drafter weights (``pre_projection``,
``post_projection``, the 4 blocks, ``norm``, ``embed_tokens`` [= the
tied draft lm_head, 262144×256], and ``masked_embedding``). It holds a
tied REFERENCE to the VERIFIER's scaled ``embed_tokens`` (1536-wide) for
the input-embedding hook — that is Gemma's own embedding, so the scale
is definitionally correct.

Q-only attention shares the verifier's KV via arbi's existing cross-
layer primitive (``LayerSpec.skip_kv_proj`` + ``kv_source_layer`` view
aliasing + backend ``skip_kv_write``); the head builds
:class:`AttentionBlock` with ``skip_kv_proj=True`` per layer, so there
is no k/v projection and the attn op reads but never scatters. See
``docs/gemma4-mtp/DESIGN.md`` §3/§5.

Centroid-masked sparse logits over the draft lm_head weight.

Projects the draft hidden to ``num_centroids`` scores, selects the
top-``k`` centroids, and scores ONLY the tokens belonging to those
centroids against the draft lm_head weight. ``token_ordering`` maps
``centroid c`` to its ``vocab // num_centroids`` vocab ids via
``view(num_centroids, per_centroid)[c]``. Mirrors vLLM
``Gemma4MTPMaskedEmbedder``.

One Gemma-4 draft block — four-norm sandwich, Q-only attention.

Composes arbi's shared components (``GemmaRMSNorm``,
``AttentionBlock(skip_kv_proj=True)``, ``GatedSiLUMLP`` with
``gelu_tanh``) — no reimplemented Gemma math. Clean per-block
signature ``(x, positions, state_view, attn_meta, rope_cache,
attn_op)`` so the head drives it directly (mirrors
``_Qwen3_5MTPLayer``).

Gemma-4 EAGLE3 assistant draft head (one token per call).

``n_draft`` is the tokens produced per head call (1 — the driver
chains for K>1). ``num_internal_layers`` (4) is the block count run
per token; it is decoupled from the draft DEPTH.

FULL ``(T, vocab)`` logits, non-selected positions = dtype min.

This is the LOSSLESS verify-path distribution. The ``-inf``-band
fill keeps a full-width row so the rejection sampler / argmax see
the true vocab index space.

Build the head.

``verifier_embed`` is a REFERENCE to the verifier's scaled
embedding module (Gemma ``GemmaScaledEmbedding``, 1536-wide) —
held out of the parameter graph so the loader binds it once on
the verifier. ``kv_source_layers[i]`` is the verifier layer
whose KV draft layer ``i`` reads (None in a standalone/test
build where the caller supplies attn ctx directly).
``marker_layer_base`` is the ``layer_idx`` the first draft marker
spec occupies on the engine's per-layer slab (serve path).

``use_ordered_embeddings`` selects the draft-lm_head projection:
True → the centroid-masked sparse path (``num_centroids`` /
``token_ordering``, the E2B checkpoint); False → a plain dense
``hidden @ embed_tokens.weight.T`` over the full vocab (the 31B
checkpoint has a full ``embed_tokens`` and no centroid tensors).
``num_global_kv_heads`` is the K/V-head count on FULL-attention
draft layers (the verifier's ``num_global_key_value_heads``); when
None the sliding count applies to every layer (E2B, where the
verifier has no distinct global-KV count).

Per-parameter shard spec binding the head's own tensors.

Maps the checkpoint's ``gemma4_assistant`` keys onto the head's
modules under ``dst_prefix`` (the attribute path the head occupies
on the engine model). The VERIFIER's ``embed_tokens`` is NOT bound
here — it is a tied reference the verifier's own weight map loads.
TP sharding: q/gate/up column-parallel (``shard_dim=0``),
o/down row-parallel (``shard_dim=1``); low-rank projections,
norms, the draft lm_head weight and the centroid tensors are
replicated. ``token_ordering`` is an int64 index buffer —
``keep_src_dtype`` so it is not cast to the engine dtype.

One draft step — greedy (default) or true-stochastic.

Greedy (``sampling_params is None``): returns ``(draft_token
(B,), feedback_hidden (B, backbone))`` and, when
``return_logits``, the FULL ``(B, vocab)`` logits third. The
draft token is the centroid sparse-argmax fast path
(:meth:`to_token`) — no full-vocab tensor materialized.

TRUE-STOCHASTIC (``sampling_params`` provided, one per row): the
draft token is SAMPLED from the head's own centroid-masked
filtered distribution (the request's temperature / top_k / top_p
/ min_p) via the SAME graph-safe
:meth:`GraphSafeRejectionSampler.sample_drafter_token` the Qwen
head uses — no argmax, no Gemma-specific sampler. Returns
``(draft_token (B,), feedback_hidden (B, backbone), q_full
(B, vocab))``; the dense per-row drafter distribution ``q_full``
feeds the verify pass's residual rejection sampler (Leviathan-2023
Thm 1 holds for any ``q``). The centroid-sparse ``q_full`` puts
the scatter-floor tokens at ``q ≈ 0`` after softmax — the
rejection sampler's ``q = 0 ⇒ accept iff p > 0, resample
residual`` path consumes it losslessly. ``slot_offset`` (the
chain step index) decorrelates the K counter-keyed draws from the
shared ``seed_buf``. ``seed_buf`` (the engine's rank-agreed step
seed) is REQUIRED — every TP/SPMD rank derives byte-identical
draws from it.

The ``run_stack`` output (draft-dim) is projected two ways:
through ``post_projection`` for the next step's feedback hidden,
and through the centroid path for the token / logits / q_full.

One-shot MTP K calibration — cache key + on-disk I/O.

The flow is intentionally minimal:

    1. ``arbi-serve calibrate-mtp-k --model … --backend …`` boots one
       in-process Engine, sweeps ``K`` ascending (0, 1, 2, …) up to
       ``--k-max``, records STEADY-STATE decode tokens/sec at each K,
       picks the winner, and writes a JSON to
       ``<k_calibration_dir>/<key>.json``. The curve is NOT unimodal at
       the shallow end (a chained head can dip at K=1..2 before the
       ``1 + a + a² + …`` series starts paying), so the sweep early-stops
       only after several consecutive above-noise declines — see
       :func:`sweep_k_in_process`.

    2. On every subsequent ``arbi-serve …`` boot the engine looks up
       the same cache key and pins ``cfg.mtp.n_draft`` to the cached
       ``best_k`` — IFF the operator did NOT pass ``--mtp-n-draft`` on
       the command line (``MtpConfig.n_draft_explicit == False``).

The cache key is a SHA-256 digest over a canonical dict of every input that
can change the picked K: the absolute model path, the SHAPE that checkpoint
is served in (weight quant + dtype, registered and active backends, TP,
context window, drafter), the GPU model name, the prompt corpus, and the
thinking regime. Deployments that differ on any of those coexist as separate
files; see :class:`KCalibrationKey` for the field-by-field reasoning and for
what is deliberately absent.

NOTE: this module is import-cheap (stdlib only) so the engine boot
path can call :func:`load_calibration_for` without dragging in torch.

Digest the named top-level functions of ``src``.

Docstrings and comments are stripped before hashing: prose about the
sweep does not change the number the sweep returns, and a guard that
fires on a comment edit trains its readers to re-pin without thinking.
The digest therefore moves only on a change to executable structure.

Stable for one Python minor version (``ast.unparse`` normalizes the
source); the guard runs in the pinned image, which is the environment the
pin is taken in.

Raises ``LookupError`` when a name is absent — a renamed or deleted
function silently shrinks the digest's coverage, the one failure mode a
digest cannot report as a mismatch.

Split-K register MTP kernel ``block_m`` ceiling (= K+1) — single source:
tkv's ``mtp_max_block_m`` flag (env ``TKV_MTP_MAX_BLOCK_M``).

This is the boundary BELOW which verify runs on the fast split-K register
kernel; it is NOT the served verify-K cap (the Turbo prefill verify route covers above it — see
:func:`tkv_max_verify_block_m`). Still the single source for sizing the
split-K cudagraph CAPTURE sweep (the Turbo prefill high-K verify shapes run eager).

Largest MTP-verify ``block_m`` (= K+1) the installed tkv engine SERVES.

Reads turbo-attn's :func:`tkv.runtime.attention.mtp_verify_max_block_m`
(the SINGLE source of truth): turbo-attn's decode-vs-prefill routing cap when
the Turbo prefill high-K verify route is present, so verify chunks above the
split-K register ceiling are served on the Turbo prefill kernel rather than asserting.

Route-aware fallback (fail-loud, never over-promise): an OLDER turbo-attn
without the Turbo prefill high-K route lacks this symbol — fall back to the split-K
register ceiling (:func:`tkv_max_block_m`) so the served cap can NEVER
exceed what the engine can actually run. Cached (the import drags torch;
this module stays import-cheap on the boot path).

Smallest verify ``block_m`` the Turbo prefill verify route serves; below it is split-K.

NOT the split-K ceiling. ``TKV_MTP_PREFILL_SPLIT`` (on by default in
turbo-attn) deliberately routes the whole verify band to the Turbo prefill kernel's
tensor-core mainloop, which drops this floor to 2 — so on a default
boot the split-K register kernel is not on the verify path at ANY K,
and a message that describes it as serving low K is wrong.

Asks tkv rather than inferring: the rule lives in
``tkv.runtime.attention._prefill_verify_min_block_m``. It is private, so
the flag-derived fallback mirrors it, and the last resort is the
pre-split behaviour (Turbo prefill owns block_m at or above the split-K max).
Each step degrades toward describing LESS Turbo prefill coverage, which is the
safe direction for a message: it can understate the new route, never
invent one.

Admission ceiling on MTP ``K`` (= ``block_m`` − 1) for tkv backends.

MTP verify on COMPRESSED tkv KV runs ``tokens_per_seq = block_m = K+1``.
Below the split-K register kernel's spill-free ceiling, that verify chunk
runs on the unified split-K MTP kernel (``tq_decode_splitk_mtp``,
inline-dequant on the centroid-quantized cache — NO decompress-to-bf16).
ABOVE it, ``TKVCore.forward`` re-dispatches the verify chunk to the
Turbo prefill verify path (tensor-core MMA, registers O(query-tile) not O(K)), so
``K`` scales to the decode-vs-prefill routing cap instead of tripping the
register kernel's ``block_m`` assert (the Turbo prefill high-K verify route).

The admission ceiling is therefore the SERVED verify ``block_m`` ceiling
minus one — read from turbo-attn's :func:`mtp_verify_max_block_m` (the
SINGLE source of truth: the Turbo prefill routing cap when the high-K route is
present, else the register kernel ceiling). It is NOT
:func:`tkv_max_block_m` (the split-K register bound), which now only marks
the split-K↔Turbo-prefill boundary, not the served K cap.

Route-aware + fail-loud: if turbo-attn predates the Turbo prefill high-K route (no
``mtp_verify_max_block_m`` symbol), fall back to the split-K register
ceiling so the cap can NEVER exceed what the installed engine can actually
serve.

Operational MTP-K ceiling for cudagraph capture (≤ the kernel serve cap).

See :data:`_MTP_CUDAGRAPH_CAPTURE_MAX_K`. Clamped to :func:`tkv_max_mtp_k`
so the capture ceiling can NEVER exceed what the installed kernel can serve
(e.g. an older turbo-attn whose served ceiling is below 15).

True if any spec routes the ``(B, K+1)`` MTP-verify cudagraph capture.

BOTH tkv families do — the compressed codec (``paged_kv:tkv-k<K>v<V>``)
and raw bf16 bypass (``paged_kv:tkv-bypass``). See
:mod:`arbi_serve.engine.capture_admin.decode`, which routes the verify
capture for "compressed codec AND tkv-bypass raw-bf16", so a bare
``paged_kv:tkv-`` prefix IS the right test for this question.

Deliberately NOT :func:`tkv_active`. That one answers "is the CODEC in
use", excludes bypass by design, and is correct for codec questions.
Using it to gate a CAPTURE cost is a category error: it let bypass
escape a ceiling it pays in full, while bypass's own scratch sizing
(``backends.tkv_bypass_backend._q_prescale_block_m_ceil``) was written
assuming the ceiling applied to it.

The largest servable MTP ``K`` for this config, and WHY it is that.

Two different ceilings, and which one binds depends on whether the boot
captures the full K ladder:

* ``--mtp-capture-full-k-ladder`` ON — one verify cudagraph per
  grows linearly and the ceiling is
  :func:`mtp_cudagraph_capture_max_k` (15).
* OFF (the DEFAULT) — only the served depth is captured, so the boot
  pays for ONE graph regardless of K and the ladder argument does not
  apply at any K. The binding ceiling is then the KERNEL's
  (:func:`tkv_max_mtp_k`, ~127).

cost the default config does not pay: a homogeneous server only ever
replays verify width ``S = n_draft + 1``. It also means a fixed-shape
tree, which has exactly ONE ``block_m``, is bounded by the kernel and
not by the ladder.

True if any registered backend spec is a tkv CODEC (``paged_kv:tkv-k{K}v{V}``).

Shared by the per-request admission gate (``engine/submission.py``) and
the boot-time gate (``cli/config_builder.py``) so the K cap can't drift.

``paged_kv:tkv-bypass`` (raw bf16, NO codec) shares the ``tkv-`` brand
prefix but must NOT count as codec-active — a bare ``tkv-`` prefix
match would misclassify it. Codec specs are precisely the
``tkv-k<K>v<V>`` grammar (see ``backends._parse_tkv``), so ``tkv-k``
is the exact signature.

Inputs that determine a deployment's pinned K. Two boots whose key
fields collide read each other's cached winner; anything that can change
the optimum K must live here or stale picks will silently leak across
deployments.

Three groups, and the boundary between them is the point:

  * WHAT was measured — the checkpoint and the shape it is served in
    (path, weight quant, dtype, registered + active backends, TP, context
    window, drafter) and the hardware it ran on (``gpu_name``).
  * WHAT it was measured ON — ``prompt_corpus_hash`` + ``regime``.
  * WHAT the measurement MEANS — ``calib_version`` (the manual lever for
    kernel/plumbing changes, guarded by :data:`KCAL_VERDICT_DIGEST`) and
    ``schema`` (the record's shape).

Kernel and library versions are deliberately absent; see
:data:`KCAL_VERDICT_SOURCES` for why and for what stands in their place.

Digest the checkpoint's own quantization declaration.

Reads ``<model_path>/config.json`` and hashes ``quantization_config``
together with the declared weight dtype. That dict is where every format
states its bit budget — exl3 ``bits`` / ``head_bits`` / ``mtp_bits``, AWQ
and compressed-tensors ``num_bits`` / ``group_size`` / ``format`` — so one
digest covers every quantizer without this module learning any of them.

``""`` when the file is absent or unreadable (a fresh checkpoint dir, a
test fixture): the field then carries no claim rather than a false one,
and the remaining key fields still separate the deployments.

Name the drafter a boot will actually run, as ``kind:path[@dtype]``.

The proposal source is the dominant term in accept-length-vs-K, and the
three sources are mutually exclusive at boot (see
``mtp_attach.build_mtp_driver``): an explicit DFlash head or external
draft model wins over the checkpoint's bundled MTP head. Reported in that
same precedence order so this string names the drafter that RUNS, not
every path that happens to be configured.

``""`` when MTP is off — there is no drafter and no K to calibrate.

The deployment-shape half of :class:`KCalibrationKey`, derived once.

The sweep WRITES a cache file and a later boot READS it; both sides must
hash identical values or every boot misses. Deriving those values in two
places is how the two sides come to disagree, so both go through here.

``max_context`` is normalized to an int: the boot resolves an ``"auto"``
window to a concrete number before MTP attaches, and a non-int reaching
this point means the caller had no resolved window to offer (``0``).

Rebuild the key a record was written under from its ``meta`` block.

``None`` when ``meta`` does not carry every field of the CURRENT key —
a record written by an older build cannot be re-keyed, because the
fields it never recorded are exactly the ones that would decide the
match. Keeping this the inverse of :meth:`KCalibrationKey
.as_canonical_dict` (rather than a hand-listed field copy) is what stops
a newly added field from being silently dropped on the read path.

Hash a prompt list canonically. Order-sensitive (a different
ordering picks a different cache slot, which is the desired
behaviour — sample order can shift the throughput average enough
to flip K=0 vs K=1 on borderline workloads).

Write a calibration record + return the on-disk path.

Creates ``root`` if missing. Atomic write via temp-rename so a
concurrent reader never sees a partially-written file.

The record carries the whole measured CURVE (``results``), the raw
per-arm ``samples``, and the ``measurement`` shape that produced them —
not just the winner. A cached ``best_k`` is a verdict that persists
across every later boot of this deployment; when it is wrong (e.g. a
degenerate curve whose winner was picked purely because ``max()`` breaks
ties by insertion order) the ONLY way to tell from the outside is to read
the numbers it was derived from. ``best_k`` alone is unfalsifiable.

Build the SamplingParams the sweep measures with — the SERVED config,
not a greedy proxy.

The optimum K is regime- and sampling-dependent: measuring greedy when
the deployment serves stochastic optimises the wrong objective. So this
pulls the model's own served sampling from ``eng.generation_defaults``
(the same ``generation_config.json`` values the request path backfills
onto an omitting request) — for Qwen3.6 that is temperature=1.0 /
top_p=0.95 / top_k=20, i.e. stochastic. A model that ships no
generation_config keeps the engine's code defaults (which are
themselves stochastic, temperature=1.0); calibration never silently
forces greedy.

Greedy is not the "clean" choice it looks like — it is the STRICTEST
verify rule and therefore the one that UNDER-measures speculation.
Greedy verify accepts a drafted token only when it is the argmax;
rejection sampling accepts any token with ``p(x) >= q(x)`` and accepts
probabilistically below that, so it accepts strictly more often. A
greedy sweep would therefore pick its K off a curve below the one the
deployment actually runs, on top of optimising a rule the server never
applies.

Output length is PINNED with ``ignore_eos``: every arm must decode the
same number of tokens or the arms are not comparable, and a request that
stops early can fall short of the discarded warmup window and contribute
nothing at all (see :func:`_measure_k_in_process`).

Refuse a sweep whose arms would all silently collapse to K=1 decode,
and return the sampling the arms will actually run (for the record).

``request_factory`` rewrites ``sampling.mtp_k = 0`` for any STOCHASTIC
request when ``cfg.mtp.rejection_sampling_enabled`` is off — a
correctness fallback (greedy verify on a temperature>0 request would
serve the wrong distribution). The sweep submits at the model's own
served sampling, which is stochastic for every modern instruct model, so
with that knob off EVERY arm K=1..k_ceiling is rewritten to K=0. The
sweep then measures the same K=0 decode path at every depth, reports a
flat curve, and concludes MTP does not pay — while never once having run
the drafter.

That is not a measurement, so it is refused here rather than resolved:
silently switching the arms to greedy would optimise a verify rule the
server does not apply (see :func:`_served_sampling_params`), and silently
measuring the flat curve produces an equally wrong verdict.

Render ``prompt`` through the model's chat template so the sweep
measures the SERVED prompt shape and inherits the deployment's thinking
regime.

A raw-string submission exercises NEITHER regime: it skips the chat
template entirely, so the drafter never saw the peaked reasoning prefix
(thinking) or the flat prose prefix (non-thinking) that actually drive
accept rate. Leaving ``thinking=None`` lets the model's own template
default govern (Qwen3.6: undefined ⇒ thinking ON), which is the DEFAULT
regime the server will serve — so the calibrated K matches what the
deployment actually runs.

Falls back to the raw prompt (with a one-time warning) only if the model
ships no chat template — never silently, so a template-less model does
not masquerade as a regime-calibrated one.

The resolved thinking regime the sweep calibrates for — a component of
the cache key so two deployments serving different default regimes never
share one cached K. ``""`` when the model has no regime axis (one cache
slot is then correct).

Model-agnostic probe: render the template's generation prompt three ways
(undefined / true / false) and compare by string identity — no hardcoded
``<think>`` marker. ``undefined == true`` ⇒ the model's default is
thinking; ``undefined == false`` ⇒ prose. It returns ``""`` (single
regime — one K is right for the whole deployment) when the template is
not thinking-aware (``true == off``) or its default matches neither. This
is DERIVED, not a fallback value: a model with one regime genuinely wants
one cache slot. It is a PURE function of the template, so the write and
load paths on the same deployment always agree on the label.

Watch ONE live request and return its ``(tokens, seconds)`` steady-state
decode window.

The window OPENS the first time the request is observed to hold more than
``warmup_tokens`` output tokens and CLOSES when ``finish_reason`` is set.
Both stamps are taken HERE, on one monotonic clock
(:func:`time.perf_counter`). That is deliberate and load-bearing:

  * It never reads ``Request.first_token_time``. That field is stamped
    with ``time.time()`` (epoch), while a ``perf_counter()`` finish stamp
    is monotonic — subtracting one from the other produces a meaningless
    (and possibly hugely negative) interval. Owning both stamps makes
    that whole class of bug unrepresentable.
  * Opening the window LATE excludes the per-request approach to steady
    state (first drafter-chain build, first touch of this K's captured
    verify/drafter graphs, allocator growth for the wider verify chunk) —
    a cost that grows with K and therefore INVERTS the ranking when it is
    charged to a short arm. See :data:`SWEEP_WARMUP_TOKENS`.

``tokens`` is measured across the window (``final - at_open``), not from
the request's total, so a coalesced ``new_token_event`` that delivers
several tokens at once is still counted exactly.

Returns ``(0, 0.0)`` when the request never reached the window (it
finished at or below ``warmup_tokens``); the caller aggregates and fails
loud if NO request contributed.

One measured sweep arm: the scored rate + the accept-length witness.

``tok_s`` is the decision variable. ``accept_len`` is diagnostic — see
:func:`_measure_k_in_process` for why an arm that cannot show a rising
accept length is not credible even when its tok/s looks plausible.

Submit ``prompts`` at speculation depth ``k`` against a LIVE engine
and return STEADY-STATE decode tokens/sec.

Drives the engine through ``eng.asubmit`` — i.e. the REAL serving
path — so under TP>1 the per-step plan broadcast carries the work to
every rank exactly as a normal request would; there is no separate
TP1 measurement engine. Each prompt gets a fresh random salt so the
radix prefix cache cannot turn a later sweep cell into a cache hit.

Measures the SERVED config (stochastic canonical sampling + chat-template
-rendered, regime-resolved prompts) and scores by Σ tokens / Σ seconds
over each request's steady window (:func:`_watch_steady_window`) — never
``out/total-wall``, which folds in the prefill/TTFT the house rules
forbid crediting to decode throughput, and never a window that includes
the K-dependent cold ticks.

Every handle is watched CONCURRENTLY. Watching them in submission order
would stamp a request's finish only after every EARLIER request has
finished, so requests 2..n would be charged the tail of request 1 — an
inflation that grows with how unevenly the arms complete.

Fails LOUD when no request produced a window: a 0.0 return is
indistinguishable from "this K is infinitely slow", and an all-0.0 result
dict would silently look like a valid measurement.

Also returns the arm's ACCEPT LENGTH (mean tokens committed per decode
tick — see :func:`_arm_accept_len`, ``0.0`` at K=0). It is not part of the
decision — tok/s is — but it is the independent witness that the arm ran
the drafter at the depth it claims. Accept length must RISE with K toward
saturation; an arm whose tok/s moves while accept length stays pinned at
0 or flat is measuring something other than speculation.

``(row_tokens, row_ticks)`` — the EFFECTIVE accept accounting, counted
over every decode row-tick including cold ones.

``(0, 0)`` when the counters are absent or never fired (they are bumped on
the SPMD commit path; a TP1 eager boot leaves them flat), which the caller
reads as "fall back to the per-request ratio".

Mean tokens committed per decode tick over this arm.

Prefers the EFFECTIVE counters (``mtp_row_tokens / mtp_row_ticks``,
deltas across the arm) — the honest, TPOT-determining number, and the one
the hardware table in ``docs/mtp-k-sweep-calibration.md`` reports.

Falls back to the per-request ``mtp_accepted / (mtp_proposed / k) + 1``
when those counters did not move (they are bumped at the SPMD rank-0
commit, so a TP1 eager boot has none). That fallback is a WARM-ONLY
average and reads HIGH: a cold row drafts 0 and commits 1, adding nothing
to either counter while still consuming a whole tick. It is diagnostic
either way — the decision is tok/s — but the two must not be conflated,
so the fallback is not silently presented as the effective number.

Run discarded measurement passes at ``k`` until the arm is warm.

Returns ``(passes_run, converged)``.

"Warm" is DETECTED, not assumed after a fixed count. A fixed count is not
sufficient: the residue it leaves is K-dependent — a deeper arm compiles
strictly more verify widths and drafter rungs — so it biases the ranking
in the one direction that flips a rising curve into a falling one.

Each pass is IDENTICAL in shape to a scored pass (same prompts, same
``max_tokens`` / ``warmup_tokens``). That is required twice over: a
different-length warmup would not compile the shapes the scored pass uses,
and the pass-to-pass comparison below would not be measuring the same
quantity it is trying to declare stable.

``min_passes <= 0`` disables warmup entirely (test harnesses that inject a
canned curve, where there is nothing to warm).

Sweep ``K`` ascending against a built, running engine via per-request
``mtp_k`` and return ``(best_k, {k: median_tok_s}, {k: accept_len})``.

``accept_len`` is diagnostic, never a decision input — but a correct
sweep MUST show it rising with K toward saturation, so it is what
distinguishes a real curve from a plausible-looking artifact.

Works at ANY TP because it drives the already-built engine at its
real topology — it never constructs a measurement engine. The engine
MUST have been built with an MTP driver whose ``max_k >= k_ceiling``
(so every swept K hits a captured graph; calibrate / boot both build
at ``k_ceiling``). Picks the winner by c=1 STEADY-STATE decode tok/s —
the same quantity a dedicated ``--mtp-n-draft K`` boot reports.

Arm shape — two independent discards, for two different cold costs:

  * ``warmup_tokens`` decode tokens are dropped from the head of EVERY
    scored request (each request's own approach to steady state — see
    :data:`SWEEP_MAX_TOKENS`);
  * whole passes are discarded until the arm is warm (the once-per-process
    JIT/autotune of this K's shapes — see :data:`SWEEP_MIN_WARMUP`).
    Warmup is ADAPTIVE, not a fixed count: it runs at least
    ``num_warmup`` passes and keeps going until two consecutive agree
    within ``warm_tol``, capped at ``max_warmup``.

Both discards matter and neither substitutes for the other. Critically,
warmup passes are IDENTICAL in shape to scored passes (same prompts, same
``max_tokens``, same ``warmup_tokens``) — a warmup that ran a different
length would not compile the shapes the scored pass then uses, and would
warm nothing that counts. It is also what makes comparing a warmup sample
against the next one a valid convergence test.

Then ``num_runs`` scored passes; the arm's score is their MEDIAN, never a
mean — a mean over few samples is fully exposed to one cold outlier, which
is precisely the residue this warmup exists to exclude.

Early stop: the sweep gives up after ``decline_patience + 1``
CONSECUTIVE declines, and a drop within ``decline_rel`` of the previous
arm is NOT a decline (it is noise — the top of the curve is flat because
accept length saturates). Patience is 2, not 1, for the same reason it
was ever more than 0: the tok/s curve is not unimodal at the low end once
a one-block head is CHAINED. The win is ``1 + a + a² + … + a^K`` while
the cost is a verify at ``S = K+1`` plus ``K`` drafter steps, so the
shallow depths buy the least-compounded tokens at nearly the full
overhead and can sit BELOW K=0, while the series only starts paying at
K=2..3 and peaks at a deeper K before falling again as accept saturates
and drafter cost keeps growing. Patience 1 tolerates a one-rung dip; a
TWO-rung shallow dip would still stop the sweep at the exact failure
this patience exists to prevent —
pinning ``best_k=0`` without ever measuring the depths that win. The
price of patience 2 is one more measured arm; the price of being wrong
is the whole speedup, cached, forever. Past the real peak the curve is
monotone, so it still terminates early.

Refuses a degenerate curve (all arms zero, or every arm identical)
rather than returning ``max()``'s first-key-on-ties — which is ``K=0``
and would otherwise report "MTP detached" from a measurement that had
silently produced no numbers at all.

The swept ``K`` with the highest measured tok/s — or a loud refusal.

Standalone from :func:`sweep_k_in_process` so the DECISION is testable
against a canned curve without a live engine, and so the refusals below
guard every caller.

Two degeneracies are refused rather than resolved, because ``max()`` on a
dict resolves ties by INSERTION ORDER and the sweep inserts ``K=0``
first — so any measurement that fails to distinguish the arms
automatically "wins" with MTP off, which is a verdict that looks
identical to a real one and gets cached:

  * no arm measured above zero, and
  * every arm measured EXACTLY the same (a harness that is not reading
    the clock it thinks it is reading produces this, not a real engine).

A genuine exact tie between two working arms is broken toward the
SMALLER K — less VRAM for the capture ladders, same throughput.

DSpark ``VanillaMarkov`` head — low-rank bigram logit bias for DFlash.

A parallel block-diffusion drafter (:class:`~arbi_serve.spec_decode.dflash.
DFlashDraftModel`) predicts every block position in ONE forward from
mask-token inputs, so position ``k`` never sees the token actually chosen
at ``k-1`` (measured "suffix decay": per-position top-1 falls steeply with
block depth). The Markov head corrects that cheaply with a per-position
additive logit bias conditioned only on the previous REALIZED token:

    B(x_{k-1}, ·) = W2( W1[x_{k-1}] )     W1 ∈ R^{V×r},  W2 ∈ R^{r×V}

    corrected_logit_k = base_k + B(x_{k-1}, ·)

i.e. a rank-``r`` factorized bigram (transition-matrix) bias. At inference
the block is produced LEFT-TO-RIGHT (semi-autoregressive): position ``k``'s
logits are biased by the token chosen at ``k-1``; position 0's predecessor
is the verified anchor token before the block. Cost per position: one
``r``-wide embedding lookup + one ``[r×V]`` GEMV + add.

The head is self-contained — it carries its OWN ``markov_w1`` /
``markov_w2`` and never touches the (borrowed) target ``embed_tokens`` /
``lm_head``. Parameter names match the DSpark checkpoint keys
``markov_head.markov_w1.weight`` / ``markov_head.markov_w2.weight`` exactly
so the strict state-dict load binds them without remapping.

The semi-AR draft loops that consume this bias live in
:meth:`~arbi_serve.spec_decode.dflash_driver.DFlashDrafter.draft`; the
verify/accept machinery is untouched (greedy verify re-checks every drafted
token against the target argmax regardless of how the draft was formed, and
the rejection sampler is lossless for any reported full ``q``).

The one method the semi-AR block realization needs from a head.

Both DSpark head shapes satisfy it — :class:`VanillaMarkov` (DFlash
checkpoints) and the DeepSeek-V4 in-model ``DSparkMarkovHead`` — which
is what lets the two drafters share one realization loop. The returned
``(B, V)`` must be a FRESH tensor: the loop accumulates the base logits
into it in place.

``base_row + B(prev)`` at ``base_row``'s dtype, in the head's own buffer.

``compute_step_bias`` returns a FRESH unaliased ``(rows, V)`` (a Linear
output), so it is the accumulator: ``bias.add_(base_row)`` is the same
elementwise kernel as ``base_row + bias`` over the same operands, and
IEEE-754 addition is commutative, so the biased logits — and therefore
the realized argmax — are bit-identical to the two-tensor form. Only ONE
full-vocab tensor exists per position instead of three, which is what a
fixed-shape capture of this chain would otherwise have to keep resident.

The cast is guarded rather than unconditional: a head bound at the
base-logit dtype (:meth:`VanillaMarkov.bind_logit_dtype`, and what
``DFlashDraftModel.from_checkpoint`` already leaves behind) materialises
no second full-vocab tensor here at all.

LEFT-TO-RIGHT greedy block realization → ``(rows, steps)`` int64.

Step ``j`` is ``argmax(base_j + B(prev))`` where ``prev`` is the token
realized at ``j-1`` (``first_prev`` at ``j == 0``). Two invariants hold:
rows are INDEPENDENT (a row's argmax reads only its own logits), and the
chain is a PREFIX in the step axis (running further steps never rewrites
an earlier one), so the ``[:n, :k]`` sub-block of a ``(rows, steps)``
call equals the same call made at exactly ``(n, k)``.

Device-only and fixed-shape — no ``.item()``, no host reads, no
data-dependent control flow — so the whole chain is cudagraph-capturable.
Per step the only device work is the rank-``r`` lookup, the ``[r×V]``
GEMV, the in-place bias accumulation, and an ``argmax`` that writes its
ids straight into the result column: no full-vocab temporary beyond the
head's own bias buffer, and no per-step ``(rows,)`` temporary or copy.

``toks_out`` writes the realized ids into a destination the caller owns
(a cudagraph capture's persistent output buffer) instead of a fresh one.

Return the Markov head for a checkpoint config, or ``None`` when
``markov_rank == 0`` (plain DFlash — no head, byte-identical build).

Only the ``vanilla`` head is implemented; any other declared
``markov_head_type`` fails loud rather than silently degrading to a
bias the checkpoint was not trained for.

Per-row bias over the full vocab: ``(B,) int64 → (B, V)``.

Emitted in the head's parameter dtype. :meth:`bind_logit_dtype`
binds that dtype to the base-logit dtype once, so the semi-AR hot
path never casts a ``(B, V)`` tensor per drafted position.

The returned tensor is always FRESH and unaliased (a Linear
output), which is what lets the caller accumulate the base logits
into it in place.

Bind the head's parameters so ``compute_step_bias`` emits ``dtype``.

Invariant after this call: ``compute_step_bias(...).dtype == dtype``,
i.e. the bias lands on the base-logit manifold with no per-position
cast. Idempotent, and a strict no-op (no copy, no rebind) when the
head already carries ``dtype`` — which is the state
:meth:`DFlashDraftModel.from_checkpoint` leaves it in, since it
materialises the whole draft model at the engine dtype.

MTP draft heads + speculative-decoding driver — public facade.

The implementation is split across sibling modules to keep each file
under the line ceiling; this module re-exports the genuine public (and
the load-bearing private) surface so existing importers and tests that
do ``from arbi_serve.spec_decode.mtp import …`` keep resolving:

  - :mod:`arbi_serve.spec_decode.mtp_draft_meta` — pure host-side
    helpers (``_nvtx``, ``mrope_pos_offset``,
    ``build_draft_step_metas_from_lists``, ``draft_step_bypass_safe``,
    ``decorate_draft_metas``).
  - :mod:`arbi_serve.spec_decode.mtp_base_head` — :class:`MtpHead` +
    :class:`_MtpDecoderBlock`.
  - :mod:`arbi_serve.spec_decode.mtp_driver` — :class:`MtpDriver`
    (draft + verify + accept).
  - :mod:`arbi_serve.spec_decode.mtp_verify` — the verify-pass plan
    builder + verify loop + recurrent / block-table rollback helpers
    (kept co-located; rollback ordering is correctness-critical).
  - :mod:`arbi_serve.spec_decode.mtp_strategy` — :class:`MtpStrategy`
    (alias :class:`MtpSpecDecode`), the engine binding.

MTP draft head + decoder block.

:class:`MtpHead` is N transformer blocks reusing the main model's
``embed_tokens`` and ``lm_head``. N draft blocks, and produces K draft proposals per sequence.
:class:`_MtpDecoderBlock` is one such block (DeepSeek V3 layout — a
standard decoder layer plus the ``eh_proj`` mixing projection).

One MTP draft transformer block.

DeepSeek V3's MTP heads are structurally identical to the main
model's decoder layer with one extra "context" linear that mixes
the previous-step hidden state with the next-token embedding
before the standard pre-norm + attention + MLP block runs.

Self-contained: the engine layer code does not special-case MTP
layers — they run a standalone forward over their own small,
disposable attention state independent of the main model's KV
cache. Construction takes a ``layer_factory`` (the host model's
decoder-layer constructor) so :class:`MtpHead` works across model
families without duplicating block code.

N draft transformer blocks reusing the main model's embed + lm_head.

Construction:
  - ``n_draft`` is the number of speculative tokens this head
    proposes per main-model step (the K).
  - ``embed_tokens`` and ``lm_head`` are references to the main
    model's modules — NOT copies. The whole point of MTP is
    reusing those tied weights.
  - ``decoder_layers`` is a list of N decoder modules. Constructed
    fresh (the MTP draft blocks are NOT the same instances as the
    main model's blocks), but type-compatible so the per-layer
    ``AttentionBlock`` forward runs unchanged.

Forward path. Given the main model's ``last_hidden`` at the K=0
position and the next-token id at K=0, run K MTP blocks
autoregressively, producing K draft token ids. Each block emits a
new hidden state which (a) feeds the next MTP block (when k+1 < K)
and (b) projects through ``lm_head`` for that step's draft logits.

``forward`` returns ``(K, num_seqs)`` int draft tokens plus the
final ``(K, num_seqs, hidden)`` hidden states (the latter is kept
for diagnostics; the engine doesn't consume it).

Assumes the main model's "last hidden state" is the post-final-
norm activation just before ``lm_head``. Models whose ``forward``
does not expose that path (e.g. fused rms-norm + lm_head kernels)
must surface it via a hook before an MTP head can plug in.

Run K draft steps and return ``(K, B)`` int draft tokens.

``state_views`` / ``attn_metas`` / ``attn_ops`` are length-K
sequences; the engine builds them by allocating K draft cache
pages per request before calling :meth:`forward`. The driver
owns this allocation; the head is purely the math.

Sampling. When ``sampling_params`` is provided the head samples
each per-step draft token from the request's
``temperature/top_p/top_k/min_p`` distribution via the
graph-safe :class:`GraphSafeRejectionSampler` (batched
Gumbel-max with a persistent seed buffer; no
``torch.multinomial`` per row) and returns the FULL per-slot
drafter distribution ``q_full`` alongside the tokens (shape
``(K, B, V)`` for q_full, ``(K, B)`` for tokens). The
rejection-sampling residual sampler in
:func:`verify_and_accept` consumes ``q_full``.
``temperature == 0`` rows collapse to argmax with a point-mass
on the picked token. Default is greedy argmax with no
probabilities returned (legacy contract).

``seed_buf`` is the engine-owned ``int64 (1,)`` step seed
(typically ``Engine.mtp_seed``); shared with the verify pass
for cross-step determinism.

Pure host-side helpers for the MTP draft chain.

These are the stateless building blocks the :class:`MtpDriver` and the
TP-worker bridge share: the NVTX sub-phase range, the per-step
:class:`AttnPagedKVMeta` reconstruction, the host-side TKV
``bypass_safe`` predicate + per-step TKV decoration, the stochastic
replay input packer, and the M-RoPE position offset. Pulled out as free
functions so both rank 0's live driver and the worker-rank captured
replay derive byte-identical shapes / values across ranks.

Host list → device tensor WITHOUT a per-step pageable-H2D sync.

``torch.tensor(list, device=cuda)`` stages the Python list into PAGEABLE
host memory and issues a SYNCHRONOUS H2D (the CUDA runtime serializes a
pageable copy as ``cudaMemcpyAsync`` + ``cudaStreamSynchronize``), so the
host blocks until the copy stream drains. On the MTP decode path these
builds run AFTER the verify forward is enqueued, so the first such copy
stalls the host for the whole forward every step on the dense c=1 decode
path.

Building the source in PINNED host memory instead lets
``.to(device, non_blocking=True)`` run as a true async copy: no host
stall, byte-identical device values (bit-neutral). The pinned block is
owned by the CUDA caching host allocator, which defers its reuse until the
recorded copy event completes, so it is safe even though the CPU tensor is
a temporary. Kernels that consume the result are enqueued on the same
(current) stream after the copy, preserving ordering.

Falls back to a plain ``torch.tensor(..., device=...)`` for a non-CUDA
device (``pin_memory`` requires CUDA; CPU/worker-reconstruct paths keep
identical values).

``src.index_select(0, idx)`` that SKIPS the launch on an identity select.

The MTP verify seed path re-slices the all-rows frontier-repair gathers
(``last_committed`` / ``bonus_hidden`` / ``repair_slots`` / ``repair_pos``
/ ``block_table``) per drafter K-bucket. Under the R1 uniform-K invariant a
verify step is a SINGLE bucket over all MTP rows, so that per-bucket ``idx``
is the IDENTITY ``[0, 1, ..., src.shape[0]-1]`` — and an ``index_select``
with a contiguous identity index still launches a gather/copy kernel and
materialises a fresh tensor. Returning ``src`` unchanged in that case is
VALUE-IDENTICAL (an identity gather is a pure copy, so the accept-critical
drafter inputs stay bit-exact) and removes the launch. A strict sub-bucket
(mixed per-request K) still gathers via a pinned-staged :func:`h2d_async`
index — the SAME device values as before, bit-neutral.

``idx_host`` is the HOST list of source rows (known at plan time), so the
identity test is a pure host comparison — no device read, no sync.

NVTX range for an MTP sub-phase, read off the engine's profiler.

Returns the profiler's ``nvtx_range`` context (which only emits a real
``range_push/pop`` under ARBI_PROFILE=nvtx) or a cheap no-op context
when no profiler is attached (ARBI_PROFILE=off — the production path).

Per-request RoPE position offset for spec-decode tokens.

Vision (M-RoPE) requests compress image-token positions, so the
decode-continuation absolute position is ``flat_len + _mrope_delta``,
not ``flat_len`` (see ``ModelRunner._ensure_media_encoded`` /
``_attach_multimodal`` decode branch). Text-only requests (no
``_mrope_delta``) return 0 — the spec-decode position arithmetic
stays bit-identical to the pre-vision path. The three spec-decode
position builders (drafter chain metas, verify flat batch, and the
TP broadcast mirror) add this offset so draft / verify tokens RoPE
at the SAME absolute position the main-model decode would use.

Construct K per-step :class:`AttnPagedKVMeta` from CPU-side lists.

The arithmetic mirrors :meth:`MtpDriver._build_draft_step_metas` —
one query per row, ``seq_lens[i] = pre_lens[i] + s + 1``,
``positions[i] = pre_lens[i] + s``, ``slot_mapping[i] =
per_req_slots[i][s]``. Pulled out as a free function so the
TP-worker bridge can reconstruct identical metas from the
broadcast payload (same shapes / values across ranks; only the
KV-shard widths differ on each rank's local layer slab).

Slab slot index of a request's LAST committed KV cell.

``committed_len`` is the request's committed KV length (page-table
``length`` with no draft slots outstanding); the frontier cell is
sequence index ``committed_len - 1``. Slot index follows the page
table convention ``page_id * block_size + offset_within_page``
(:meth:`FlatPageTable.allocate_slots`), derived from the request's
block-table row so rank 0 and TP workers compute the identical slot
from the broadcast payload.

Shift the drafter chain onto the fill position convention.

The MTP head's training pairing places ``(hidden(p), embed(token at
p+1))`` at RoPE position ``p`` — the HIDDEN's position. The verify
fill pass honours that (cell ``p`` gets the pair at position ``p``).
This helper derives the per-step layout from the SAME inputs both
ranks already share:

  - chain step ``s`` sits at position ``pre_len - 1 + s`` with
    ``seq_len = pre_len + s`` (feed the unshifted lists through
    :func:`build_draft_step_metas_from_lists` with
    ``pre_lens - 1``),
  - step 0 writes INTO the frontier cell (slot of sequence index
    ``pre_len - 1``) and steps ``1..K-1`` use the first ``K-1``
    allocated draft slots (cells ``pre_len .. pre_len+K-2``). The
    LAST allocated draft slot goes unused (the allocation count is
    kept so page-table extension / rollback and the TP broadcast
    shapes stay identical).

Returns ``(pre_lens_shifted, per_req_slots_shifted)``.

One-query-per-row :class:`AttnPagedKVMeta` targeting each request's
LAST committed MTP-KV cell.

The verify/decode MTP fill pass pairs cell ``p`` with
``embed(token[p+1])`` (next-token shift) but pairs each request's
last in-span cell with its OWN embedding (the successor is unknown
at fill time). After verify+accept that cell — the per-request
frontier, sequence index ``committed_len - 1`` — permanently holds a
K/V computed from the wrong embedding; the next span starts one
position later and never rewrites it. This meta drives the repair
rewrite (:meth:`Qwen3_5MtpHead.refill_kv_cell`): ``slot_mapping``
points at the frontier cell's EXISTING slot, ``positions`` is the
cell's own RoPE position (matching what the fill pass used), and
``seq_lens`` covers the committed history so the (discarded)
attention read stays in-bounds. Reuses
:func:`build_draft_step_metas_from_lists` (k=1) so the shapes /
dtypes are byte-identical to a drafter chain-step meta across ranks.

Device-tensor variant of :func:`build_frontier_repair_meta`.

Used by the async (optimistic-advance) verify path, where the
frontier cell index depends on the DEVICE ``n_accepted`` tensor:
the caller gathers ``slot_mapping`` / ``positions`` from the verify
span's staged view at flat index ``cu_q[row] + n_accepted[row]``
(no host sync). ``seq_lens_host`` only bounds the discarded
attention read; any host value within the row's allocated pages and
covering the frontier cell is valid (the async caller passes the
optimistic page-table length).

DIAGNOSTIC gate.

``ARBI_MTP_FRONTIER_EAGER=1`` routes the frontier MTP-KV cell write
through the EAGER ``_repair_kv_frontier`` path on the SYNC drafter
paths (both ranks), and redirects the chain's step-0 scatter OFF the
frontier cell (see :func:`redirect_step0_off_frontier`) so the eager
write is what every chain step reads. Chain geometry (positions /
seq_lens / cu_seqlens) is byte-identical to the production path;
only the frontier cell's WRITER changes (eager kernels outside any
captured graph, at M=B, instead of the captured chain's in-graph
step-0 write). Default OFF = byte-identical production behaviour.
Read per call (eager paths only) so a container env flip needs no
code reload.

DIAGNOSTIC (see :func:`frontier_eager_diag_enabled`): send chain
step 0's K/V scatter to the row's UNUSED last allocated draft slot
instead of the frontier cell.

Under :func:`shift_chain_to_fill_convention` the last allocated
draft slot (sequence index ``pre_len + k - 1``) goes unused, and no
chain step's read window (``seq_len = pre_len + s``, max ``s = k-1``
⇒ cells ``0 .. pre_len + k - 2``) ever reaches it — so the
redirected write is a harmless scatter into an allocated page. The
frontier cell then keeps whatever the eager repair wrote.

Host-side TKV ``bypass_safe`` for one drafter chain step.

The drafter decodes one query per row, so ``q_len_i == 1`` and
``seq_len_i == pre_lens[i] + step + 1``. ``bypass_safe`` (first-chunk
prefill: ``q_len_i == seq_len_i`` for all i) therefore holds iff
every ``pre_lens[i] + step == 0`` — known purely from host ints, no
device read. Bit-identical to ``compute_bypass_safe`` on the device
tensors, computed without the D2H sync.

Decorate drafter per-step metas with the PAGED_KV per-step buffers.

The drafter chain builds bare :class:`AttnPagedKVMeta` objects (one
query per row) that never route through the PAGED_KV metadata
builder's ``build`` / ``_finalize``, so the per-step page buffers the
decode kernel needs (the TKV ``_tq_*`` mirrors; the tkv-bypass
``_bt_*`` CSR triplet) default to ``None`` and the decode loader
crashes. This dispatches to whichever PAGED_KV builder exposes a
``decorate_external_meta`` method (``TkvMetadataBuilder`` and
``TkvBypassMetadataBuilder`` both do; tkv-bypass / GDN builders
don't and are left untouched), applying the full per-step wiring to
EACH step's meta. The per-step scratch is overwritten by the next
step, so each step re-decorates (the K steps run serially on the live
chain).

``bypass_safe`` is consumed only by the TKV decorator (which threads
it into ``TQRunState``); the tkv-bypass decorator ignores it (raw bf16
KV has no bypass-vs-compressed gate).

MTP speculative-decoding driver.

:class:`MtpDriver` wraps the engine step:
  1. Draft phase: run :class:`MtpHead` over the main model's last
     hidden state to produce K draft tokens per sequence.
  2. Verify phase: run the main model with the K draft tokens
     appended to each sequence's input (extends batch by K
     tokens × B sequences).
  3. Accept phase: longest matching prefix per sequence, plus the
     "bonus / recovery" token from the main model (see
     :mod:`arbi_serve.spec_decode.verify`).

The driver copes with mixed-K batches: requests with the same ``K``
run together; ``mtp_k == 0`` requests fall back to the K=1 path in
the same step. :class:`Engine` slot-routing decides whether the step
needs the driver at all, or can use the fast K=1 path.

The two large method clusters — draft dispatch (captured / stochastic /
live routing) and the captured-chain replay + per-step draft-meta
machinery — live in sibling mixin modules
(:mod:`arbi_serve.spec_decode.mtp_driver_draft`,
:mod:`arbi_serve.spec_decode.mtp_driver_replay`), and the module-level free
functions + flag-truth counters in
:mod:`arbi_serve.spec_decode._mtp_driver_ops`. The public import surface
(``from arbi_serve.spec_decode.mtp_driver import MtpDriver`` /
``effective_row_k`` / ``stochastic_draft_temp_threshold`` /
``stochastic_draft_row`` / ``DrafterChainBreakerError``) is unchanged.

Wraps the engine step with draft + verify + accept phases.

Handles a step that mixes MTP requests
(``sampling.mtp_k > 0``) and non-MTP requests (``mtp_k == 0``):
requests sharing the same K group into one verify pass, and
``K == 0`` requests run through the engine's K=1 path inline. A
request's effective K is capped at :attr:`max_k` — the DRAFT DEPTH
this boot was built for, NOT ``head.n_draft`` (the blocks the
checkpoint ships). The two are distinct: :meth:`draft` chains a
single block autoregressively, so ``max_k > head.n_draft`` is a
first-class path and is what makes speculation pay on the one-block
Qwen 3.5 / 3.6 heads (the win is ``1 + a + a² + …`` — it needs
depth). See :mod:`arbi_serve.engine.mtp_depth`.

The draft-dispatch methods (:meth:`draft` and its captured / stochastic /
live sub-paths) live on :class:`_MtpDraftDispatchMixin`; the captured-chain
replay + per-step draft-meta machinery on :class:`_MtpDraftChainMixin`. Both
reference ``self.<attr>`` bound in :meth:`__init__` (duck-typed). The mixins
precede :class:`DrafterCircuitBreakerMixin` in the MRO.

Whether THIS boot gates the chain's depth per slot.

Read fresh rather than cached at construction so a live config
override flips the drafter and the verify plan together — a plan
still reading a short chain as a cache miss would collapse every
gated step to plain decode.

Set the maximum speculative draft-token count per step.

NOT a runtime serving-depth raise on its own. ``max_k`` is a FIXED
boot capacity: the drafter+verify cudagraphs are captured for
``(B, K+1)`` shapes keyed on it (``cudagraph_admin``), and the
verify snapshot pool is sized ``T = max_k + 1``. Bumping this
attribute alone advertises a depth with no captured graph and (for
tkv) no kernel above block_m=12, so a request at the higher K would
miss its graph (TP>1 cudagraph → the c=1 deadlock guard in
``_refuse_or_warn_eager_miss``) or trip the kernel block_m assert.
Raising the SERVABLE depth requires a reboot with a larger
``--mtp-n-draft`` (re-capturing graphs at runtime on a 27B TP2
engine wedges the NVIDIA driver — a known reboot-only hazard). This
setter exists only for the deferred-KV-resize reattach path and
tests; the snapshot pool's idempotent-grow
(``attach_mtp_snapshot_buffers``) covers ONE of the ~4 sub-steps a
true runtime raise would need, not the cudagraph re-capture.

Drafter vocab size (the bundled head reuses the main lm_head).

Per-shard width: ``weight.shape[0]`` on a dense head. A
weight-quantized head (AWQ mirror) carries no dense ``.weight``;
its ``out_features`` IS the local shard width (quant TP mixins
store LOCAL out), so the two reads are the same quantity.

FULL (logical) vocab width the drafted argmax tokens index.

Under TP the lm_head is VOCAB-PARALLEL: ``head._lm_head.weight``
holds only this rank's ``vocab / tp_size`` shard, but the head's
forward all-gathers + trims so the returned logits — and thus the
argmax draft tokens — span the FULL vocab ``[0, dims.vocab_size)``.
Sizing the one-hot ``q`` from the sharded :attr:`vocab_size` would
scatter a full-vocab token index into a half-width buffer → CUDA
``index out of bounds``. Prefer the engine model's logical
``dims.vocab_size``; fall back to the (TP=1-correct) sharded width.

Width of the row THIS DRAFTER forms — what a ``q`` buffer must hold.

:attr:`full_vocab_size` unless ``--draft-vocab-prefix`` gave the
drafter its own lm_head over ids ``[0, N)``, in which case the sampled
proposal is ``N`` wide and the verify seam pads it (see
:func:`~arbi_serve.spec_decode.drafter.resolve_draft_vocab_size`).
Sizing a captured drafter buffer from :attr:`full_vocab_size` while the
head emits ``N`` is a shape error at the first chain step, not a
harmless over-allocation.

INFO-log a STRUCTURAL live-chain fallback once per reason.

Capability boundaries (TP>1 / min_p — see
:meth:`_draft_captured_stochastic`) are not failures: output is
identical, only the per-step drafter cost differs. One INFO at
first use makes the perf envelope visible without per-step spam.

No-op for the bundled head.

The bundled head writes K/V into the main pool's draft slots
via :meth:`MultiStatePool.allocate_draft_slots`; the engine's
verify driver already calls
:meth:`FlatPageTable.finalize_draft_slots` after acceptance,
so rollback is already done by the time we reach here.

Cross-check the bundled head reuses main embed_tokens / lm_head.

Runs once at construction (when engine is known), keeping the
bundled-head invariant baked into the type rather than a
side-condition of the engine attach hook. Skipped when engine
is None (test path constructs without an engine and re-checks
during ``attach_mtp_driver``).

A low-rank head that carries its OWN draft-dim lm_head
(``bundled_lm_head_identity is False`` — the Gemma-4 assistant)
does not satisfy the shared-embed/lm_head convention; its own
build path validates its verifier-embedding binding instead.

True iff the head's attention block calls a real per-layer
``AttnOp`` (Qwen3_5MtpHead). Stub heads + the
no-op path return False.

Detection is structural: the head exposes a ``layers``
ModuleList whose first entry has a ``self_attn`` whose
forward signature accepts ``(hidden, positions, state_view,
attn_meta, rope_cache, attn_op)``. The marker
attribute ``_uses_real_attention = True`` on the attention
block signals real attention; absence falls through to the
legacy no-op chain (test-stub heads).

Resolve the MTP marker spec's ``layer_idx`` on the engine model.

The arch (Qwen3.5 / 3.6) appends a single :class:`LayerSpec`
with ``is_mtp_layer=True`` at the tail of ``layer_specs``;
its ``layer_idx`` indexes the per-layer pool slab + attn_op.

Update per-request accept counters from a verify pass.

Engine-side commit (token append, page-table extend, stop
condition checks) is still owned by the engine — the driver
only updates the speculative-decoding metric counters here.

Draft-dispatch method cluster for the MTP driver.

:class:`_MtpDraftDispatchMixin` holds the
:class:`~arbi_serve.spec_decode.mtp_driver.MtpDriver` methods that route a
draft request through the captured / captured-stochastic / live chains
(:meth:`draft` and its sub-paths). The methods reference
``self.<attr>`` bound in :meth:`MtpDriver.__init__` and methods provided by
the sibling :class:`~arbi_serve.spec_decode.mtp_driver_replay._MtpDraftChainMixin`
(duck-typed via the composed class).

Repeat ``t``'s last row until it has ``width`` rows.

Pad rows carry a real row's values, so the padded chain reads in-bounds KV
and its extra outputs are discarded by the caller. ``t`` unchanged when it
is already ``width`` rows.

Stamp the cuMem ``capture.cudagraphs`` tag over a capture region.

``torch.cuda.graph(pool=...)`` routes the graph-private working set through
the named pool's allocator without passing through ``NamedMemPool.use()``,
so the tag the freeze capped has to be stamped separately — the same pairing
the boot capture sweep uses (``engine/build.py``). Untagged capture bytes
would sit outside every cap.

Invoke the head and return ``(K, B)`` int64 draft tokens.

``frontier_repair_meta`` overrides the page-table-derived
frontier-repair meta (see :meth:`_repair_kv_frontier`). The
default derivation assumes the page-table length equals the
committed length at call time — true for every sync caller
(verify commit tail, seed-draft path). The async optimistic-
advance verify path calls ``draft`` BEFORE freeing rejected
draft slots, so it must pass its own device-gathered meta
(built in :func:`run_verify_step_async`, accept-aware with no
host sync).

``k`` is the number of speculative tokens to propose
(defaults to ``self.max_k``). For models whose bundled head
has ``n_draft == 1`` (Qwen 3.5 / 3.6 — all sizes ship a
single trained block) we **autoregressively chain** the same
block ``k`` times to produce ``k`` drafts:

  - draft_0: head(last_committed_token, last_committed_hidden)
             → (token_0, hidden_0)
  - draft_1: head(token_0, hidden_0)
             → (token_1, hidden_1)
  - …
  - draft_{k-1}: head(token_{k-2}, hidden_{k-2})
                 → (token_{k-1}, hidden_{k-1})

Each chained call's input is the previous draft's
(token, hidden) pair. The ``return_hidden=True`` path on
:class:`Qwen3_5MtpHead` exposes the post-norm hidden so this
chaining round-trips through the trained block exactly as
the next-step draft would.

Drafter accept rate compounds (drops geometrically) with k —
each chained draft is conditioned on the previous draft's
own output, so any drafter error compounds. Operators tune k
per-deployment; the engine's :func:`_mtp_verify_step`
consumes whatever drafts the driver returns.

For models with truly multi-block bundled heads
(DeepSeek V3 ships ``n_draft == 2``) the chain length above
equals ``head.n_draft``, and ``k > head.n_draft`` further
autoregresses the LAST block to reach the requested k. For
``head.n_draft >= k`` the multi-block path runs natively.

Sampling. When ``sampling_params`` is provided (one per
sequence in the batch) the driver by default STILL drafts
GREEDILY (the same argmax chain greedy requests use) and
returns ``(tokens, POINT_MASS_Q)`` — the proposal is a point
mass on each drafted (argmax) token, carried as the
:data:`~arbi_serve.spec_decode.drafter.POINT_MASS_Q` marker
and NEVER materialized as a dense ``(K, B, V)`` one-hot (that
densification plus its per-row mirror clones would add real
device-memory churn per spec step). The verify pass's
:class:`GraphSafeRejectionSampler` consumes the marker via its
index-form fast path — bit-identical to the dense one-hot —
and corrects the deterministic proposal back to the target
distribution: lossless for ANY proposal ``q`` (Leviathan-2023
Theorem 1; a point mass is a valid ``q``), so the served
output distribution equals the target. Drafting greedily for
stochastic requests lets them reuse the greedy-captured
drafter cudagraph — there is NO separate stochastic drafter
capture (a separate capture would need per-(B, K) full-vocab
sampling buffers at every ladder rung). ``seed_buf`` is not consumed by
the greedy draft (the rejection sampler at verify owns the
RNG). Default (``sampling_params=None``) returns just the
tokens, matching the legacy greedy contract.

TRUE-STOCHASTIC drafting (``ARBI_TRUE_STOCHASTIC_DRAFT``,
tri-state — ``0`` / ``1`` / ``auto`` with ``auto`` the default).
When the slate has at least one row that qualifies under the
resolved per-row temperature threshold (mode ``1``: any
``temperature > 0`` row; ``auto``: any row with ``temperature >=
ARBI_TRUE_STOCHASTIC_DRAFT_TEMP`` — see
:func:`stochastic_draft_temp_threshold`), each chain step
SAMPLES its draft token
from the drafter's own filtered distribution (temperature /
top_k / top_p / min_p — the same chain verify applies to
``p_target``) via one counter-keyed Gumbel-max draw per step,
and ``q_full`` is the REAL per-step ``(K, B, V)`` distribution.
Acceptance becomes ``sum_v min(p_v, q_v)`` instead of the
one-hot's ``p(argmax q)`` — temperature-flat instead of
decaying with entropy. Still lossless (Thm 1 holds for any ``q``); the
flag only moves the accept rate. The captured drafter chain is
greedy-argmax, so true-stochastic slates run the LIVE chain
(see :meth:`_draft_captured`); ``seed_buf`` (the rank-agreed
engine step seed) is REQUIRED — every TP/SPMD rank derives the
byte-identical draws from it, so drafts stay rank-symmetric
with no token broadcast. Greedy rows inside a mixed slate —
including ``auto``-mode rows BELOW the threshold, which are
re-encoded as greedy clones before routing
(:meth:`_effective_draft_params`) — collapse to their argmax
with a point-mass ``q`` (bit-equal tokens to the greedy
chain). The ``(K, B, V)`` ``q_full`` is
the same transient allocation scale the one-hot path already
pays — no VRAM regression.

Install the drafter's penalty-alignment state on the head; idempotent.

MUST run before the drafter chain is captured. The captured graph
bakes the addresses of whatever tensors its head forward touched at
record time, so a state installed afterwards would penalise the live
chain and leave every replay proposing from the un-penalised
distribution — the exact split this fix exists to close.

Returns whether the state is installed. Declining is always SAFE:
the drafter is a free parameter (Leviathan-2023 Thm 1), so a
penalty-blind proposal costs acceptance and nothing else.

Bring the drafter's penalty state up to the committed prefix.

Called once per :meth:`draft`, not once per chain step: the
committed term is the same for all K positions, and the per-position
difference (this step's own drafts) is applied on device inside the
chain where those tokens are realized.

True iff this slate drafts through the SAMPLED chain
(``ARBI_TRUE_STOCHASTIC_DRAFT``) instead of the greedy-argmax +
one-hot proposal: at least one row qualifies under the resolved
per-row temperature threshold (mode ``"1"`` = any ``temperature
> 0`` row; ``"auto"`` = any row at/above
``ARBI_TRUE_STOCHASTIC_DRAFT_TEMP``). A slate with NO qualifying
row — all-greedy, or all below the auto threshold — keeps the
one-hot path byte-identical to mode ``"0"``.

AUTO-mode per-row routing: re-encode below-threshold rows as greedy.

Only meaningful on a slate that IS routing stochastic
(:meth:`_true_stochastic_active` True). Rows whose ``temperature``
is ``> 0`` but BELOW the threshold draft greedy one-hot INSIDE the
sampled chain: they are substituted with a greedy clone
(``temperature=0.0, min_p=0.0``), which every drafter encoding
collapses to a point mass on the row argmax —
:meth:`DrafterSamplingTensors.encode_params_host` encodes
``temperature <= 0`` as ``temp=1.0, top_k=1, top_p=1.0`` for the
captured chain, and the live samplers short-circuit it to argmax +
one-hot ``q``. The drafted token and its one-hot ``q`` are exactly
the mode-``"0"`` values for that row, and the verify pass consumes
the SAME per-row ``q`` the drafter drew from (the tuple returned by
:meth:`draft` is threaded straight through), so losslessness holds
row-by-row. ``min_p`` is zeroed with the temperature because a
point-mass proposal has nothing for min_p to filter — this keeps a
below-threshold ``min_p > 0`` row from forcing the whole slate off
the captured chain (the request's real ``min_p`` still shapes
``p_target`` at verify).

Substitution happens ONCE at the top of :meth:`draft`, before the
captured/live routing and before the TP DrafterOp broadcast, so
every consumer (captured tensors encoding, live per-row sampler,
worker mirror) sees the same effective params. Idempotent: a
substituted row has ``temperature == 0`` and is left alone on a
second pass (the TP worker re-enters :meth:`draft` with the
broadcast params). Mode ``"1"`` (threshold ``0.0``) returns the
input unchanged — byte-identical to the pre-tri-state behaviour.

This captured chain's depth policy and replay form.

Returns ``(gate_run, segmented)``, or ``None`` when an armed gate
cannot ride this capture and the slate has to take the eager chain
instead. ONE resolver for both pools: the gate's signal is the
drafter's logit row, which the greedy chain and the sampled one
both form, so which pool a slate routes to must not decide whether
it can be gated.

The gate needs a host seam BETWEEN depths, which a fused K-step
graph does not have — it replays every depth it recorded in one
launch. A capture with recorded segments has that seam and the
margin buffer they fill; without both, the only way to serve the
gate is the eager chain.

Whether an armed gate rules out the pad-up graph for this slate.

The pad-up replays a WIDER capture and discards the padding rows'
drafts. The gate's slate reduction is a MAX across rows, so those
rows would decide the depth for the real ones — and a padding row
carries whatever the last replay left in its margin slot. Rather
than reduce over a subset (a second, quieter definition of the
signal), an armed slate that misses its exact shape takes the eager
chain, which gates at the shape it actually has.

Captured drafter-chain fast path; ``None`` ⇒ fall through to live.

The shipped MtpHead today is Qwen3_5MtpHead (single block,
n_draft == 1). When ``sampling_params`` is set, the greedy chain
returns ``(tokens, POINT_MASS_Q)`` — the point-mass proposal
marker, zero-cost — while the true-stochastic pool returns the
real stacked ``(K, B, V)`` per-step distribution for the
verify-pass recovery sampler.

Captured-graph fast path. Real-attention heads (Qwen3_5MtpHead)
capture cleanly: the captured pool's persistent buffer set threads
K per-step ``slot_mapping`` / ``seq_lens`` / ``cu_seqlens_k`` /
``positions`` slices that the live driver :meth:`copy_` ``s`` into
before each replay (see :class:`DrafterChainGraph`). The K draft
slots are still allocated by the live path so the host page-table
bookkeeping stays consistent across captured + live calls; only
the kernel-launch chain runs inside the graph.

TRUE-STOCHASTIC captured-chain fast path; ``None`` ⇒ live chain.

Replays the sampled-carry chain captured at boot into
``engine.drafter_graphs_stoch`` (see
:func:`~arbi_serve.runtime.capture.drafter.capture_drafter_chain`
with ``stochastic=True``): per-step counter-keyed Philox draws
(same ``OFFSET_DRAFTER + 17 * step`` lanes as the live chain, so
captured and live draw byte-matching noise for the same seed),
tensor-driven sampling params, and the per-step filtered ``q``
read back from the persistent ``q_chain`` buffer. Returns
``(tokens (K, B), q (K, B, V))``.

TP>1: the captured stochastic chain rides the SAME
rank-lockstep machinery as the greedy captured chain — rank 0's
replay broadcasts a captured stochastic ``DrafterOp``
(``replay_captured=True`` + the slate's per-row sampling params) so
every worker replays its OWN captured stochastic graph. The graph
bakes the identical per-step o_proj all_reduce + lm_head all_gather
on every rank (captured under the TP-aware ``graph_capture``), so the
collective count rendezvous 1-to-1 — no live fallback, no ~1ms/step
tax. The counter-keyed draws are byte-identical across ranks from the
rank-agreed ``seed_buf`` (no token exchange), exactly as the live TP
chain already proved. See :meth:`_replay_captured_chain` +
:func:`~arbi_serve.distributed.worker_bridge._worker_replay_captured_chain`.

Falls through to the live chain (``None`` — correct at any
shape, no refusal) when:

  - the slate has any ``min_p > 0`` row AND the fused draw is off
    (the unfused captured chain doesn't apply min_p) — symmetric
    under TP (every rank reads the same env);
  - no engine / no seed / non-real-attention stub head;
  - the pool has no exact ``(B, k)`` graph (miss counted; warned
    once per shape — no pad-up in v1). Under TP the miss is
    symmetric (every rank captures the same ladder), so rank 0's
    live fallback broadcast keeps the worker in lockstep.

Shared-KV drafter chain for the Gemma-4 EAGLE3 assistant.

The head's N draft layers are Q-only and READ the verifier's KV
(their marker specs alias the verifier's donor layers via the pool;
the attn ops carry ``skip_kv_write``). No draft slots are allocated
and the query position is HELD at the frontier across all K chain
steps (``advances_positions is False``); chaining is purely through
the ``post_projection`` hidden feedback. A bare read meta (one query
per row, ``seq_len = committed_len``) is built per attention-window
variant and decorated with the backend per-step page metadata.

Sampling. Greedy default (``sampling_params is None`` or a slate
with no true-stochastic row): drafts GREEDILY through the captured
held-position chain (or its eager fallback) and returns
``(tokens, POINT_MASS_Q)`` for a stochastic slate — the verify
pass's rejection sampler corrects the point mass losslessly.
TRUE-STOCHASTIC (:meth:`_true_stochastic_active`,
``ARBI_TRUE_STOCHASTIC_DRAFT``): each chain step SAMPLES from the
head's own centroid-masked filtered distribution and the driver
returns the REAL stacked ``(K, B, V)`` ``q_full`` — byte-shared
wiring with the Qwen own-slab :meth:`_draft_live` stochastic path
(same head contract, same shared sampler). The captured greedy
chain bakes the argmax carry, so a true-stochastic slate runs the
LIVE (eager) chain — exactly as the Qwen path falls back when no
stochastic drafter graph was captured; ``seed_buf`` (the
rank-agreed engine step seed) is REQUIRED so every TP rank draws
byte-identical carry tokens with no token exchange.

Build the per-window ``kv_ctx`` assembler for the shared-KV chain.

The head's draft layers are Q-only and READ the verifier's KV, so the
context is ``(layer view, meta, rope cache, attn op)`` per marker spec
and depends on the engine and the head, never on the live batch. Shared
by the live chain and the boot pre-capture so both assemble the chain
from identical parts.

Record + warn ONCE per shape that this shape serves eager for good.

A shared-KV shape with no captured graph re-runs the K-step chain eager
on every spec step. That is correct but costs throughput for the life of
the process, so it is a named warning and a must-not-fire counter, never
an absorbed fall-through.

The shared-KV draft chains this driver owns, for the sleep walk.

The engine's walk introspects dataclass entries of its cudagraph
pools; these chains live in a plain dict of plain dicts, so the
ownership is declared here (see
``arbi_serve.engine.sleep._GRAPH_OWNER_ATTRS``). Empty on a head with
no shared-KV chain.

Capture width to replay for ``(B, k)``, or ``None`` to run eager.

``B`` itself when that shape is captured; otherwise the smallest
captured width above it, whose extra rows the caller fills with a real
row and discards — the same pad-up the decode and Qwen-drafter lookups
do, and what keeps an off-rung batch off the eager chain. Before the
freeze seals the surface an uncaptured shape is captured on demand, so
the width is ``B``; once a capture has been denied it is not retried,
because the layout it did not fit does not grow back.

Close the shared-KV capture surface at the Phase-2 freeze.

After the freeze every cuMem pool is capped at its live mapped bytes, so
a capture has no budgeted home. Sealing turns an uncaptured shape into a
pad-up (or, failing that, a warned eager fall) instead of a capture that
the cap denies mid-request.

The ``(B, K)`` ladder the boot sweep pre-captures. Empty when unused.

The shared-KV chain's only varying dimensions are the row count and the
draft depth: the block-table width is pinned at the pool page ceiling,
the attention-window set comes from the head's marker specs, and the
held query position is a per-replay input. So the ladder is exactly the
drafter-chain ``(B, K)`` cross-product the Qwen own-slab sweep uses.

Capture the shared-KV ``(B, K)`` ladder. Returns the number captured.

Called by the boot capture sweep, inside the sweep's capture tag, so the
chains land in the pools the freeze then caps. Without it the first live
request at each shape captures against a frozen layout, is denied, and
serves eager for the life of the process.

Uses the production replay path with synthetic single-cell rows: the
captured graph reads its persistent buffers, which every replay refills
from the live batch, so only the SHAPE has to match here. A bucket that
fails to capture is warned and counted, and the remaining buckets are
still attempted.

CUDA-graph-captured shared-KV draft chain (SOTA path — never eager).

Persistent per-``(B, K)`` buffers hold the chain inputs (last token,
backbone hidden) and the constant read meta (block table, seq_lens,
positions, frontier slots) — CONSTANT across the K held-position
steps, unlike the Qwen own-slab chain. The tkv-bypass CSR page
metadata lives in the backend builder's own persistent buffers
(baked by ``data_ptr`` at capture); re-decorating before each replay
refills them for the live request. First call per shape warms up +
records; subsequent calls ``copy_`` fresh inputs, re-decorate, and
replay.

Every allocation a capture makes — persistent buffers, the warmup
transient, the graph-private working set — routes into the engine's
named capture pools (``capture.io_buffers`` / ``capture.cudagraphs``),
so the bytes sit inside the tags the Phase-2 freeze caps. The own-slab
drafter sweep declines shared-KV heads, so the ladder is captured here
instead — by :meth:`precapture_shared_kv_chains` during the boot sweep,
before the freeze. A capture attempted after it is denied at map time
and surfaces as :class:`torch.OutOfMemoryError` for the caller to warn
on and serve eagerly.

Live autoregressive drafter chain (captured-graph miss / no pool).

For real-attention heads we pre-allocate K draft slots per request,
run the head K times (each step attends over the request's full KV
history + prior chain steps and writes ONE new K/V into the next
slot), then free the K slots so the next step's verify builder
allocates fresh tail+draft slots from the same post-commit length.
Pool exhaustion bubbles up — the caller (run_step / verify driver)
wraps :meth:`draft` in a try/except that degrades to the cold
next-step path; there is no half-wired no-history fallback.

True-stochastic slates (:meth:`_true_stochastic_active`) sample
each step's carry token from the drafter's own filtered
distribution via the head's counter-keyed sampling path
(``slot_offset=step`` decorrelates the K draws from one shared
``seed_buf``) and return the REAL stacked ``(K, B, V)`` ``q``;
everything else — slot allocation, per-step metas, frontier
repair, slot free — is byte-shared with the greedy chain.

Captured-chain replay + per-step draft-meta machinery for the MTP driver.

:class:`_MtpDraftChainMixin` holds the cohesive cluster of
:class:`~arbi_serve.spec_decode.mtp_driver.MtpDriver` methods that build
per-chain-step attention metas, allocate/free draft slots, run the frontier
MTP-KV repair, and replay the captured (exact + padded) drafter chains. The
methods reference
``self.<attr>`` bound in :meth:`MtpDriver.__init__` (duck-typed via the mixin).

Host block-table rows for ``requests``, padded to the batch max page count.

Returns ``(max_pages, rows)`` where ``max_pages`` is the largest
per-request page count in the batch (``>= 1``) and ``rows[i]`` is
request ``i``'s page-table row padded to ``max_pages``. Shared by
the draft-meta builder, the TP broadcast, and the captured-chain
replays so all four read the table identically.

Free each request's K per-step draft slots (rollback, ``k_accepted=0``).

Shared ``finally``-block helper for the live + captured-replay
chains: the next step's verify pass allocates fresh tail+draft
slots from the post-commit length, so the chain's slots must be
returned first.

Build per-chain-step :class:`AttnPagedKVMeta` + slot lists.

For each request we allocate ``k`` draft slots tail-extending
its current length, then construct ``k`` per-step metas on the
fill / vLLM position convention
(:func:`shift_chain_to_fill_convention`). At chain step ``s``
(0..K-1):

  - ``cu_seqlens_q = [0, 1, 2, ..., B]`` — one query per row.
  - ``seq_lens[i] = pre_len + s`` — covers committed history
    plus this step's own K/V write.
  - ``cu_seqlens_k = [0, sum(seq_lens)]`` cumulative.
  - ``slot_mapping[i]`` — step 0 rewrites the FRONTIER cell
    (sequence index ``pre_len - 1``); steps 1..K-1 use the
    first ``k - 1`` allocated draft slots.
  - ``positions[i] = pre_len - 1 + s`` — the pair sits at its
    HIDDEN's RoPE position (the pairing the head was trained
    on; matches the verify fill pass and vLLM's proposer).
  - ``block_table`` = per-request page table rows.

Returns ``(metas, per_req_slots, pre_lens, repair_meta)`` where
``metas`` is the list of per-step :class:`AttnPagedKVMeta`,
``per_req_slots`` is the list of per-request K-slot lists in
request order, ``pre_lens`` is the host pre-chain length per
request (used to compute the per-step TKV ``bypass_safe`` bool
host-side, no D2H), and ``repair_meta`` is the one-query-per-row
frontier-repair meta (:func:`build_frontier_repair_meta`)
targeting each request's last committed MTP-KV cell — derived
from the SAME ``pre_lens`` + block-table rows, so it is only
valid when the page-table length equals the committed length at
call time (true on every sync caller; the async optimistic path
passes its own device-gathered override instead). The caller
frees the slots after the chain.

True iff chain step 0's own K/V write subsumes the eager repair.

On the SYNC drafter paths (verify commit tail, seed-draft path)
the page-table length equals the committed length, so
:func:`shift_chain_to_fill_convention` places chain step 0 INTO
the frontier cell (slot of sequence index ``pre_len - 1``) at the
cell's own RoPE position, computing ``fc(cat(embed(token),
hidden))`` + ``layers[0]`` from the SAME ``(last_token_ids,
prev_hidden)`` inputs the eager :meth:`_repair_kv_frontier` uses —
a byte-identical K/V write to the identical slot. The attn op
scatters the new K/V into the cache BEFORE attending
(write-then-read), so step 0 even reads its own corrected cell;
the eager prologue is a pure idempotent duplicate and is skipped.

The skip is keyed ONLY on ``frontier_repair_meta is None`` — the
sync/async discriminator. On the ASYNC optimistic-advance path the
caller passes a device-gathered override targeting the REAL
last-accepted cell while step 0 (built from the optimistic, not-
yet-freed page-table length) targets the OPTIMISTIC span end — a
DIFFERENT cell — so the eager repair is still required there and
must NOT be skipped.

Plan-driven + rank-identical at TP>1: on the async override path
rank 0 ships the device-gathered frontier slots / positions /
seq-lens in the ``DrafterOp`` (``_broadcast_drafter_chain`` →
``bridge.broadcast_drafter(frontier_repair_meta=...)``) and the TP
worker (``_run_worker_drafter``) reconstructs a byte-identical
override meta and runs the SAME eager ``_repair_kv_frontier``
BEFORE its chain — so the extra repair collectives rendezvous
1-to-1 and the per-step count stays in lockstep. Every sync caller
passes ``frontier_repair_meta=None`` on all ranks (step 0 subsumes
the repair), so both ranks take the same branch.

DIAGNOSTIC: under ``ARBI_MTP_FRONTIER_EAGER=1`` the
sync-path subsumption is DISABLED — the eager repair always runs
(chain step 0 is simultaneously redirected off the frontier by
:func:`redirect_step0_off_frontier`, and the TP worker mirrors
both changes from the same env, keeping the repair collectives
paired 1-to-1).

Rewrite each request's frontier MTP-KV cell from the corrected pair.

The verify/decode fill pass leaves each request's LAST committed
cell paired with the wrong embedding (self on full accepts, the
rejected draft on partial accepts — see
:meth:`Qwen3_5MtpHead.refill_kv_cell`). This runs BEFORE the
chain's first step so step 0 already attends over a fully
repaired slab. ``last_token_ids`` / ``prev_hidden`` are the
chain's own step-0 inputs — the committed bonus/recovery token
and the main model's hidden at the last-accepted cell — which
are exactly the corrected ``(embed, hidden)`` pair the fill pass
lacked; the repair adds no new data dependencies.

TP>1: the worker mirror (``_run_worker_drafter``) performs the
identical repair from the broadcast payload (slot derived from
the shipped ``pre_lens`` + block-table rows via the same pure
helper), so the per-repair collectives (vocab-parallel embed
all_reduce + o_proj / MLP all_reduces) rendezvous 1-to-1 with
rank 0's. CUDA graphs: the repair always runs EAGER — on the
captured-replay paths it executes between graph replays, never
inside a captured region. ``committed_lens`` (host) feeds the
TKV ``bypass_safe`` predicate; ``None`` (the async device-
gathered override, whose frontier index is data-dependent)
conservatively disables the bypass, which is always correct.

Send a ``("drafter", DrafterOp)`` to worker ranks.

Used by both the live chain AND the captured-replay branches.
Rank 0 issues K cross-rank collectives per chain step (one
``all_reduce`` per layer ``o_proj`` + one ``all_gather`` per
``ColumnParallelLinear`` lm_head). The worker rank must fire a
matching peer collective for each — in the captured-replay case it
replays its OWN captured chain (``replay_captured=True``); otherwise
it runs the live chain — so the per-step collective COUNT matches
peer-for-peer, else the TP forward desyncs. (The control plane rides
a CPU shm queue, not NCCL, so a miscounted collective cannot mis-pair
with a control broadcast — but the count must still match.)

``replay_captured`` controls HOW rank 1 supplies the peer:

  - True (rank 0 took the captured-graph replay): rank 1 REPLAYS
    its own captured chain for the same ``(B, K)`` bucket — both
    graphs bake the identical K all_reduce + False (rank 0 ran the live chain, or a worker capture miss):
    rank 1 runs the live eager chain, which fires the same K
    collectives step-by-step.

Eager frontier MTP-KV repair before a captured-chain replay.

Shared by the exact + padded captured-replay paths. On the sync
path (``frontier_repair_meta is None``) the captured chain's step 0
rewrites the frontier cell in place, subsuming the repair, so this
is a no-op; the async override path supplies a device-gathered meta
(and ``None`` committed lens) and runs the eager repair. See
:meth:`_repair_kv_frontier` and :meth:`_chain_subsumes_frontier_repair`.

Captured-chain replay for real-attention heads.

Allocates K per-step draft slots via
:meth:`_build_draft_step_metas`, packs the live per-step
:class:`AttnPagedKVMeta`s + the shared block_table into
:meth:`DrafterChainGraph.replay`, returns the cloned ``(K, B)``
tokens. The K slots are freed after replay regardless of
outcome — the verify pass's tail+draft allocator runs from
the post-commit length and must see the pre-chain state.

TP>1: ``bridge`` carries the rank-0 :class:`RankZeroBridge` so
we can broadcast a ``("drafter", DrafterOp)`` BEFORE
``captured.replay()`` fires its captured K all_reduces +
K all_gathers. Worker ranks then run the live chain with the
same per-row slot allocation, matching rank 0's collective
sequence on the TP group 1-to-1. See
:meth:`_broadcast_drafter_chain` for the rationale.

``segmented`` replays the capture's per-step graphs at full depth —
the same kernels, K launches instead of one — which is what prices
the seam a gated walk needs, separately from the cut it makes.

``gate_run`` (a per-slot depth policy) routes the replay onto the
capture's per-step graphs and is consulted in the seam between them,
so the returned block can be SHORTER than ``k`` by decision. It
rides the GREEDY capture and the stochastic one on the same terms —
the margin it reads comes off the drafter's logit row, which both
chains form to make their pick. The K
per-step draft slots are still allocated and freed as a block: the
verify pass's tail+draft allocator runs from the post-commit length
and must see the pre-chain state whatever depth the walk reached.

``sampling_params`` + ``seed_buf`` (TRUE-STOCHASTIC captures
only): threaded to :meth:`DrafterChainGraph.replay`, which
encodes them into the persistent param buffers before the
replay; the return becomes ``(tokens, q)``. The slot / meta /
block-table / repair / free machinery — including the
kv_pages_bucket length-safety of the per-step metas built by
:meth:`_build_draft_step_metas` — is byte-shared with the
greedy replay.

Read the live ``ARBI_DECODE_PAD_CUDAGRAPH`` gate.

Reads the runner's runtime-toggleable
``_decode_pad_cudagraph`` attribute (same flag the decode /
verify pad-up paths gate on, flippable via
``/v1/admin/decode_pad_cudagraph/{state}``) so a same-process
ON/OFF bit-exactness harness drives the drafter pad-up in
lockstep with the verify pad-up. Defaults OFF when the runner
or attribute is absent.

Tally a drafter-chain captured-graph ``hit``/``miss``/``pad`` at ``(B, k)``.

Gated on ``ARBI_DEBUG_CAPTURE_LOOKUP`` (set at construction). The
histogram is surfaced via ``/v1/admin/capture_hist`` so the
hetero probe can confirm the bundled drafter chain now HITS (or
pad-HITS) at off-ladder B instead of running eager.

Refuse LOUD under TP>1+cudagraph; otherwise warn once.

Under ``world_size > 1`` with cudagraphs enabled the live EAGER
drafter chain does NOT just run slow — its per-step collectives
DEADLOCK the NCCL group (the captured verify/decode path replays
graphs with no host-side collective, so a rank that drops to the
eager chain desyncs the others → 0 tokens, GPUs pinned). This is
the c=1 27B TP2 MTP cudagraph deadlock. A captured-graph miss
here is therefore a hard error, never a silent fall-through to a
deadlocking path. Single-GPU (or eager-everywhere) keeps the
legacy slow-but-correct eager chain via :meth:`_warn_eager_miss`.

Warn ONCE per ``(B, k)`` that the drafter chain runs eager.

Every spec step at this shape re-runs the K-step live chain eager
instead of replaying a captured graph. The capture sweep
(:func:`arbi_serve.engine.cudagraph_admin._drafter_chain_buckets`)
covers K=1..n_draft, so a steady miss means an out-of-ladder shape
or a trimmed sweep — surface it loudly either way.

Smallest captured drafter ``B' > B`` at the same ``k``, or ``None``.

Enumerates the drafter pool's ``(B', K')`` keys for the smallest
``B' > B`` with ``K' == k``. The caller has already bounded the
live batch by ``cfg.batch.max_batch`` (the scheduler never admits
more than ``max_batch`` rows), so every captured drafter shape is
``<= max_batch`` by construction and no explicit ceiling is
needed here. Returns ``None`` when no larger same-``k`` shape was
captured (caller falls through to the eager chain).

Pad-up greedy captured-chain replay for real-attention heads.

Mirrors :meth:`_replay_captured_chain` but for a live ``B``-row
batch that missed the exact ``(B, k)`` drafter graph: pad UP to a
captured ``(b_prime, k)`` graph (``b_prime > B``), replay, and
slice the returned ``(K, b_prime)`` draft tokens back to
``(K, B)``.

Correctness contract (identical reasoning to the decode / verify
pad-up in :mod:`arbi_serve.runtime.decode_pad`):

  * The drafter chain is per-row independent — each row's K-step
    autoregression attends only within that row's own KV history
    (per-row ``cu_seqlens_q`` / ``seq_lens``), and the head's
    embed / RMSNorm / per-row sampling are pointwise. A padding
    row therefore cannot change what a real row ``i < B``
    drafts. The only coupling is the GEMM batch dim ``M``
    (``B ⇒ b_prime``) in the head's ``Linear`` projections, whose
    cuBLAS split-K reduction order shifts the last bf16 bits —
    intrinsic nondeterminism, not corruption. Gated behind
    ``ARBI_DECODE_PAD_CUDAGRAPH``.
  * Real rows ``[0:B]`` keep their REAL allocated draft slots +
    block-table rows (built by :meth:`_build_draft_step_metas`),
    so their KV writes + attention reads are byte-faithful to the
    exact replay.
  * Padding rows ``[B:b_prime]`` carry a null page-0 block-table
    row, ``slot_mapping = 0`` (the reserved null KV slot — same
    sentinel the decode pad-up scatters into), ``seq_lens = 1``,
    ``positions = 0``. Their per-row attention reads a single
    null token and their K/V scatters into the null slot; the
    captured graph processes them for shape only and their draft
    tokens are sliced off before return.

No draft slots are allocated for the padding rows — they reuse
the null slot 0 — so the post-chain ``free_draft_slots`` is over
the real rows only (matching :meth:`_replay_captured_chain`).

Extend each per-step :class:`AttnPagedKVMeta` to ``b_prime`` rows.

Real rows ``[0:B]`` keep their values; padding rows ``[B:b_prime]``
get null-slot scatter (``slot_mapping = 0``), single null token
(``seq_lens = 1``), ``positions = 0``, all-zero block-table row.
``cu_seqlens_q`` becomes ``arange(b_prime + 1)`` (one query/row),
``cu_seqlens_k`` the cumsum over the padded ``seq_lens``.

Extend the per-request K-slot lists with ``b_prime - B`` null rows.

Used only on the TP>1 broadcast path so the worker reconstructs a
``b_prime``-wide chain matching rank 0's padded replay. Padding
rows scatter into the null slot 0 each step (``[0] * k``).

The drafter chain, expanded into a tree.

A chain runs the head ``K`` times, one row per request, feeding each
step's argmax token and output hidden into the next. A tree runs it once
per DEPTH, on one row per previous-depth node THAT EXPANDS, and takes
the head's top-``W`` at each of them.

The recursion's shape, which is the part that is easy to get subtly
wrong: expanding node ``i`` runs the head on ``(t_i, h_parent(i))`` and
yields its children, which share ONE output hidden — children differ by
TOKEN, not by state. So depth ``d + 1``'s hiddens are depth ``d``'s
output hiddens indexed by which ROW each next parent came from, and its
tokens are depth ``d``'s candidate slate indexed by which CANDIDATE each
next parent IS. Those are one and the same permutation for a product
tree, where every candidate expands and every row spawns the same
number; a spine tree's truncated tails break that, which is why the two
gathers are carried separately (:attr:`TreeSpec.draft_carry`).

One ``topk`` runs per depth, at the widest width any row there takes,
and :attr:`TreeSpec.draft_take` cuts the slate down to that depth's
nodes. A per-row width would be a data-dependent kernel shape; the
gather is not, and it costs nothing a captured graph objects to.

## The attention approximation, stated plainly

Correctness does not depend on anything in this module. The drafter
PROPOSES; the verify pass decides, and it decides against the target
model. A worse proposal costs acceptance, never correctness — which is
why byte-identity against non-spec greedy decode holds regardless of
what the drafter does here.

That matters because the drafter cannot give each node its own attention
prefix. Paged attention bounds a row's keys by a contiguous ``seq_len``,
so two siblings — alternatives for ONE position — cannot both own that
position's KV cell. The verify block solves this with an ancestors-only
mask; the drafter runs one query per row and has no block to mask.

So every node at depth ``d`` attends over the SPINE: the committed
prefix plus the top-1 branch's cells at depths ``0..d-1``. Its own token
and its parent's hidden are exact; its ancestors' KV is the spine's.
Off-spine nodes are therefore drafted from a slightly wrong context and
should propose slightly worse tokens. Non-spine rows write their KV to
SINK slots past every row's ``seq_len``, so nothing races the spine's
cell and nothing reads what they wrote.

``mtp_tree_draft_spine`` counts the approximation rather than leaving it
in a comment: a boot can tell how many node expansions took it.

Per-depth row layout for the tree drafter's attention metas.

Pure host arithmetic, separated from the meta build so it can be
asserted without a device — the slot schedule is where a tree drafter
goes wrong silently, and it is all integers.

One entry per depth ``d`` with:

  ``rows``      — parents that EXPAND at this depth (1 at depth 0)
  ``pre_lens``  — the fill convention's ``pre_len - 1``, per BATCH row
  ``slots``     — flat ``B * rows`` slot mapping, request-major
  ``seq_len_of``— ``pre_lens[b] + d``, per batch row

Slot schedule, in the same fill convention the chain uses
(:func:`~arbi_serve.spec_decode.mtp_draft_meta.shift_chain_to_fill_convention`):

  * the SPINE row (parent 0) of depth 0 writes the FRONTIER cell,
    and of depth ``d`` writes draft-slot ordinal ``d - 1`` — exactly
    the cells chain step ``d`` writes, so a width-1 tree reproduces
    the chain's writes byte-for-byte;
  * every other row writes a SINK, ordinal ``depth - 1`` and up.
    Those live at logical positions ``>= L + depth - 1``, which is
    past the widest ``seq_len`` any depth uses (``L + depth - 1``
    covers logical ``0 .. L + depth - 2``), so a sink is never read.
    Writing them at all — rather than pointing every off-spine row at
    one scratch cell — is what keeps the rows from racing each other
    inside a single fused scatter.

``[depth][parent_row] -> node ids of that row's root-to-self path``.

Depth 0 has no history (its parent is the committed token, which is
not a drafted node), so entry 0 is empty. Depth ``d``'s rows are
:attr:`TreeSpec.draft_rows` — the depth-``d-1`` nodes that expand, in
node order, matching the head's row order.

Every static index one geometry's expansion gathers through.

Built ONCE per geometry, which is a process-global flag: building
these per depth per step instead put a PAGEABLE host-to-device copy on
the live path, and a pageable copy is a host stall — the CUDA runtime
serializes it behind everything already on the stream, which on this
path is the previous step's verify forward.

Per-depth flat gather index into the draft tensor's node rows.

Entry ``d`` is :func:`_ancestor_paths` at depth ``d``, flattened to the
``index_select`` the drafter runs — so the live path performs a gather
and no transfer.

Draft one tree per request. Returns ``(num_nodes, B)`` int64.

Row ``i`` of the result is node ``i``'s proposed token, in the
breadth-first order every other part of the tree machinery uses —
the verify plan's draft slots, the ancestor mask, the accept
walk's node ids and the KV compaction's ordinals all index by it,
so a different order here is wrong everywhere at once and
raises nowhere.

The proposal is a POINT MASS per node at every temperature: the
picks are the head's top-W, taken deterministically. That is the
premise the multi-candidate accept rests on — it inverse-CDFs the
TARGET's mass over the candidate set, which is lossless for
deterministic candidates and is not the rule for sampled ones.

Allocate the tree's draft slots and build one meta per depth.

The returned frontier-repair meta is ``None`` under
``want_repair_meta=False`` — the caller's depth-0 write subsumes the
repair, and building a meta nothing reads is four host-to-device
copies a step.

Refuses a slate whose widest depth would exceed the batch the
backend's persistent page-metadata buffers were sized for: a tree
depth runs ``B * parents`` attention rows where a chain step runs
``B``, and the CSR triplet those rows index is allocated once at
boot for ``max_batch`` rows. Overrunning it writes past a
persistent buffer, which is not a shape error anywhere — it is a
wrong page index handed to the kernel later.

Qwen 3.5 / 3.6 MTP draft head — single-block transformer drafter.

The HF Qwen 3.5 / 3.6 release ships a 1-block MTP head bundled inside
the main safetensors under ``mtp.*``:

    mtp.pre_fc_norm_embedding.weight    # RMSNorm on embed(token)
    mtp.pre_fc_norm_hidden.weight       # RMSNorm on prev hidden
    mtp.fc.weight                       # (H, 2H) — concat([embed, hidden]) -> H
    mtp.layers.0.input_layernorm.weight
    mtp.layers.0.post_attention_layernorm.weight
    mtp.layers.0.self_attn.{q,k,v,o}_proj.weight
    mtp.layers.0.self_attn.{q,k}_norm.weight
    mtp.layers.0.mlp.*                  # FFN — arch-specific shape
    mtp.norm.weight                     # final norm before lm_head

The FFN shape follows the arch's own decoder layers: a dense gated MLP
on ``Qwen3_5ForConditionalGeneration`` (``mlp.{gate,up,down}_proj``), a
sparse MoE block on ``Qwen3_5MoeForConditionalGeneration``
(``mlp.gate`` router + ``mlp.experts.{E}.*`` + ``mlp.shared_expert*``).
The arch constructs that module and hands it to :class:`Qwen3_5MtpHead`
as ``mlp=``, together with the ``mlp_weight_map=`` callable that binds
its keys — everything else about the head is shape-identical across the
two.

Forward path (per draft step, K=1):

    embeds = embed_tokens(committed_token_id)           # (B, H)
    e = pre_fc_norm_embedding(embeds)
    h = pre_fc_norm_hidden(prev_hidden)
    fused = fc(cat([e, h], dim=-1))                     # (B, H)

    # one Qwen-3.5-style decoder block (pre-norm + attn + MLP), with
    # REAL paged attention over the MTP layer's per-request KV cache:
    x = fused + self_attn(input_layernorm(fused), positions, layer_view, attn_meta, attn_op)
    x = x + mlp(post_attention_layernorm(x))

    out = norm(x)                                       # (B, H)
    draft_logits = lm_head(out)                          # (B, V)
    draft_token = argmax(draft_logits)                   # (B,)

The MTP layer's KV cache is sized as one extra entry on the engine's
per-layer slab (see ``Qwen3_5Model.from_safetensors`` — appends a
marker :class:`LayerSpec` with ``is_mtp_layer=True`` whose pool /
attn-op infrastructure mirrors the main attention layers). The slab
is populated by :meth:`Qwen3_5Model.forward` running the MTP layer
over the SAME flat batch the main model just saw (vLLM-equivalent
"draft prefill" — same attn_meta, same slot_mapping, same positions,
DIFFERENT layer_view + attn_op for the MTP slab). Each draft step
then runs the head over a single new position whose K/V scatters
into the next free slot in the MTP slab; attention reads the entire
history.

This shape ships with K=1 support (``mtp_num_hidden_layers == 1``).
:class:`MtpDriver.draft` chains the head autoregressively for K>1,
building one fresh per-step ``MtpDraftStepMeta`` (one query token,
attn over the MTP slab's running history) per chain step.

``enable_if`` predicate gating the MTP draft-head compile.

Evaluated at head CONSTRUCTION time (inside the engine's
:func:`~arbi_serve.compile.set_compile_context` window) — reads
``ARBI_MTP_HEAD_COMPILE`` via the runtime-flags registry. Default
OFF: when False the ``@support_torch_compile`` decorators below are
inert (the decorator's ``__init__`` wrapper returns before any
trampoline / marker install), so the constructed head is
structurally identical to the eager head — same bound
``forward``, no ``_arbi_compile_state``, and the drafter-chain
capture's ``compile_capture_ctx(head)`` gate stays closed.

Why default OFF: the compiled head's Inductor reduction-order drift
can flip a draft argmax, and a K>1 autoregressive draft chain
amplifies one early flip into losing the whole chain suffix — a net
tok/s regression despite the faster per-step head. Small models can
opt in via ``ARBI_MTP_HEAD_COMPILE=1``. See
``RuntimeFlags.mtp_head_compile``.

Single MTP decoder layer — pre-norm + paged attn + pre-norm + MLP.

Signature matches :class:`arbi_serve.models.qwen3_5._Qwen3_5AttentionDecoderLayer`
so the same call sites (main-model forward MTP-fill pass; per-step
draft chain in :class:`MtpDriver`) can invoke it uniformly.

Decorated ``@support_torch_compile(level="block")``
so — exactly like the main model's ``_Qwen3_5AttentionDecoderLayer`` —
the layer's ``forward`` is routed through the same Inductor trampoline
when constructed inside :func:`~arbi_serve.compile.set_compile_context`
(the head is built from ``Qwen3_5Model.__init__``, which the engine
wraps in that context). Without this trampoline (or when ``compile_on``
is False) the MTP draft head runs entirely eager: at c=1-with-MTP the
driver's live autoregressive chain (:meth:`MtpDriver._draft_live`) calls
the head once per speculative step, each firing unfused eager RoPE
``torch.cat`` (5 ``CatArray``/step), eager RMSNorm reduce kernels, and
per-step bf16 casts — the same eager-glue gap the main forward's own
RoPE/RMSNorm kernels avoid (surfaced as ``triton_poi_fused_*``). Routing
the layer through the trampoline fuses them into Triton kernels
for BOTH the live chain (c=1) and the captured drafter chain
(:class:`~arbi_serve.runtime.capture.drafter.DrafterChainGraph`, whose
warmup passes JIT the compiled callable before the ``torch.cuda.graph``
record so the fused kernels bake into the captured graph too).

The attention op stays an opaque custom-op boundary inside the compiled
region (Inductor treats ``arbi_serve``'s kernels as black boxes), so the
whole ``forward`` — including the ``self_attn`` call — compiles as one
graph, matching the main attention layer's non-split-attn surface.
The decorator is a strict no-op when ``compile_on`` is False (the layer
stays eager, bit-identical to the eager path).

Gated by ``ARBI_MTP_HEAD_COMPILE`` (default OFF) via the
``enable_if`` predicate: the compiled head's Inductor reduction-order
drift regresses accept_len on K>1 chains (see
:func:`_mtp_head_compile_enabled`), so the trampoline is opt-in.

Qwen 3.5 / 3.6 single-block MTP draft head.

An instance of :class:`MtpHeadBase`: a full-rank head that reuses the
main model's embed/lm_head, writes its own draft KV, and advances
positions per draft token. The composable-hook methods
(:meth:`embed` / :meth:`mix` / :meth:`run_stack` / :meth:`to_feedback`
/ :meth:`to_logits` / :meth:`to_token`) express the SAME math as the
hand-tuned :meth:`forward` and give the head-agnostic driver one
surface across both families; the production draft path still runs the
inlined :meth:`forward` (unchanged) for its fused sampling / captured-
chain machinery.

Decorated ``@support_torch_compile(level="model")``
which — because the head owns a block-decorated descendant
(:class:`_Qwen3_5MTPLayer`) — installs a *marker* trampoline (eager
passthrough; the per-layer decorator does the real Inductor work),
mirroring how :class:`~arbi_serve.models.qwen3_5.Qwen3_5Model` relates
to its decoder layers. The marker makes the head report
:func:`~arbi_serve.compile.capture_bridge.is_compile_eligible` True so
the drafter-chain capture scaffold's ``compile_capture_ctx(head)`` gate
opens exactly as it was designed to (the bridge is otherwise a no-op on
an undecorated head). ``head.forward`` behaviour is unchanged — the
marker trampoline just calls the original body. No-op when
``compile_on`` is False.

Gated by ``ARBI_MTP_HEAD_COMPILE`` (default OFF) via ``enable_if``
— see :func:`_mtp_head_compile_enabled`. With the flag off neither
the marker here nor the block trampoline on
:class:`_Qwen3_5MTPLayer` installs, so the head is the plain
eager head and the drafter-capture ``compile_capture_ctx`` gate
stays closed.

Reuses the main model's ``embed_tokens`` and ``lm_head`` (held
behind ``_tied_modules`` so they don't get registered twice in
the parameter graph; the loader binds them on the main model).

The head's own weights — ``pre_fc_norm_embedding``,
``pre_fc_norm_hidden``, ``fc``, the single decoder block, and
the final ``norm`` — are loaded from the safetensors' ``mtp.*``
keys via :meth:`weight_map`.

Build the MTP decoder layer: pre-norms, gated attention block, and
the caller-supplied FFN MODULE.

``mlp`` is a module, not a size: the arch supplies whatever FFN its
bundled MTP block carries — a dense :class:`GatedSiLUMLP`, a sparse
MoE block, anything with a ``(N, H) -> (N, H)`` forward. A
size-typed parameter here could only ever build the dense MLP,
which is exactly the assumption that locks MoE arches out of MTP.

Build the single-block Qwen 3.5/3.6 MTP head (pre-FC norms, fc
projection, decoder layer, final norm), holding tied references to
the main embed/lm_head.

``mlp`` is the decoder block's FFN as a MODULE — the arch supplies
it (dense :class:`GatedSiLUMLP`, sparse MoE block, …) so the head
can represent any bundled MTP block's FFN, and MTP composes with
the arch instead of assuming a dense MLP.

``mlp_weight_map`` supplies the FFN's weight-map entries as a
callable ``(dst_prefix, src_prefix) -> {param: WeightShardSpec}``
(prefixes name the mlp module path on the head and in the
checkpoint respectively). ``None`` is valid ONLY for a
:class:`GatedSiLUMLP` (the standard ``{gate,up,down}_proj``
binding is built in); any other FFN must bring its map — refused
at construction, not at load time.

The lm_head the DRAFTER projects through.

The tied main lm_head unless a drafter-scoped override is
installed (:meth:`set_draft_lm_head_override` — head-quant
"drafter" mode). Verify NEVER reads this: it projects through
``engine.model.lm_head`` directly, so an override here can only
change DRAFT tokens (accept-length risk), never emitted-output
quality.

Install a drafter-scoped lm_head replacement.

Head-quant "drafter" mode: the drafter reads an RTN-quantized
copy of the vocab shard (5×/step bandwidth cut) while verify
keeps the checkpoint-dense lm_head. Held OUT of the module graph
(same 1-element-list trick as ``_tied_modules``) — the engine
owns the override module. The ``MtpDriver`` identity assert
(``head._lm_head is engine.model.lm_head``) is untouched: the
override never replaces ``_lm_head`` itself.

Run the single Qwen MTP block over its own KV slab → ``(B, H)``.

``kv_ctx`` is a length-1 sequence of
``(state_view, attn_meta, rope_cache, attn_op)`` (Qwen ships one
MTP block; the driver chains it for K>1).

Run the head and return ``(B,)`` int64 draft tokens.

``prev_hidden`` is the main model's (or a previous head call's)
post-final-norm hidden state at the slot whose committed token
is ``last_token_ids``; the head consumes both, fuses them, runs
its single decoder block, projects through the main model's
``lm_head``, and returns the argmax over the vocabulary.

``return_hidden=True`` additionally returns the ``(B, H)``
post-norm hidden — used by :class:`MtpDriver` to chain K>1
drafts autoregressively (feed draft_k's hidden in as
``prev_hidden`` for draft_{k+1}). Qwen 3.5 / 3.6 ship a single
head block, so true K>1 multi-block heads are not present —
autoregressive chaining of the same block produces additional
speculation tokens at the cost of compounding drafter error.

Sampling. When ``sampling_params`` is provided AND any
per-request ``temperature > 0``, the head samples its draft
token via :class:`GraphSafeRejectionSampler` (batched
Gumbel-max with a persistent seed buffer; no per-row
``torch.multinomial``) and ALSO returns the FULL drafter
distribution per row (shape ``(B, V)``) as the third tuple
element — the rejection-sampling residual sampler in
:func:`verify_and_accept` consumes ``q_full`` (Leviathan-2023
Theorem 1). ``temperature == 0`` rows collapse to argmax with
a point-mass on the picked token. ``seed_buf`` is required
when ``sampling_params`` is set; ``slot_offset`` lets a
K-step chain produce independent draws across calls.

Default (``sampling_params=None``): greedy argmax everywhere,
no probabilities returned — the legacy contract.

``sampling_tensors`` (mutually exclusive with
``sampling_params``) is the graph-CAPTURABLE sampling path: the
per-row temperature/top_k/top_p knobs and the step seed are read
from device tensors (:class:`DrafterSamplingTensors`), so a
captured drafter chain replays with fresh parameters via
``copy_()`` instead of baking Python values at record time.
Returns ``(tokens, hidden, probs)`` under ``return_hidden=True``
— the same triple as the ``sampling_params`` path.

``wm_row_seeds`` (``sampling_params`` path only) routes the draw
through the keyed-watermark coupled sampler — the slot token
shares its Gumbel noise with the verify pass's keyed choice at
this depth. See ``sample_drafter_token``.

``draft_prefix`` is the ``(j, B)`` int64 block of tokens the
CALLING chain has already drafted this step, at chain position
``j``. It exists so the drafter's proposal carries the same
presence / frequency penalty the verify slate applies at that
position (:mod:`arbi_serve.spec_decode.draft_penalties`) — the
verifier penalises position ``j`` against the committed prefix
plus ``drafts[:j]``, so the proposal must too or every drafted
token is measured against a distribution it was not drawn from.
``None`` (the default, and chain position 0) means the history is
the committed prefix alone.

``top_w > 0`` (greedy path only) returns the drafter's top-``W``
at each row as ``(B, W)`` instead of the ``(B,)`` argmax — one
tree node's worth of candidates per row. The SHAPE differs, so a
caller cannot receive it by accident; ``W = 1`` returns exactly
what the argmax path returns.

``margin_out`` is a ``(B,)`` fp32 buffer the per-slot depth gate
reads its signal from (:meth:`MtpHeadBase.write_slot_margins`).
Filled off the SAME row the pick reduces over, in every branch, so
the gate's threshold means one thing whether this slate drafted
greedily or sampled. ``None`` (the default, and every unarmed
boot) computes nothing.

Recompute + rewrite the MTP layer's K/V at one existing slab cell per row.

The verify/decode fill pass (:meth:`Qwen3_5Model.forward`) pairs
cell ``p`` with ``embed(token[p+1])`` (the next-token shift the
head was trained on) — EXCEPT each request's last in-span cell,
which it pairs with its OWN embedding because the successor token
is unknown at fill time. After verify+accept the committed
history keeps that frontier cell paired with the WRONG embedding
(the rejected draft on partial accepts, itself on full accepts),
and the next verify span starts one position later, so the cell
is never rewritten — one permanently stale drafter-KV key per
verify cycle.

This method runs the SAME fuse + decoder-layer path the fill
pass uses — embed + pre-FC norms + ``fc`` + ``layers[0]`` (same
projections, q/k norm, RoPE at the cell's own position, same
KV-cache write path / codec via ``attn_op``) — with the
now-known committed successor in ``last_token_ids``, so the
rewritten cell is indistinguishable from a correctly-filled one.
The layer's attention output is discarded; only the K/V scatter
into ``attn_meta.slot_mapping`` matters. Idempotent: the write
is a pure function of ``(token, hidden, position)``.

Append one per-step head record to the ARBI_MTP_HEAD_DUMP JSONL.

Records the per-step head INPUTS (committed token id, prev_hidden
mean/L2 fingerprint), the fused-vector + post-norm output
fingerprints, the RoPE positions, and the picked draft token —
enough to diff arbi-serve's per-step feeding against vLLM's
reference without shipping full hidden tensors. No-op (single flag
read) when the dump path is unset. Diagnostic only — runs eager,
outside any captured graph (the probe boots eager).

THE drafter logit row every pick reads — penalty included.

One chokepoint for all three picks in :meth:`forward` (greedy
argmax, the live per-request sampler, the graph-capturable
tensor-driven sampler), so the drafter's distribution and the
verify slate's cannot drift apart one branch at a time. Identity
— the bare ``_draft_lm_head`` row, bit-for-bit — until a
:class:`~arbi_serve.spec_decode.draft_penalties.DraftPenaltyState`
is installed on the head.

Kept in the head's native dtype: the samplers re-cast to fp32
before the softmax, and the accumulator entry the bias comes from
is already in that dtype.

Global greedy top-``width`` over the sharded lm_head.

The tree analogue of :meth:`_distributed_greedy_argmax`. A chain
needs one token per slot and that method never materializes the
full row — it all-gathers a ``(max_value, local_index)`` pair. A
tree needs the drafter's top-W at every node, which is exactly the
information that optimization discards, so this forms the row.

Returns ``(B, width)`` int64, ranked. ``width == 1`` returns the
same token :meth:`_distributed_greedy_argmax` would, which is what
makes a width-1 tree a drop-in for the chain.

REFUSES at TP > 1. A correct global top-W needs a cross-rank
reduction over per-rank top-W pairs — the generalization of
:meth:`_reduce_gathered_argmax` — and taking the local top-W
instead would silently propose tokens that are not the global
top-W on any rank but rank 0. Wrong drafts are not a correctness
bug (verify rejects them) but they ARE a silent quality
regression, which is worse than a refusal.

Global greedy argmax over the sharded lm_head WITHOUT the full gather.

Bit-identical to ``self._lm_head(x).float().argmax(-1)`` but
all-gathers only the per-rank ``(max_value, local_index)`` pair
instead of the full padded-vocab logit row.

Tie-break matches ``torch.argmax`` over the gathered full-vocab
row exactly: lowest GLOBAL vocab index wins. ``max(dim=-1)``
already returns the lowest LOCAL index on a within-rank tie; the
cross-rank reduction picks the lowest ``tp_rank`` on an equal max
value (lowest rank ⇒ lowest global index, because the ranks' column
ranges are contiguous and ascending in rank order, so every id on
rank ``r`` is below every id on rank ``r+1``).

Padded vocab columns are EXCLUDED from the local max: a short
last shard's tail rows are zeroed by the weight loader, and the
gather path trims them via ``[..., :org_vocab_size]`` before the
argmax, so a zero-valued pad column must never be selectable.

``draft_prefix`` opts the pick into the drafter's presence /
frequency penalty (:mod:`arbi_serve.spec_decode.draft_penalties`).
It lands at the PICK, on the row each branch already formed, so
every branch's matmul, its dtype, and its pad trim stay exactly
what they were — with no penalty state installed, or with one
that no row is asking anything of, the returned tokens are
bit-identical to the pre-flag path.

``margin_out`` takes the per-slot depth gate's signal off the row
this pick already reduces over. REFUSED at TP > 1: the branch
below holds a per-rank SHARD, whose top-1/top-2 gap is not the
global row's and would read as a confident drafter on any rank
that happens to own two close logits. The gate is TP1-only for its
own reasons (:mod:`arbi_serve.spec_decode.depth_gate_boot`), so
this refuses rather than reconstructing a cross-rank top-2.

Pick the global-argmax token id from the gathered per-rank pairs.

``gathered`` is the ``(tp_size*B, 2)`` all_gather of each rank's
``(max_value, index)`` pair (fp32). Returns ``(B,)`` int64
global vocab ids. Ties on the max value resolve to the LOWEST
rank (⇒ lowest global vocab index) to match ``torch.argmax`` over
the gathered full-vocab row.

``indices_are_global`` says the packed index is ALREADY the global
vocab id; otherwise it is rank-local and ``rank * per_shard`` maps it
up, which is correct only while every shard is the same width.

Per-parameter shard spec for the head's ``mtp.*`` keys.

The arch-side weight map (``Qwen3_5Model.weight_map`` /
``Qwen3_5MoeModel.weight_map``) binds the main-model parameters;
this function returns the head's own keys, prefixed with the
attribute path under which the head sits on the engine
(``mtp_head.*``). The arch merges the two maps so head + main
weights load in one :func:`load_model_weights` pass.

The whole head binds here, FFN included — the FFN half through
the ``mlp_weight_map`` the arch injected at construction (whose
key set is a function of the arch's FFN shape: dense gated MLP vs
sparse MoE with a per-rank expert slice). Keeping the FFN binding
on the same object that owns the FFN module is what makes the
constructor's refusal total: a head that built cannot then fail
to load, because there is no second place for the arch to forget
to add the keys.

Generalized MTP draft-head base.

A draft head is defined by composable hooks over one fixed pipeline, so
the head-agnostic driver / verify / rejection-sampler code drives every
family through the SAME surface. Two families are instances:

  * :class:`arbi_serve.spec_decode.mtp_head.Qwen3_5MtpHead` — reuse-main-
    embed, ``fc(cat)`` mix, one gated-attn block over its OWN KV slab,
    tied ``lm_head`` logits; advances positions, writes draft KV.
  * :class:`arbi_serve.spec_decode.gemma4_mtp_head.Gemma4MtpHead` — a
    low-rank Gemma-4 EAGLE3 head: verifier-embed input, ``pre_projection``
    mix (2×backbone → internal), N Gemma blocks with Q-only attention
    reading the VERIFIER's KV, ``post_projection`` feedback + centroid
    logits; constant positions, no draft KV.

The two capability flags remove all family branching from the driver
body: the driver reads :attr:`advances_positions` and
:attr:`writes_draft_kv` instead of an ``isinstance``.

Base for bundled MTP draft heads — composable-hook contract.

A concrete head sets the four descriptor attributes and implements
the hooks. :meth:`draft_step` composes them into one speculative
step; the head's own ``forward`` may call :meth:`draft_step` or run
an equivalent inlined body (the production Qwen head keeps its
hand-tuned forward and only exposes the hooks + flags for the
driver).

Descriptor attributes (set on the instance in ``__init__``):

  * :attr:`n_draft` — draft tokens PRODUCED PER HEAD CALL. Both
    shipped families produce 1 (Qwen: one block; Gemma: one 4-layer
    pass); the driver chains for K>1.
  * :attr:`internal_hidden_dim` — the head's internal working width
    (Qwen: main hidden H; Gemma: 256).
  * :attr:`backbone_hidden_dim` — the width of the hidden state the
    head consumes from and feeds back to the engine (Qwen: H;
    Gemma: 1536). Equal to ``internal_hidden_dim`` for a full-rank
    head, larger for a low-rank one.
  * :attr:`advances_positions` — True iff each chained draft token
    occupies a NEW absolute position (Qwen). False iff every draft
    step predicts from the SAME frontier position (Gemma
    ``constant_draft_positions``).
  * :attr:`writes_draft_kv` — True iff the head writes the drafted
    token's K/V into a slab it owns and the prefix grows (Qwen).
    False iff the head only READS the verifier's KV and writes
    nothing (Gemma Q-only, ``skip_kv_write``).

Run the head's internal decoder blocks → ``(B, internal_hidden_dim)``.

``kv_ctx`` is a per-internal-layer sequence of the opaque
``(state_view, attn_meta, rope_cache, attn_op)`` tuples the
engine threads in. The head owns how many internal layers it
runs to produce ONE draft token (Qwen: 1; Gemma: 4).

Project the block-stack output to FULL ``(B, vocab)`` logits.

This is the LOSSLESS path the verify / rejection sampler
consumes — it must span the full vocab (Gemma's centroid
masking fills non-selected positions with ``-inf`` but keeps a
full-width row). Qwen: tied ``lm_head``.

Fast greedy draft token ``(B,) int`` — may use a sparse
argmax (Gemma centroids) that never materializes full logits.

Default: argmax over :meth:`to_logits`. Overridden by heads with
a cheaper argmax path.

Fill ``out (B,)`` with this slot's per-row top-1/top-2 margin.

The ONE place a head hands the per-slot depth gate
(:mod:`arbi_serve.spec_decode.depth_gate`) its signal. An
out-parameter and not a return value because the captured drafter
chain has to land the margin in a persistent buffer its recorded
kernels write and the host reads between segments; a returned
tensor would be a fresh allocation the graph cannot bake.

``logits`` is the row the head's pick already formed — penalties
included, so the margin describes the distribution the drafted
token actually came from. ``out=None`` is the unarmed path and
costs one branch.

MTP per-step metadata — what the engine threads through one verify pass.

The engine builds an :class:`MtpStepPlan` per step that sees any MTP-
opted request. The plan carries:

  - per-request K (the number of draft tokens the head proposed for
    that request);
  - the per-request draft slot indices (from
    :meth:`MultiStatePool.allocate_draft_slots`);
  - the row-flat per-position draft mask (1 for slots that are draft
    tokens being verified, 0 for the committed prefix entry).

Per-block forwards do NOT branch on the mask — the verify pass
relies on the standard attention path: when a request is MTP-opted,
the engine emits ``K+1`` flat tokens (``[last_committed, draft_1,
..., draft_K]``) and ``K+1`` slot indices for that request. The
attention kernel reads/writes the per-position slots exactly as for
a normal multi-token chunk; no per-position branching needed inside
the model. The mask + per-request K live on
:class:`ScheduledBatch.mtp_meta` for the post-forward verify +
gather + free path.

The previous-generation :class:`MTPMeta` in :mod:`arbi_serve.engine.batch`
is kept as a placeholder dataclass; this module supersedes it. The
two are bridged via :func:`build_step_plan`, which the engine calls
once per MTP-routed step.

Per-step MTP verify plan.

The engine constructs one of these per MTP step before running
the main model's verify pass. Lives on
:class:`ScheduledBatch.mtp_meta` (typed-down to :class:`MTPMeta`
via :meth:`as_batch_meta`).

Fields:
    per_req_k: ``(B,)`` int — per-request K (draft count). Always
        ``>= 0``. Under the R1 uniform-K invariant the engine
        coerces every row in a verify-pass slate to a single
        ``step_k`` (or all rows to ``0`` for a cold-path step)
        before constructing the plan; :func:`gather_verify_logits`
        asserts this. The dataclass itself is permissive — older
        unit tests exercise mixed-K layouts to characterize the
        flat-batch arithmetic.
    draft_token_mask: ``(N_tokens,)`` bool — True for positions
        in the flat batch that are draft tokens (i.e. one of the
        K per-request entries appended after ``last_committed``).
        False for the leading committed-token entry. Engine
        consumers use this to gather verify logits.
    per_req_offsets: ``(B+1,)`` int — flat per-token boundaries
        into the verify-row layout. Mirrors ``cu_seqlens_q`` for
        the MTP step (so a request with ``per_req_k=K``
        contributes ``K+1`` flat positions).
    draft_slots: ``(N_drafts,)`` int — concatenated draft slot
        indices returned by
        :meth:`MultiStatePool.allocate_draft_slots`, one entry
        per draft token in flat order. ``len ==
        sum(per_req_k)``. Empty tensor when no row is
        MTP-opted.
    request_ids: ``(B,)`` int — the request ids, in slate
        order. Engine uses these to call
        :meth:`MultiStatePool.free_draft_slots` per-request
        after the verify result.

Build the per-step MTP plan from the engine's slate-side state.

Inputs:

  - ``requests``: slate requests in row order (length B).
  - ``per_req_k``: chosen K per row. Non-MTP rows must pass 0
    here (they contribute a single committed token only).
  - ``draft_slots_per_req``: per-row draft slot lists from
    :meth:`MultiStatePool.allocate_draft_slots`. ``None`` for
    non-MTP rows or rows whose allocation was rejected and
    degraded to K=1; entries must have length ``per_req_k[i]``
    for MTP rows.

Returns the assembled plan. All tensors are int / bool; nothing
in here owns GPU memory beyond a short-lived flat int buffer.

``device_fields=False`` leaves the five device tensors ``None`` and
keeps only the host mirrors. Nothing on a single rank reads the device
copies — the forward tests ``mtp_meta`` for presence, the gather and
the commit read the host mirrors — so a single-rank verify step skips
five H2D copies per step. A tensor-parallel driver ships them to its
workers and keeps them.

Gather per-token verify-pass logits at the slate's uniform K.

The main model returns per-seq last-token logits — but for the
verify pass we wired model.forward to return per-token logits at
every emitted token (achieved by NOT slicing on cu_seqlens_q
when running in verify mode). This function takes the flat
``(N_tokens, V)`` logits tensor and gathers the K+1 verify slots
of every MTP-opted row into the ``(K+1, B_mtp, V)`` tensor that
:func:`verify_and_accept` expects.

When ``verify_buffers`` is supplied (the served path) the gather runs
in place into the pool's persistent ``sc_verify_logits`` scratch via a
single ``index_select`` — NO per-step ``torch.stack`` + ``.permute`` +
``.contiguous`` double O(B*V) copy and NO fresh vocab-scale device
allocation (the frozen-layout / cudagraph-safe invariant). With
``verify_buffers=None`` (or a pool built without the scratch,
``vocab_size == 0``) the legacy stack path runs — byte-identical, just
allocating.

Uniform-K invariant (R1). Within one MTP step every MTP-opted
row contributes the same ``K + 1`` flat tokens — admission +
engine coercion guarantee this. Mixed K per step would force each
distinct K to produce its own ``(K+1, B_k, V)`` shape and one
verify forward per K; a single uniform shape unblocks single-bucket
capture of the verify forward.

Pre-conditions:
  - ``logits.shape[0] == plan.per_req_offsets[-1]``.
  - All MTP-opted rows in ``plan`` share the same per-request K.

Returns ``(K, mtp_row_idx_in_plan_order, main_logits_(K+1, B_mtp, V))``
or ``None`` when no row is MTP-opted (legacy K=1 / cold-path step).

Bridge to :class:`arbi_serve.engine.batch.MTPMeta`.

Returns the slim :class:`MTPMeta` dataclass for stuffing onto
:attr:`ScheduledBatch.mtp_meta`. Per-block forwards consult
only those fields (``per_req_k`` / ``draft_token_mask`` /
``per_req_offsets``); the richer plan (with draft_slots /
request_ids) lives on the engine.

Which MTP-verify widths still need the ``block_m``-keyed split-K family.

The verify band has two kernels behind one dispatch. ``DecodeAttend``
declines a verify chunk to the Turbo prefill tensor-core route whenever
``TKV_MTP_PREFILL_SPLIT`` is on (tkv's default) — every width, not just the
high-K tail — so on a speculative boot of that shape the split-K register
family is compiled and measured for a kernel the process then never
launches. That family is not one ``.so``: it is one module per
``(block_m, tile_tokens, emit_lse, reduce_nwarps, row_tile, range_policy)``
plus a MEASURED autotune table per ``block_m``, and the sweep is the only
site allowed to time it (the pick is a constexpr baked into the captured
graph). Both costs land in the boot's critical path.

This module decides, from the route tkv RESOLVED, which widths still reach
the split-K kernel. It never reads the flag: the caller hands it the floor
that :func:`~arbi_serve.spec_decode.k_calibration.tkv_prefill_verify_min_block_m`
resolved (which asks ``tkv.runtime.attention`` first and only mirrors the
flag when that private symbol is absent), or the answer the built decode
attend gives for each width via ``verify_routes_to_prefill``. With the flag
at 0 the floor is the register ceiling again and the whole family is planned
exactly as before.

A width dropped here must never be launched later. The family is
block_m-keyed, so a first use JITs it on the request path — minutes of nvcc
under a live request, and a deadlock if it lands inside cudagraph capture.
:func:`arm_splitk_family_guard` therefore converts that deferral into a
refusal at tkv's own module accessor, so a boot that skipped the family
either never asks for it or fails loudly naming the flag that would bring it
back. ``mtp_splitk_family_launched`` is the must-not-fire counter for that
edge; :func:`~arbi_serve.flag_truth.flag_truth_report` reports it.

MUST-NOT-FIRE counter: a skipped split-K MTP module was requested.

A fire means the route moved under the boot's decision — the flag was
flipped on a live server, or a layer resolved a route the plan did not
read. The guard raises on the same edge, so a non-zero count is always
paired with a refused boot/request rather than a silent JIT.

Verify widths split by which kernel this boot's resolved route uses.

``build`` are the widths whose split-K modules must be compiled and
autotuned before capture; ``skipped`` are the widths the Turbo prefill
route serves, for which neither is needed. ``prefill_min_block_m`` is
the resolved floor the split was made against — kept so the receipt can
state the evidence, not just the outcome.

Split ``block_ms`` into the widths split-K serves and the ones it does not.

Pure. Mirrors ``DecodeAttend._declines_verify_to_prefill_split``: a
verify width at or above ``prefill_min_block_m`` leaves the split-K
kernel, everything else stays on it. Conservative in the direction that
matters — a width is only dropped when the resolved floor proves the
prefill route owns it, so an unreadable route (floor falling back to the
register ceiling) plans the family exactly as before.

``tree_active`` is carried for the receipt only. A tree cannot be served
on split-K at all (that kernel takes no ancestor mask), and
:mod:`arbi_serve.spec_decode.tree_boot` already refuses the combination
at boot — so a tree plan that still builds the family means the two
refusals disagree, which raises here rather than compiling a family the
tree could never use.

Refuse, rather than JIT, a split-K MTP module the plan dropped.

Returns True when the guard was installed. A no-op when the plan built
the family (nothing to protect) or when tkv is absent. Idempotent: a
second arm re-uses the first wrap and only widens the skipped set, so
the original accessor is never lost behind two layers of wrapper.

Wait for the parked verify step's acceptance copy OFF the loop.

With device-resident drafts the host reaches the drain while the
drafter chain may still be running; the copy lands as soon as the
accept kernel has, well before the chain ends. Under admission pressure
the wait runs in the forward executor so intake and streaming keep
flowing; otherwise the drain's own event wait serves.

Bundled-head MTP strategy (Qwen 3.5 / 3.6, DeepSeek V3).

Single seam: :meth:`run_step`. Partitions the slate into MTP-opted
decode rows (routed through :func:`run_verify_step`) and the
legacy K=1 sub-slate (prefill / non-opted decode, run via
:func:`arbi_serve.engine.run_step.step`). Both halves run
sequentially within one wall-clock step so page-table mutation
stays single-threaded.

The strategy, the verify-pass plan builder (:func:`build_verify_plan`),
and the verify-pass driver (:func:`run_verify_step`) live in the
``mtp_verify`` module; the strategy is a thin engine binding around
them.

Run one slate end-to-end (mixed MTP + non-MTP rows).

Partition: MTP-opted decode rows go to the verify pass; the
rest (prefill / non-opted decode) run via the legacy K=1
``step`` free function. Both halves run within one wall-clock
step.

MTP verify pass — public entry points and re-export surface.

The verify pass is split by responsibility across sibling modules; this
module preserves the import surface so
``from arbi_serve.spec_decode.mtp_verify import X`` keeps working:

  - :mod:`arbi_serve.spec_decode.mtp_verify_plan` — the verify-pass
    :class:`StepPlan` builder (:func:`build_verify_plan`).
  - :mod:`arbi_serve.spec_decode.mtp_verify_offload` — the event-loop
    offload helpers (verify forward / recurrent rollback / drafter chain)
    and the verify-forward ``mtp_block_m`` scope.
  - :mod:`arbi_serve.spec_decode.mtp_verify_accept` — the shared eager
    verify body (:func:`_verify_accept_rollback_commit`) with its
    gather / accept / rollback / commit / drafter-seed helpers.
  - :mod:`arbi_serve.spec_decode.mtp_verify_sync` — the rank-0 synchronous
    verify pass (:func:`run_verify_step`).
  - :mod:`arbi_serve.spec_decode.mtp_verify_async` — the async (SOTA)
    verify pass (:func:`run_verify_step_async`) and its config gating.
  - :mod:`arbi_serve.spec_decode.mtp_verify_spmd` — the SPMD-rank verify
    worker (:func:`run_verify_step_spmd_worker`).

Shared verify / accept / rollback / commit body for the MTP verify step.

:func:`_verify_accept_rollback_commit` is the single eager verify core that
both the rank-0 sync pass (:func:`run_verify_step`) and the rank-symmetric
SPMD worker (:func:`run_verify_step_spmd_worker`) run after their forward:
gather the ``K + 1`` verify logits, verify + accept, let a tap-conditioned
drafter observe, roll back recurrent state on partial accepts, and commit
cold + accepted tokens. The gather / sampler-scratch / mask / drafter-seed
helpers it composes live alongside it here.

Return the engine's persistent rejection-sampler scratch sliced to
``(k, b)``, or ``None`` when no scratch is available.

``None`` (greedy-only / no-MTP stubs, or a ``vocab_size=0`` pool) makes
the sampler keep its fresh-allocation path — the served 27B stochastic
config always has scratch, so the persistent path is exercised in
production; the fallback keeps the CPU stub fixtures working unchanged.

Keyed on ``sc_p_target``, the member every stochastic verify step on the
gathered path writes. ``sc_draft_probs`` is not a pool-wide predicate:
a boot that cannot carry a dense drafter q skips it while still owning
the rest of the scratch, and asking about it here would drop the whole
slate back to fresh vocab-scale allocation.

Return a zeroed ``(k, b, V)`` draft-probs buffer for the verify slate.

Uses the persistent scratch slice when the pool carries one (zeroed in
place, then filled per-row by the caller — no fresh vocab-scale alloc),
else allocates fresh on ``ref``'s device.

A pool built for a boot that cannot carry a dense drafter q has no
``draft_probs`` member, and neither does a CPU stub with no scratch at
all — both land on the fresh allocation. Reaching it in production would
mean a drafter densified on a boot whose
:func:`~arbi_serve.spec_decode._mtp_driver_ops.dense_draft_q_possible`
said it could not; the step stays CORRECT (this buffer is written before
it is read, and the verify tail is eager, outside any capture) and pays
one out-of-layout allocation rather than failing.

Resolve the slate's drafter proposal ``q`` for the rejection sampler.

Returns ``(draft_probs_operand, scratch)``. The single dispatch every
stochastic accept path (sync ``_accept_slate``, async / SPMD-overlap
``_async_verify_accept``) shares:

  * ALL rows carry :data:`POINT_MASS_Q` (the served greedy-draft
    default) → the marker passes straight through — no ``(K, B, V)``
    zeroing, no per-row fills, no dense q at all. The sampler's
    index-form fast path consumes ``draft_tensor`` directly
    (bit-identical accept/recovery/bonus to the densified one-hot).
  * Any dense per-row carry present (true-stochastic rows) → the
    classic densified buffer: zero the persistent scratch slice, copy
    dense rows in, scatter point-mass rows from their draft tokens
    (byte-identical to the retired materialized one-hot carry, so a
    MIXED slate — dense + point-mass carries joined across steps — is
    exactly as correct as before).

Caller guarantees every row satisfies :func:`row_carries_draft_probs`
(the ``have_probs`` gate).

Run the sampler's registered logits-processor chain over the slate.

The speculative draw samples the ``(K+1, B, V)`` slate directly and never
calls :meth:`~arbi_serve.sampler.sampler.Sampler.sample`, so a constraint
wired only into that chain would reach the ``K=1`` decode path and nothing
else. This carries the WHOLE chain onto the slate, in the sampler's own
order: grammar, the n-gram block, the omni speech guard, the
thinking-budget force-close — and whatever is registered next.

Every member's
:meth:`~arbi_serve.sampler.processors.LogitsProcessor.apply_to_verify_slate`
is read DIRECTLY off the processor. A processor that does not implement it
raises here, so a newly registered constraint is either wired to the
speculative draw or it stops the first speculative step — it is never a row
served with less than it asked for.

Apply EVERY verify-time logit transform the main sampler would apply on
the speculative draw — in ONE place, so a per-token constraint can never be
silently dropped on the MTP path.

Ordering is :meth:`arbi_serve.sampler.sampler.Sampler.sample`'s own:
penalties, then logit bias, then the registered processor chain — the
chain read from the sampler rather than named here, so the two draws
constrain a request with the same set. ``_build_p_target`` applies
temperature and the top-k/top-p/min-p clips afterwards, so the slate
reaches the accept test on the same processed-logits manifold the K=1 draw
would have used.

``window`` restricts every applier to one vocab-parallel rank's slice, so
the shard-resident verify draw reaches this SAME entry point rather than
carrying a second list of what to apply. One place still means one place
when there are two draws.

``eff_drafts`` is indexed by SLATE row; every consumer below is indexed by
BUCKET column. The two agree only when every slate row speculates, so the
re-index happens ONCE here rather than in each consumer — one of them
doing it and the other not is how a cold row ahead of an MTP row came to
hand the penalty correction its neighbour's drafted tokens.

The accept test's ``(k, B)`` draft tensor, and the host lists to mask with.

A plan built with the drafts on the device (``StepPlan.drafts_on_device``)
left ``eff_drafts`` holding
:data:`~arbi_serve.spec_decode.device_drafts.PENDING` placeholders and
filled :attr:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers.draft_tensor`
itself, ordered behind the drafter chain by event — the accept reads that
buffer in place. The host values are pulled for that step only when a
consumer reads token IDS: a per-row constraint, or a logits processor
whose slate form says it reads the drafts.

EVERY verify path resolves the slate's drafts through here. Which path
runs is a DRAFTER property (``requires_sync_verify``, which for the
DFlash drafter is ``not dflash_device_slots``) and says nothing about
where the drafts live, so a path that resolved them its own way is a path
that can hand a placeholder to an accept test that indexes by token id —
a device-side out-of-bounds gather, which kills the CUDA context.

Per verify row, its REAL recurrent slab row + ``n_accepted``.

Used by the shared sync verify core (:func:`_verify_accept_rollback_commit`,
run on both rank 0 and the SPMD worker) recurrent rollback; the async path
resolves the same pair on device in :func:`_rollback_recurrent_offloaded`.
Returns ``([], [])`` when there is no recurrent pool or the plan carries no
recurrent row mapping.

EVERY row the verify forward advanced is returned, full-accept and cold
(``k == 0``) rows included, because the replay-mode forward commits no
state: a row left out keeps the state it entered the step with. A cold row
is not in ``mtp_rows`` and has no accept result; its one committed token is
offset 0.

Rollback MUST target the REAL slab row (not the batch-row index), so each
row is mapped through the plan's recurrent ``state_meta`` entry — the SAME
mapping the reconstructed verify batch threaded onto its ``state_indices``.
Replaying the wrong (or sentinel) row would corrupt recurrent state.

Batch the cold-row (``k_eff == 0``) frontier argmax into ONE D2H.

Shared by the SYNC + SPMD per-row commit loops. Gathers the cold rows'
flat offsets from the CPU plan mirror (``per_req_offsets_host`` — no
device sync), runs a SINGLE argmax over their logits, and pulls the
result once (instead of a per-row ``argmax().item()`` GPU sync). ``[]``
when no row is cold.

Verify + accept the gathered slate, then commit the accepted rows.

Greedy slates route the device argmax-match op; a stochastic slate
whose rows all carry the prior step's drafter q_full routes the batched
rejection sampler. Advances the engine seed once, calls
``driver.commit_results``, and returns the per-row ``VerifyResult`` list
(``[]`` when no rows resolved). ``nvtx`` toggles a profiling range.

DIAGNOSTIC (accept-deficit E2): per-slot draft-vs-target-top1 agreement.

Appends one ``{"e2": ...}`` JSONL record per verify step to the
``ARBI_MTP_ACCEPT_DUMP`` file: per row, the per-slot bit
``draft_token == argmax(target_logits)`` (teacher-forced agreement —
the verify logits ARE the target's distribution at each draft
position), plus the PROCESSED-target probabilities ``p'(top1)`` and
``p'(draft)`` per slot via the exact sampler helper
(:func:`~arbi_serve.spec_decode.verify._apply_sampling_params_to_logits`),
the row's accepted count, ``path`` — the accepted ordinals, or
``None`` for the chain, so a reader can tell a scattered tree accept
from a prefix one rather than assuming — and the token ids themselves:
``draft`` per slot and ``top``, the target's ranked candidates whose
first entry IS the argmax the agreement bit is computed against. Together these decompose the
per-slot accept rate
``alpha = P(hit)*E[p'(top1)|hit] + P(miss)*E[p'(draft)|miss]``.
Raw-logits argmax equals processed-p' argmax: temperature
scaling is monotone and top-k/top-p masking always keeps the argmax.

No-op (one flag read) when the dump path is unset. Eager host path
only — per-step D2H + per-slot softmax; NEVER enable in a perf arm.
Gathered-lane logits only: run diagnostic cells with
``ARBI_SHARDED_VERIFY=0`` so every verify row routes here.

Let a tap-conditioned drafter (DFlash) stash per-row context from
the verify forward before the seed-draft tail consumes it.

No-op for the bundled head / external drafter. A failure is COUNTED and
rate-reported by the shared observe-failure seam (the affected rows
cold-path their next step); ``log_label`` names the caller's path in the
first traceback.

Commit the accepted-prefix recurrent state for this step's rows.

Maps every verify row to its real slab row (see
:func:`_compute_partial_rollback_rows`) and applies the pool rollback.
When ``broadcast`` is set (rank 0 under TP>1) the ``(rows, n_accepted)``
pair is sent through the worker bridge first so every rank commits its
shard in lockstep. No-op only on a model with no recurrent state.

``(row_idx, j) -> LogprobEntry`` for every token this step will commit.

Returns ``None`` — and touches nothing — unless the surface is armed AND a
row asked for logprobs, so a deployment that never serves the surface pays
one attribute read and one ``any`` over the slate.

A committed token's distribution lives at the slate position that produced
it: the frontier row for a cold commit, ``cu_q[row] + j`` for the j-th
accepted token. Scoring them together keeps this to ONE head pass per
step rather than one per token.

Commit cold (k=0 frontier argmax) + accepted-prefix tokens per row.

Appends committed tokens to each row's output, frees rejected draft
slots on active rows, and collects the ``(req, last_committed,
bonus_hidden)`` head inputs for the next-step drafter. ``cold_argmax``
is the pre-computed frontier argmax for the k=0 rows (from the gathered
full-vocab logits on the gathered path, or the candidate global-argmax
on the shard-resident path — both byte-exact and one D2H).
``post_token_fn`` runs per committed token when given (a row that
finishes stops its accepted loop); pass ``None`` to skip the streaming
side effect. Returns ``(head_inputs, accepted_per_row)`` (the latter is
each row's committed-token list).

The shallowest drafted block a bucket at ``k_val`` may legally return.

``k_val`` normally. A per-slot depth gate stops the chain by DESIGN, so
on a gated boot a shorter block is the drafter's decision rather than a
truncated buffer — but only down to the gate's own minimum: the first
``lag + 1`` slots have no predictor behind them and are never gated
away, so anything below that is a truncated buffer whatever the flag
says and the refusal must still catch it.

Read as the verify plan reads it — the arming flag AND the drafter's
declared capability, in that order — so no boot can have two call sites
disagreeing about whether a short chain is a decision or a fault, and an
unarmed boot never consults the drafter at all.

Run the offloaded drafter chain for one K-bucket, store next drafts.

Writes each row's ``mtp_next_drafts`` (and ``mtp_next_draft_probs`` from
the drafter q_full for a stochastic bucket, else ``None``); a stochastic
bucket advances the engine seed once. On a chain failure the bucket's
rows are reset to cold-path (empty drafts).

Shared eager verify body — gather, accept, observe, roll back, commit.

The single sync verify core both the rank-0 verify pass
(:func:`run_verify_step`) and the rank-symmetric SPMD worker
(:func:`run_verify_step_spmd_worker`) run after their forward. Given the
forward's ``hidden_full`` + ``per_token_logits``, it gathers the ``K + 1``
verify logits, verifies + accepts the slate, lets a tap-conditioned
drafter observe the forward, rolls back recurrent state on partial-accept
rows, and commits cold (``k == 0`` frontier argmax) + accepted tokens.

``rank_symmetric`` selects the SPMD-worker behavior. Every rank derives
the same partial-accept ``(slab_row, n_accepted)`` from the rank-symmetric
mirror, so the recurrent rollback skips the rank-0 bridge broadcast (the
symmetry IS the broadcast) and the rank-0-only nvtx range is suppressed.
When unset (rank 0) the rollback broadcasts the rows so every worker rolls
its shard back in lockstep. ``post_token_fn`` streams each committed token
on rank 0 and is ``None`` on the worker (no detok / stream there).

Returns ``(results, mtp_rows, head_inputs, accepted_per_row)``.

The bucket's ``(depth, len(bucket_reqs))`` drafted block as
host lists, refusing any other shape.

Every row reads ``drafts[step][col]`` at its own column, so a buffer
that is NARROWER than the bucket would hand rows another row's
tokens (or index past the end deep in the commit tail). Reject it
here: the bucket cold-paths and the breaker accounts it. The row
count is checked against ``floor``, which is the whole difference
between a gated boot and an ungated one.

Async (SOTA) MTP verify pass + its boot-time config gating.

SCOPE: this path runs on the NON-SPMD engine loop only. The canonical
SPMD-TP serve (``ARBI_SPMD_TP``, default ON at TP>1) routes every rank —
including rank 0 — through the rank-symmetric verify core in
:mod:`arbi_serve.spec_decode.mtp_verify_spmd` (a blocking per-step
``_batched_tolist`` egress unless ``ARBI_SPMD_VERIFY_OVERLAP`` is set),
and never reaches this module.

:func:`run_verify_step_async` is the default verify path on that
non-SPMD loop: it
optimistically advances each row as if all K drafts accepted, launches the
next-step drafter chain from GPU-resident outputs, and defers the accept
host pull + commit to the drain. :func:`async_mtp_eligible` /
:func:`assert_async_mtp_supported` decide (per step / at boot) whether the
engine runs this path or the synchronous reference path.

Ring-only watermark advance for an accept that drew no keyed token.

No keyed draw happened, but the context ring must track the committed
tokens or a later keyed step would hash a stale window. D2D from the
on-device tensors, so it keeps the no-host-pull contract.

Device-resident verify accept — greedy OR stochastic, no host pull.

Returns ``(n_accepted (B,) int64, committed (K+1, B) int64,
path_rows (depth+1, B) int64 | None)`` where ``committed[0:n]`` is
the accepted-draft prefix and ``committed[n]`` is the bonus (full
accept) / recovery (partial) token. All stay on the verify device so
the async path can defer the host pull to the drain —
temperature-blind.

``path_rows`` is ``None`` for a chain, whose accepted node at depth
``d`` IS verify-block row ``d + 1``; a TREE's is scattered, and the
drafter seed, the frontier repair and the KV compaction all address
by row rather than by count.

Greedy slate (every row ``temperature == 0``): routes the
``mtp_verify_greedy`` op; ``committed`` IS the per-slot ``main_argmax``
(for ``k < n`` greedy accept means ``argmax == draft`` by
construction, so the prefix equals the accepted drafts).

Stochastic slate (any row ``temperature > 0``): routes the batched
Leviathan-2023 rejection sampler, which is already fully
device-resident (:meth:`GraphSafeRejectionSampler.sample_batched`).
The committed prefix is the accepted DRAFT tokens (not the residual
recovery draws), so we splice ``where(k < n, draft, residual/bonus)``.

Falls back to greedy when a stochastic row's drafts carry no q_full
yet (the FIRST verify step after a greedy / sync seed) — identical to
the sync :func:`run_verify_step` ``any_stochastic and have_probs``
gate. That single step accepts greedily; the next step has probs.

Stochastic TREE slate: routes
:func:`~arbi_serve.spec_decode.tree_accept.tree_accept_stochastic`,
the lossless multi-candidate accept. It is checked FIRST because a
tree's siblings are alternatives for one position and the chain's
per-slot rule would accept them as successors. It carries no
``have_probs`` condition: the tree drafter's proposal is a point mass
per candidate, so the accept needs the target's distribution and
nothing from the drafter, and there is no cold step to degrade.

True iff this verify slate runs the async-verify (SOTA) path.

The async path is the DEFAULT served path for the bundled MTP head at
any TP, with cudagraph, GREEDY or SAMPLED. It optimistically advances
each row as if all K drafts accepted, launches the next-step drafter
chain from the GPU-resident verify outputs BEFORE the host bookkeeping
(intra-step head overlap — the c=1 idle-head kill), defers the
``num_accepted`` host pull + the per-row commit / ``post_token`` /
draft-slot reclaim by one engine tick, and corrects in the drain. At
TP>1 the recurrent rollback is broadcast to worker ranks
(``_rollback_recurrent_offloaded`` → ``broadcast_rollback``), the
device-side equivalent of the sync path's ``RollbackOp`` — so every
rank rolls its slab shard back in lockstep.

SAMPLED (stochastic) is a FIRST-CLASS async case, not a
silent fall-to-sync degrade. The accept decision is the batched
Leviathan-2023 rejection sampler (:meth:`GraphSafeRejectionSampler.
sample_batched`), which runs fully DEVICE-RESIDENT and returns the
per-row ``num_accepted`` + recovery / bonus draws on the verify device
— the same no-host-pull property the greedy op has. So the
optimistic-advance + deferred-commit machinery is temperature-blind:
:func:`_async_verify_accept` produces the device ``(n_accepted,
committed)`` pair the drain consumes for both modes. (A stochastic row
whose drafts carry no q_full yet — the first step after a sync / greedy
seed — accepts greedily for that one step, then has probs.) Since the
honest serving workload uses a sampler (greedy degenerates the instruct
model), routing sampled to sync would silently leave the SOTA async
path unused in practice — exactly the degrade this removes.

This returns ``False`` ONLY for the transient per-step COLD step — a
step where every row resolves to ``K < 1`` (drafter miss / slot-alloc
fail). That is the K=1 decode-via-verify path: there are no drafts to
verify, so there is nothing to advance optimistically. Per-step, not a
config gate.

CONFIG-LEVEL incompatibilities (the async path genuinely cannot run
for the whole engine: a DFlash drafter that taps the sync verify
step, or a recurrent pool with no device-side rollback) are NOT
handled here by a silent reroute — :func:`assert_async_mtp_supported`
REFUSES them LOUDLY at boot so an operator never silently runs the
slow sync path below SOTA (kill-footguns). By the time the run loop
reaches this gate, the config is known async-capable; only the
transient cold step above routes to the synchronous path.

The engine config cannot run the async MTP SOTA verify path.

Raised LOUDLY at boot (``assert_async_mtp_supported``) instead of
silently routing every served request to the slow synchronous verify
path. A silent sync reroute is invisible to the API caller and is a
SOTA regression (kill-footguns) — so a config that truly cannot run
async refuses the boot rather than degrading silently. Resolve by
fixing the config (wire the device-side rollback) or by explicitly
opting into the reference sync path with ``ARBI_ASYNC_OUTPUT=0``.

Boot-time gate: REFUSE configs that cannot run async MTP.

Called once at MTP-driver attach. The async path is the served SOTA
path; we never silently fall back to the sync path mid-flight (that
degradation is invisible to the caller). So at boot, with the driver
and pool known, either the config can run async — or we raise
:class:`AsyncMtpUnsupportedError` so the operator sees it and fixes
the config (or explicitly opts into the sync reference path).

No-op when MTP is not attached, or when the operator has explicitly
selected the sync reference path with ``ARBI_ASYNC_OUTPUT=0`` (the
sanctioned opt-out — logged, not silent).

Per-slot accept diagnostic for the async path (``ARBI_MTP_ACCEPT_DUMP``).

The sync path writes the ``e2`` record from its host-side results; the
async path holds the accept count on the device, so this pulls it once
and hands the same record writer the same shape. One flag read when the
dump is off; a host sync per step when it is on, so never in a perf arm.

Async verify pass — GREEDY or SAMPLED, attention-only OR hybrid.

A's K=1 deferral for the spec-decode verify commit:

  1. Build the verify plan + run the main-model forward (same as
     sync). Run :func:`_async_verify_accept` to obtain the DEVICE
     ``n_accepted`` ``(B,)`` and ``committed`` ``(K+1, B)`` — NO host
     pull. Greedy slates route the ``mtp_verify_greedy`` op; sampled
     slates route the device-resident batched rejection sampler.
     ``committed`` carries the accepted-draft prefix + the bonus /
     recovery frontier token; the drain reads ``committed[0:n+1]``.
  2. Seed the next-step drafter chain from GPU-GATHERED inputs:
     ``last_committed = committed[n_accepted[row], row]`` and the
     bonus hidden ``hidden_full[cu_q[row] + n_accepted[row]]``. The
     gather reads the device ``n_accepted`` directly, so the drafter
     forward (which produces the host ``mtp_next_drafts`` for the
     NEXT plan) does not wait on the verify host pull. Sampled rows
     also stash the drafter's per-slot ``q_full`` in
     ``mtp_next_draft_probs`` so the next step's rejection sampler has
     the full distribution (Leviathan-2023 Theorem 1).
  3. OPTIMISTICALLY advance each row as if all K accepted (the
     ``_mtp_optimistic_len`` advance — placeholder tokens never reach
     ``output_token_ids``) and do NOT free the rejected draft slots.
  4. Issue the non-blocking D2H of ``(n_accepted, committed)`` on the
     side stream + record a CUDA event, and PARK a
     :class:`PendingVerifyStep` on ``eng._pending_verify``.

The drain (:func:`run_step.drain_pending_verify`, top of the next
step) materializes ``n_accepted``, TRUNCATES the optimistic advance to
the real ``n_accepted + 1`` committed tokens, frees the ``K - n``
rejected draft slots, runs ``post_token`` (the streaming side effect —
R1: confirmed tokens only) and updates the accept counters. A row
finished mid-flight drops its deferred output.

Caller MUST have checked :func:`async_mtp_eligible`; this function
asserts active K and falls back to the sync path on a violated
precondition (cold / non-uniform-K step).

The async pass from its forward's outputs to the drafter seed — steps 1
(accept) and 2 (seed) of :func:`run_verify_step_async`, plus the
device-resident recurrent reconcile. Runs under the caller's
``_mtp_block_m_scope`` + ``inference_mode``, exactly as it did inline.

Split from the forward so the SAME tail serves a verify slate whether its
forward ran alone (:func:`run_verify_step_async`) or fused with the
step's prefill chunk (:mod:`arbi_serve.runtime.fused_mixed_step`); the
outputs are indexed by the plan's own ``cu_q``, which start at 0 in both
layouts. Returns what :func:`_async_verify_park` needs, or ``None`` when
no verify row resolved.

Event-loop offload machinery for the MTP verify step.

The verify forward, the recurrent rollback, and the drafter chain each
run OFF the asyncio event loop (in the engine's forward executor) so the
FastAPI intake / detok / SSE coroutines keep flowing while the blocking
GPU section runs. :func:`_mtp_offload_active` decides per step whether to
offload; :func:`_mtp_block_m_scope` bumps the verify block-M on the attn
ops for the forward.

True iff the per-step MTP forward/drafter/rollback should OFFLOAD.

The verify-forward + drafter-chain + recurrent-rollback offloads
exist to free the asyncio event loop while the
blocking GPU section runs, so FastAPI intake / detok / SSE coroutines
keep flowing under load (the Little's-Law concurrency fix). Each
offload, however, costs a ``run_in_executor`` thread hop PLUS a forced
``torch.cuda.synchronize()`` (a full GPU drain) inside the worker — and
an MTP step pays this TWICE (verify forward + drafter chain) vs plain
decode's once.

At SINGLE-STREAM (no other request to admit / stream) there is nothing
to overlap, so that thread-hop + double full-drain is pure per-step
latency on the critical path. Running the forward + drafter INLINE at
single-stream drops the thread hop and lets the next host work
implicitly serialize on the compute stream — no hard global drain —
recovering MTP's single-stream latency without changing any token.

The decision uses :func:`arbi_serve.engine.run_step._has_admission_pressure`
— offload only when the scheduler has requests queued for admission
(``waiting`` / ``_deferred_admits``), i.e. when the freed loop has real
intake work to absorb. Without pressure (single-stream / steady-state)
the offload is skipped and the helpers fall back to their byte-identical
inline path.

No scheduler wired (unit / single-thread paths) ⇒ keep the prior
always-offload behaviour (offload when an executor is present).

Whether an offloaded launch waits for its kernels in the worker.

With device-resident drafts the step's one wait is the drain's copy
event, taken off the loop at the head of the next step; every launch
before it returns as soon as it is enqueued so the accept, the chain
and the next plan build all overlap the GPU. Without them the wait
stays with each launch, as before.

Run the verify main-model forward + ``lm_head`` OFF the event loop.

The single biggest blocking section of an MTP step is the verify
forward (``forward_plan``) plus the ``lm_head`` over the ``K+1`` flat
tokens. Run inline on the asyncio loop, it holds the loop for ~10 ms
per step, so the FastAPI intake / detok / SSE coroutines can't run —
new arrivals never reach ``scheduler.add`` and in-flight streams
stall. The K=1 decode path already fixed this for plain decode via
:func:`arbi_serve.engine.run_step.step_async` (``run_in_executor``);
the MTP verify path was explicitly left inline (see ``step_async``
docstring) and is the dominant cause of MTP's heterogeneous-traffic
concurrency collapse (~3 active streams vs no-MTP's ~12).

This mirrors the K=1 fix: the blocking forward + ``lm_head`` run in
``eng._forward_executor`` (the SAME single-worker thread the K=1 path
uses, so the single-threaded-CUDA invariant holds — at most one
forward thread is ever in flight because the engine ``await``s each
step before the next), with a ``torch.cuda.synchronize()`` inside the
worker so the GPU compute COMPLETES in-thread (releasing the GIL for
the duration). The loop is free to admit / stream during the forward.
The returned device tensors are then consumed by the verify + commit
tail back ON the loop thread — those touch per-request state /
``asyncio`` events that are loop-only.

The ``mtp_block_m`` scope is entered on the loop thread (it only sets
attributes on the attn ops, read by the forward in the worker) and
held across the ``await`` so it's still active when the worker runs;
the caller resets it after this returns. Falls back to a synchronous
in-thread call when offload is disabled or no executor is wired (test
/ single-thread paths) — byte-identical, just not loop-friendly.

GDN recurrent accepted-prefix commit, issued OFF the loop.

Runs for EVERY verify step, full-accept included: the replay-mode verify
forward commits no state, so a row that is not replayed keeps the state
it entered the step with. The commit is issued INSIDE the forward
executor (the same single-worker thread the verify forward + drafter
chain use), exactly like :func:`_verify_forward_offloaded`, so the loop
stays free to admit / stream.

Ordering: the commit is issued before this coroutine returns, hence
before the caller parks the deferred step, hence before step N+1's GDN
forward — and the worker's ``synchronize`` forces it to COMPLETE in
thread, so it lands on the compute stream ahead of step N+1.

Run the next-step drafter chain (``MtpDriver.draft``) OFF the loop.

The symmetric layer-4 analogue of :func:`_verify_forward_offloaded`:
the residual MTP host-overhead floor is the bundled drafter chain +
seed-draft. Without this offload, the ``K``-step autoregressive head
forward (plus the per-step argmax / rejection sampler) runs INLINE on
the asyncio loop — tens of kernel launches/step that hold the loop,
starving the FastAPI intake / detok / SSE coroutines exactly like an
inline verify forward would. ``MtpDriver.draft`` itself is

``sync=False`` skips the in-worker completion wait: the caller keeps the
drafts on the device and orders its consumers behind them by event, so
the host returns as soon as the chain is enqueued.

``draft_kwargs`` are the keyword arguments for ``driver.draft(...)``;
the call returns the device draft tokens, or the ``(tokens, q_full)``
tuple for the stochastic contract. It is run in
``eng._forward_executor`` — the SAME single-worker thread the verify
forward + K=1 decode use, so the single-threaded-CUDA invariant holds
(the engine ``await``s each step before the next; at most one forward
thread is ever in flight). A ``torch.cuda.synchronize()`` inside the
worker forces the head-chain kernels to COMPLETE in-thread (releasing
the GIL), so the loop-free window covers the real compute and not just
the async launches. The returned device tensors are consumed
(``.cpu().tolist()`` + per-request ``mtp_next_drafts`` assignment) back
ON the loop thread, where the request state lives.

``draft`` does host-side per-row draft-slot allocation / TP broadcast /
slot-freeing against ``eng.pool``; running the whole call in the worker
is safe for the same reason the verify forward is — admission-time pool
mutation is gated by ``_forward_in_flight`` while the worker runs, so no
``add`` coroutine touches the pool mid-chain. Falls back to a
synchronous in-thread call when offload is disabled or no executor is
wired (test / single-thread paths) — byte-identical, just not
loop-friendly.

MTP verify-pass :class:`StepPlan` builder and its host-side helpers.

:func:`build_verify_plan` allocates ``K + 1`` slots per row, stages the
flat batch into the persistent :class:`VerifyBuffers`, and returns the
:class:`StepPlan` the verify driver runs over.
:func:`_resolve_verify_recurrent_rows` / :func:`_verify_block_table_rows`
are the host-side support it shares with the verify driver.

Recurrent kind + per-row REAL slab-row indices for the verify slate.

Mirrors :meth:`EagerModelRunner._resolve_recurrent_rows`: pick any
recurrent-kind pool view and resolve each slate request's allocated
slab row via :meth:`RecurrentStatePool.row_for`. Returns the
``(kind, rows)`` pair so the caller can store the mapping under the
model's recurrent :class:`StateKind` in ``StepPlan.state_meta``
(the verify reconstruction routes it back onto
``ScheduledBatch.state_indices``). Returns ``None`` when the model
has no recurrent pool (pure-attention — no recurrent state to
index). A mock pool view without ``row_for`` (test path) sentinel-
routes every row to slab row 0.

Per-row padded block-table for the verify slate, cached on ``eng``.

Decode is append-only, so a request's padded row only changes when it
gains a page (or the slate's ``max_pages`` grows). The persistent
cache (``eng._verify_bt_cache``: ``req_id -> (n_pages, max_pages,
padded_row)``) skips the ``page_ids() + [0]*pad`` rebuild on the
unchanged common case. Self-cleans to the live slate each step.
Bit-identical to ``page_table.block_table_row`` per row.

Returns ``(rows, occupant_keys)`` where ``occupant_keys[i]`` is the
``(req_id, n_pages)`` pair :meth:`VerifyBuffers.load_step` uses to
skip the pinned-host rewrite of a row whose page chain is unchanged
since the prior step (the host-staging analogue of this cache).

Resolve the slate's uniform step-K + each row's cached drafts.

``effective_row_k`` RAISES (never silently clamps) if any row's mtp_k
exceeds max_k — admission already rejected those, so reaching here is a
bug. Asserts the uniform-K invariant (the scheduler bucketed at
admission via :meth:`Scheduler._bucket_mtp_k_uniform`), floors K to 1,
and pulls each row's cached ``mtp_next_drafts`` prefix.
``can_step_active`` is False (and that row's drafts ``None``) on any
drafter cache miss — the caller then collapses the whole step to K=0.

Allocate the per-row tail + draft slots (tail first, drafts second).

Returns ``(tail_slots_per_req, eff_per_req_k, eff_drafts, eff_slots)``
— the per-row effective K is uniform across the slate. On any row's
draft-slot allocation failure the WHOLE STEP collapses to K=0
(already-allocated draft slots are freed and the ``eff_*`` outputs
reflect K=0 — every row contributes one ``last_committed`` flat token
through the K=1 decode-via-verify path).

Assemble the flat verify batch over the resolved per-row K / slots.

Per row i, contribute ``K_i + 1`` flat tokens at positions
``[pre_len_i - 1, pre_len_i, ..., pre_len_i + K_i - 1]`` using slots
``[tail_slot_i, draft_slots_i[0], ..., draft_slots_i[K_i-1]]``. Returns
``(flat_ids, flat_positions, flat_slots, cu_q, seq_lens,
per_req_max_pages)``. Raises on a per-row page-table desync (a row whose
``seq_len`` outran its allocated pages — the verify attention kernel
would read past the block_table row).

Whether the step plan materializes its five device tensors.

A tensor-parallel driver ships them to its workers; a CPU engine (the
stub engines in tests) has no H2D to save. A single CUDA rank reads only
the host mirrors, so it skips them.

Build the verify-pass :class:`StepPlan` for ``mtp_slate``.

Per request, allocate ``K + 1`` slots (one tail slot for the
re-write of ``last_committed`` plus ``K`` draft slots), stage the
flat batch into the engine's persistent :class:`VerifyBuffers`,
and return the :class:`StepPlan` the driver runs over. The plan
transports the per-step routing header (paged-KV tensors + scalar
attn fields + recurrent slab-row mapping) in ``state_meta`` and
carries the :class:`MtpStepPlan` + per-row K / draft / cu_q
bookkeeping as private payload fields (see :class:`StepPlan`
docstring). :meth:`EagerModelRunner.forward_plan` reconstructs the
transient :class:`ScheduledBatch` from ``state_meta`` — there is no
staged ``ScheduledBatch`` payload.

Uniform-K invariant. The slate is uniform-K by construction (the
scheduler bucketed at admission via
:meth:`Scheduler._bucket_mtp_k_uniform`); we assert the invariant
here. On any row's drafter cache miss OR draft-slot allocation
failure, the WHOLE STEP collapses to ``K=0`` — every row
contributes one ``last_committed`` flat token through the K=1
decode-via-verify path.

Engine wiring for shard-resident MTP verify sampling (``ARBI_SHARDED_VERIFY``).

Replaces the per-cycle full-vocab verify ``all_gather`` — the K+1 verify
rows' :class:`~arbi_serve.models.linear.VocabParallelLMHead` gather — with
a tiny exchange of per-shard TOP-CANDIDATE ``(value, global_index)`` pairs,
then runs the SAME accept + recovery/bonus sampling over the compacted
candidate rows, rank-symmetrically. The pure, process-group-free selection
math lives in :mod:`arbi_serve.spec_decode.verify_sharded` (CPU-testable);
this module is the engine glue: the sharded lm_head shard-matmul + pair
all_gather, the verify-slot candidate gather, the ``VerifyResult`` assembly
matching :mod:`arbi_serve.spec_decode.verify`, and the tie-overflow metric.

A per-token constraint does NOT force the slate off this path: it restricts
exactly to a rank's vocab slice, so :func:`_slate_constrainer` applies it to
the LOCAL logits BEFORE the per-shard top-C and the selection ranks
already-transformed values. It reaches the constraints through
:func:`~arbi_serve.spec_decode.mtp_verify_accept.apply_verify_masks` — the
gathered path's OWN entry point, taking a vocab window — so the two draws
cannot come to constrain a request with different sets.

The gathered path (:func:`_verify_accept_rollback_commit`) stays the honest
general-support sampler — this is the specialized fast path for the common
grammar-free, finite-top_k stochastic serving slate, gated on and off at
:func:`sharded_verify_active`.

Log on every OFF<->ON transition, tagged with this process's TP rank.

Level keys on whether the OPERATOR set the flag (env var or live config
override), for the same reason as :func:`announce_gate_refusal`:

  * EXPLICIT (a live A/B flip, or a ``ARBI_SHARDED_VERIFY=0`` kill-switch
    boot) ⇒ **WARNING**, which ``server/logging_config`` keeps on ALL
    ranks (it suppresses INFO/DEBUG off rank 0). "Did the live flip reach
    EVERY rank?" is exactly the question this log answers — a flip that
    reached only rank 0 would desync the TP collectives, so the A/B counts
    ``tp_size`` lines. A transition is rare (an arm switch), not spam.
  * the shipped default engaging on an ordinary boot ⇒ **INFO**. A
    default-ON path coming up is not a warnable event; a WARNING on every
    TP>1 boot is the noise that gets warnings ignored. Rank symmetry is
    still provable here from the ``sharded_steps_total()`` counter, which
    every rank exposes.

Announce ONCE that the gate refused a slate and the GATHERED path ran.

``ARBI_SHARDED_VERIFY`` ships ON, so the flag being true is NO A TP=1 boot, a grammar-
constrained row, and an oversized-top_k row all route to the gathered
sampler BY DESIGN (it is the correct general-support path, not a fallback
shim). Warning on those would fire on every ordinary boot and train
operators to ignore the line — so the level keys on WHO asked:

  * the operator EXPLICITLY set the flag (env var or live config override)
    and the gate refused anyway ⇒ **WARNING**. They deliberately requested
    the path and did not get it: a bogus A/B arm or a misconfigured boot.
    This is exactly the case the original latch existed to catch, and it
    still fires.
  * the flag is merely at its shipped default ⇒ **INFO**. Running the
    gathered path here is a CORRECT outcome. It stays observable — "which
    path ran?" is answerable from the log — without crying wolf.

Either way this is one-shot, and never a silent fallback.

``getattr(obj, name)``, raising when the attribute is absent.

The screen is only as good as its ability to SEE the field it screens on.
A plain ``getattr(obj, name, None)`` reads a shim that never carried the
field as an unconstrained row and admits it to the compacted draw — the
failure this function exists to make impossible. Absent is a wiring bug,
not a value.

Whether any row on the slate carries or requires a grammar.

``grammar_state`` is built lazily on the first sampled step, so a row can
be constrained before its matcher exists; ``grammar_required`` (set at
admit on every rank from the same spec, cleared only on a rank-identical
compile degrade) is what makes the read correct on that first step rather
than one step late. Both are rank-replicated, so every rank answers the
same and the TP collectives stay paired.

Whether this rank's vocab shard slices the packed bitmask cleanly.

Asks the WINDOW, which is the object the masker will be handed, so the
screen and the requirement cannot come to disagree about what "aligned"
means. At the served geometry the answer is always yes, so this is a
screen rather than a live restriction — but it is a property of the vocab
size and the TP degree, so it has to be checked rather than assumed.

Rank-symmetric gate for the shard-resident verify path.

True iff (all rank-replicated state, so every rank agrees and the TP
collectives stay paired):

  * ``ARBI_SHARDED_VERIFY`` is set;
  * TP > 1 (at TP=1 the gather is a no-op — nothing to eliminate);
  * ``ARBI_SPMD_VERIFY_OVERLAP`` is OFF (the device-resident overlap
    body owns its own gather; the two paths are mutually exclusive);
  * the lm_head exposes a rank-local vocab-shard logits surface — a
    dense vocab-parallel ``weight`` OR a quantized column-parallel
    ``local_logits`` (see :func:`_lm_head_vocab_shard`);
  * a grammar row lands on a vocab shard that slices the packed bitmask
    on a word boundary (:func:`bitmask_aligned`) — the mask itself IS
    served here, by masking this rank's own words;
  * every row is ACTIVE uniform-K (a cold k=0 row still routes here —
    its frontier argmax is candidate-exact — but a mixed/zero slate
    with no MTP row has nothing to shard);
  * every STOCHASTIC row has a finite ``top_k`` within the per-shard
    candidate budget (``0 < top_k <= n_cand``) so the compacted rows
    cover the kept set; an unbounded/oversized-support row cannot be
    candidate-compacted and falls back to the gathered path.

A greedy-only slate (every ``temperature == 0``) is eligible — its
per-slot argmax is candidate-exact.

The vocab-sharded lm_head module, or ``None`` if the head carries no
rank-local vocab-shard logits surface.

Resolves the SAME module the drafter's distributed greedy argmax uses — an
untied :class:`~arbi_serve.models.linear.VocabParallelLMHead` or the
sharded :class:`~arbi_serve.models.layers.VocabParallelEmbedding` behind a
tied head — and confirms it exposes the vocab-shard interface
:func:`build_flat_candidates` needs: the ``_tp_size`` / ``_tp_rank`` /
``org_vocab_size`` metadata AND one of the two shard-logits surfaces:

  * a DENSE 2-D ``weight`` (the batch-invariant ``F.linear`` shard
    matmul), or
  * a ``local_logits`` callable (a weight-quantized column-parallel head
    — EXL3 trellis, AWQ pack — whose shard-local quant GEMM is
    bit-identical to this rank's slice of the gathered full-vocab row;
    see :meth:`~arbi_serve.weight_quant.base.ColumnParallelMixin.local_logits`).

A head with the metadata but NEITHER surface — and any head missing the
metadata (``LinearBase.weight`` is a class annotation, not an attribute,
so ``getattr`` returns ``None``) — returns ``None`` and the gate routes
the slate to the gathered full-vocab verify path.

This rank's un-gathered lm_head shard logits for ``hidden``.

Dense vocab-parallel head: the batch-invariant ``F.linear`` on the shard
``weight`` — bit-identical to the per-rank matmul the gathered
column-parallel head concatenates. Quantized vocab-parallel head: the
shard-local quant GEMM (``local_logits``) — the same kernel + input the
gathered forward all-gathers, so the local row is bit-identical to this
rank's slice of the gathered full-vocab row. Either way the sharded
path's candidate values equal the gathered path's values at those
indices, preserving that head's verify equivalence class unchanged.

Return ``(shard_module, tp_size, tp_rank, per_shard, org_vocab)`` for
the vocab-sharded lm_head. The gate (:func:`sharded_verify_active`) has
already excluded a head with no vocab-shard surface via
:func:`_lm_head_vocab_shard`; this reasserts the invariant and fails LOUD
if the fast path is ever reached with a head that has no vocab shard.

``per_shard`` is the rank-LOCAL padded shard width: ``weight.shape[0]``
on a dense head; ``out_features`` on a quant TP mixin (which stores the
LOCAL width there — the same quantity).

Row order that puts the MTP verify slots first, POSITION-major.

The flat verify rows arrive interleaved — each slate row contributes its
own ``k+1`` (or, cold, one) consecutive tokens. The per-token constraint
appliers are written against the ``(K+1, B, V)`` slate the gathered path
uses, and the compacted accept math wants the same shape, so this
permutation buys BOTH: after it, the MTP slate is ``local[:n_mtp].view(K
+ 1, B, ...)`` — a contiguous VIEW rather than a per-row stack — and the
cold rows are the named positions behind it.

Derived entirely from the CPU plan mirror, which is rank-replicated, so
every rank builds the same permutation and the pair all_gather stays
paired.

Whether anything would be written to this slate.

Two sources, both asked rather than assumed: the free-function appliers
(penalties, ``logit_bias``) answer from the slate's sampling knobs, and
every registered processor answers for itself via
:meth:`~arbi_serve.sampler.processors.LogitsProcessor.constrains_verify_slate`.
A caller-side list of "which flags mean constrained" would answer for the
processors that existed when it was written; a new one would then be
skipped on THIS path alone, which is the silent-drop failure the
constraint matrix exists to make impossible.

The slate's constraint pass, bound to this step, or ``None``.

Routes to ``apply_verify_masks`` — the GATHERED path's own entry point —
with the rank's vocab window. Ordering, the re-index from slate rows to
bucket columns, and which appliers run are therefore decided in exactly
one place for both draws; a constraint wired into that chain reaches this
path by construction rather than by someone remembering to mirror it.

``None`` when nothing would be written: the appliers already early-out on
a neutral slate, but returning ``None`` also skips the
:class:`SlateLayout` reorder, so the unconstrained slate — the common one
— pays nothing at all for this capability.

Per-shard top-``n_cand`` lm_head candidates for every flat token.

Runs THIS rank's lm_head shard logits (:func:`_local_shard_logits` — the
batch-invariant dense shard matmul, or the shard-local quant GEMM on a
weight-quantized head; both bit-identical to this rank's slice of the
gathered row), takes the local top-``n_cand`` ``(value, global_index)``
pairs, and all_gathers ONLY those pairs across the TP group — the
verify-sampling generalization of the drafter's ``(B, 2)`` scalar-pair
gather. Returns ``(cand_vals (N_flat, C), cand_idx (N_flat, C) int64,
gathered_pairs (tp, N_flat, n_cand, 2))`` with ``C = tp * n_cand``,
canonically ordered by global index (lowest-global-index tie-break).

No full-vocab ``(N_flat, V)`` logits and no ``(N_flat, V)`` all_gather
are ever materialized — the exchanged tensor is ``tp * N_flat * n_cand
* 2`` fp32 (KB-scale) vs the gathered path's MB-scale row.

``constrain`` (with ``layout``) applies the slate's per-token transforms
to the LOCAL logits BEFORE the top-C selection, so the candidates are
ranked on already-transformed values and the kept set stays a pure
``top_k`` of the same values the gathered draw would have ranked. The
matmul itself runs on the UNPERMUTED rows: the permutation is applied to
its output, so the claim that a candidate value equals the gathered
path's value at that index does not rest on the GEMM being
row-order-invariant. Returned rows are in ``layout.perm`` order.

Drafter ``q`` at the draft tokens and at the slot-k candidates.

Per row, ``mtp_next_draft_probs`` is either the :data:`POINT_MASS_Q`
marker (the served greedy-draft default: q is the one-hot at the draft
token — ``q_at_draft = 1``, ``q_at_cand = 1[cand == draft]`` — no dense
``(K, V)`` tensor exists or is materialized) or a dense ``(K_carried,
V)`` per-step distribution (true-stochastic drafts), gathered at the
draft token (``q_at_draft (K, B)``) and at the residual-slot candidate
ids (``q_at_cand (K, B, C)``). A ``PAD_IDX`` candidate contributes
``q = 0`` (never a real token). Neither form reads the eliminated
full-vocab all_gather.

Verify + accept the slate over compacted candidates → ``VerifyResult``.

Greedy slate → the candidate-exact greedy op
(:func:`sharded_verify_greedy`); a stochastic slate carrying q_full →
the compacted rejection sampler (accept test + Philox-keyed
recovery/bonus draws AT the candidate indices, byte-identical noise per
token to the gathered kernel). Advances the engine seed once (matching
the gathered ``_accept_slate``'s stochastic advance), calls
``driver.commit_results``, and folds the tie-overflow count into the
single host pull. Returns the per-row ``VerifyResult`` list.

Shard-resident twin of ``_verify_accept_rollback_commit``.

Builds per-shard candidates from ``hidden_full`` (the tiny pair
all_gather in place of the full-vocab lm_head gather), gathers the K+1
verify-slot candidates, verifies + accepts over them, then reuses the
SHARED observe / recurrent-rollback / cold+accepted-commit tail
unchanged (those read ``results`` / ``hidden_full`` / the plan — never
the eliminated full-vocab logits). Returns the identical
``(results, mtp_rows, head_inputs, accepted_per_row)`` contract the
caller commits, so the drafter-chain tail is untouched.

Cold-row (k==0) frontier argmax from the flat candidates — one D2H.

Candidate-exact analogue of
:func:`arbi_serve.spec_decode.mtp_verify_accept._cold_row_argmax`: the
global argmax is always a candidate (every shard's local top-1), so the
lowest-global-index winner is byte-exact to the full-row argmax.

Reads the cold rows at their POST-permutation positions. Like the
gathered path's cold argmax, these rows are drawn from untransformed
logits: ``apply_verify_masks`` runs over the MTP bucket only, and the
shard-side twin matches it row for row rather than quietly constraining
a draw the gathered path leaves alone.

SPMD-rank MTP verify worker (``ARBI_SPMD_TP``).

Under SPMD every rank runs the SAME rank-symmetric verify work rank 0
issues, so the TP collectives match peer-for-peer. The worker has no
scheduler / request map, so :class:`_WorkerVerifyRequest` exposes the
:class:`Request` subset the shared verify code reads off a
:class:`RankSlateMirror` row, and :func:`run_verify_step_spmd_worker`
drives the shared eager verify core (:func:`_verify_accept_rollback_commit`)
plus the drafter chain, dropping the rank-0-only side effects.

A worker rank's Request-shaped view over a SPMD mirror row.

Under SPMD the worker rank has no scheduler and no real
:class:`arbi_serve.engine.request.Request` objects — only the
:class:`arbi_serve.distributed.spmd.RankSlateMirror`. The MTP verify
forward + drafter chain (:func:`build_verify_plan`,
:meth:`MtpDriver.draft`) are written against ``Request`` fields, so
this thin shim exposes exactly the subset those read off a mirror row
— keeping the worker on the SAME ``build_verify_plan`` /
``verify_and_accept`` / ``draft`` code as rank 0 (so it inherits greedy
losslessness) WITHOUT depending on the full ``Request`` schema.

Fed the identical mirror as rank 0's request map, the shim's
``request_id`` / ``total_length`` / ``last token`` / ``mtp_next_drafts``
are bit-identical to rank 0's request, so ``build_verify_plan``
allocates the SAME page-table slots in the SAME order and the verify
forward sees the identical flat batch. The worker NEVER streams /
commits / runs ``post_token`` — that stays rank-0-only.

True iff this SPMD verify step runs the device-resident overlap body.

The gate for :func:`_spmd_overlap_accept_and_draft` — the
``ARBI_SPMD_VERIFY_OVERLAP`` sync-kill that removes the intra-step
forward→drafter host bubble. It is read on EVERY rank AFTER the (rank-
symmetric) verify plan is built, and its verdict MUST be identical on
every rank so the overlap-vs-sync choice never desyncs the TP collectives
or the ``mtp_seed`` advance schedule. It reads ONLY rank-symmetric state:

  - ``runtime_flags().spmd_verify_overlap`` — a process-wide env flag, set
    identically on every rank at boot.
  - the driver kind (``requires_sync_verify`` — a DFlash drafter taps the
    SYNC verify forward's state, so it has no device-resident path); the
    driver is replicated on every rank.
  - ``eff_per_req_k`` — the per-row effective K, derived from the
    replicated mirror + the lockstep rank-local pool allocation (the
    SAME uniform-K the sync worker sees). A COLD step (any row K=0, the
    decode-via-verify path) has nothing to advance device-side and falls
    back to the synchronous inline worker; the uniform-K invariant means
    an active step is all-K.

Never reads the rank-0-only scheduler / request map, so a worker rank and
rank 0 reach the identical verdict.

Device-resident SPMD verify accept + drafter seed (the overlap body).

Runs after the SHARED verify forward (``hidden_full`` /
``per_token_logits``, computed once by
:func:`run_verify_step_spmd_worker`). It is the SPMD analogue of
:func:`arbi_serve.spec_decode.mtp_verify_async.run_verify_step_async`'s
device-resident core — but adapted to the rank-symmetric worker (mirror
rows, NO worker-bridge broadcast, SAME-tick commit):

  1. Gather the ``(K+1, B, V)`` verify logits and run
     :func:`arbi_serve.spec_decode.mtp_verify_async._async_verify_accept`
     — the SAME device op the sync path routes (``mtp_verify_greedy`` for
     greedy, the batched Leviathan-2023 rejection sampler for stochastic)
     — to get the DEVICE ``n_accepted`` ``(B,)`` and ``committed``
     ``(K+1, B)``. NO host pull, so the greedy argmax / rejection draw is
     byte-identical to the sync worker's ``verify_and_accept`` and rank-
     symmetric (identical logits + rank-agreed ``mtp_seed`` → identical
     result on every rank).
  2. Seed the next-step drafter chain from GPU-GATHERED inputs
     (``committed[n_accepted]`` / ``hidden[cu_q + n_accepted]``, read off
     the DEVICE ``n_accepted``), with the async path's device frontier-
     repair meta (the rejected draft slots are NOT freed yet, so the
     drafter's page-table-derived frontier cell would be wrong). So
     forward → accept → drafter enqueue back-to-back on the compute stream
     with NO host round-trip — killing the sync worker's forward→drafter
     bubble (the ``n_accepted`` ``.tolist()`` the sync path pays BETWEEN
     the forward and the drafter to build the drafter inputs).
  3. Conditionally roll back the GDN/Mamba recurrent slab on partial-
     accept rows (27B is a GDN hybrid), device-resident and RANK-SYMMETRIC
     — NO worker-bridge broadcast (the symmetry IS the broadcast; every
     rank derives the identical ``(slab_rows, n_accepted)`` from its
     mirror + plan and rolls its own shard). ``k_uniform`` folds full-
     accept rows into self-copies.
  4. ONE combined host pull of ``n_accepted`` / ``committed`` / drafts at
     the END. The per-tick C9 ``_spmd_agree_step_failed`` all_reduce
     already drains the compute stream every tick (``.item()`` on the
     reduced fail flag waits on the whole tick's GPU work), so this single
     pull is near-free — and every rank pulls the SAME device tensors, so
     the mirror commit (``accepted_per_row`` + ``mtp_next_drafts`` write)
     is byte-identical across ranks. The commit stays SAME-TICK (no
     optimistic advance / no cross-tick park) so the next tick's
     ``derive_verify`` reads a fully-corrected mirror — the deferred-park
     machinery the TP1 async path needs would give ~zero extra overlap
     here (the C9 drain caps the pipeline at one tick) while adding a
     rank-symmetric optimistic-advance divergence surface.

Returns ``(accepted_per_row, drafted_per_row)`` — the SAME contract as
:func:`run_verify_step_spmd_worker`, so the driver commits the mirror via
``commit_verify`` identically. Only reached for an ACTIVE uniform-K step
(the :func:`_spmd_overlap_active` gate); a cold / DFlash / flag-off step
stays on the synchronous inline worker.

Worker-rank MTP verify step under SPMD.

Runs the SAME rank-symmetric work rank 0's :func:`run_verify_step`
issues — the verify forward (row-parallel all-reduce + lm_head
all_gather) and the K-step drafter chain (per-step all-reduce +
all_gather) — so the TP collectives match rank 0 peer-for-peer and
the NCCL group never deadlocks. The worker has no scheduler / request
map / stream, so it builds Request-shaped views over the
:class:`arbi_serve.distributed.spmd.RankSlateMirror`
(:class:`_WorkerVerifyRequest`) and DROPS the rank-0-only side effects
(``post_token`` / scheduler commit / detok / stream).

Returns ``(accepted_per_row, drafted_per_row)``. ``accepted_per_row``
is the per-row accepted-token lists (``n_accepted + 1`` each) so the
caller can advance the mirror via ``commit_verify`` — identical to
rank 0's commit, since greedy accept over the SAME logits is
deterministic. ``drafted_per_row`` is the per-row effective K (``0``
for a cold row whose draft-slot allocation failed, ``step_k``
otherwise) — the rank-0 stream publishes it as ``mtp_proposed`` with
the SAME semantics as the driver's ``VerifyResult.num_drafted`` (cold
rows never reach ``commit_results`` on the driver, so they count zero
proposals there too). Both lists are computed identically on every
rank, so the accept curve is correct without a per-step host sync. The
drafter chain writes the next step's drafts back into the mirror (via
the shim's ``mtp_next_drafts`` setter) so the worker's NEXT
``derive_verify`` stays in lockstep.

Greedy OR stochastic. Both draws are rank-symmetric under SPMD:
``worker_bridge`` is ``None``, so ``advance_mtp_seed`` advances the
engine step counter LOCALLY in lockstep on every rank (no broadcast);
every rank runs THIS function identically and advances the counter the
same number of times in the same order, so ``mtp_seed`` is byte-
identical across ranks. The rejection-sampler accept uniforms + the
drafter-chain Gumbel both read that counter + fixed offsets, so the
accept/reject decisions and the sampled draft tokens are identical on
every rank — the same property the plain-decode
:func:`gumbel_exp_noise` path has, applied to the MTP sub-system.

The explicit history the verify penalties score against, or ``None``
for the generic ``prompt + output`` one.

Read DIRECTLY off the shim by ``sequence_history_parts``; carried here
so a row that opted out of the generic history is scored against the
one it asked for on the verify slate too.

The rank-local grammar matcher, or ``None`` for an unconstrained row.

Read by the chain's grammar processor on the verify slate
(:meth:`~arbi_serve.sampler.xgrammar_processor.XGrammarLogitsProcessor.apply_to_verify_slate`)
and by the sharded-verify gate. A live read off the mirror: the matcher
is stateful (advanced per committed token, advanced-then-rolled-back
across the slate's speculative positions), so the shim must hand out
the mirror's OWN object rather than a snapshot.

True when a grammar spec rode this row's ``AdmitRow``.

Set at admit time on every rank, so it is the rank-symmetric answer to
"is this row constrained" even before the lazily-built ``grammar_state``
exists. Cleared only on an explicit rank-identical compile degrade.

Rank-0 synchronous MTP verify pass (the reference verify entry point).

:func:`run_verify_step` builds the verify plan, runs the main-model forward
over the ``K + 1`` flat tokens, and runs the shared eager verify core
(:func:`_verify_accept_rollback_commit`) followed by the external-drafter
rollback and the next-step drafter seed.

Run the verify pass for an MTP-only sub-slate.

Per request, allocate K draft slots, run the main model forward
over K+1 input ids ``[last_committed, draft_1, ..., draft_K]``,
gather K+1 logits, verify, free rejected slots, commit accepted
tokens. The gather/accept/rollback/commit core is shared with the
rank-symmetric SPMD worker via :func:`_verify_accept_rollback_commit`.

``step_plan`` is the slate's ALREADY-BUILT plan, for a caller that
built one and then routed here. Building a second plan for the same
slate is not idempotent: :func:`build_verify_plan` takes each row's
tail slot through ``page_table.allocate_slots``, which EXTENDS the
row's page-table length, and nothing on the verify path gives a slot
back. A second build therefore leaves the row's KV length one ahead
of its logical length with an unwritten slot inside the window the
attention kernel reads as live. One build per slate, always.

The sync pass from its forward's outputs to the next-step drafter seed.

Runs under the caller's ``_mtp_block_m_scope`` + ``inference_mode``,
exactly as it did inline in :func:`run_verify_step`. Split from the
forward so the SAME tail serves a verify slate whose forward ran fused
with the step's prefill chunk (:mod:`arbi_serve.runtime.fused_mixed_step`):
every index into ``hidden_full`` / ``per_token_logits`` is the plan's own
``cu_q``, which starts at 0 in both layouts.

Where does the TARGET's argmax land in the DRAFTER's ranking?

A K-chain proposes exactly one token per slot: the drafter's rank-0 pick.
So the chain's acceptance ceiling at a slot is "the target's argmax IS the
drafter's rank-0 token", and a width-``m`` tree's ceiling at that same slot
is "the target's argmax is somewhere in the drafter's top ``m``". The gap
between those two numbers is the entire case for tree speculation, and it
is measurable before any of it is built.

That matters because the acceptance-vs-K sweep in
``docs/decode-stall-open-gaps.md`` (G7) shows thinking-mode acceptance is
FLAT from K=2 to K=5 — depth is exhausted. Whether *width* is also
exhausted is a different question with a different answer, and this is the
instrument that answers it.

Armed by ``ARBI_MTP_Q_RANK_PROBE=<path>``; OFF by default, and off means
one attribute load per verify step (the tensor work is behind the path
check). Diagnostic only.

Hooked on the two verify paths that carry a dense ``q``: the async
dispatch (:mod:`mtp_verify_async`, which is what a default TP1 boot
actually serves) and the gathered sync dispatch
(:mod:`mtp_verify_accept`). It is deliberately NOT hooked on
:mod:`mtp_verify_sharded`, which verifies over a top-``n_cand``
compaction rather than the full vocab — a rank histogram taken there
would be truncated by ``sharded_verify_n_cand`` and would read as better
drafter agreement than there is. Probe with ``ARBI_SHARDED_VERIFY=0``.

Reads the RAW verify logits, deliberately: temperature, ``top_p`` and
``top_k`` are all rank-preserving on the target side, so the argmax this
records is the same token the sampler would call rank 0 under any of the
served recipes.

Alongside the rank histogram it cross-tabulates rank against the DRAFTER's
own confidence at that slot — top-1 probability, and the top-1/top-2
margin. A confidence-gated tree spends a sibling only where the drafter is
unsure, so what decides whether such a gate can work is exactly whether low
confidence predicts a nonzero rank. A pooled rank histogram cannot answer
that; the cross-tab can, and it is the same D2H.

It also records whether the target's argmax was inside the drafter's
support at all. Rank is ``#{v : q[v] > q[target argmax]}``, so an argmax
the drafter filtered away scores ``rank = |support|`` — which ``top_p`` can
squeeze to 1-3 on a peaked slot, making a total MISS read as "rank 1" and
inflating every width above 1. Counting the in-support samples separates a
true rank from that pile without needing the ``top_k=20, top_p=1.0`` boot
that earlier runs used to dodge it.

Record that this slate verified GREEDY, and which predicate sent it.

A greedy-routed slate proposed a point mass, so there is no drafter
ranking to probe. Writing that fact is what keeps an empty histogram
from reading as "the drafter never ranked the target's argmax" when it
actually means "the drafter was never asked for a distribution".

Append one aggregate rank histogram for this verify step.

No-op when the flag is unset. A non-dense ``draft_probs`` (the greedy
drafter's point mass) carries no ranking to probe, so the step is
recorded as ``skipped`` rather than silently contributing K*B rank-0
hits — the histogram must never be able to read "the drafter is
perfect" because the drafter had no distribution.

Drop the bound index-key slab so a pool teardown can free it.

Every draft and fill re-binds before use, so holding it between
steps buys nothing and keeps a torn-down pool's slab reachable —
which the named pool's release seam refuses to reclaim.

Graph-safe batched rejection sampler for MTP verify + drafter.

Avoids a per-row Python loop in
:func:`arbi_serve.spec_decode.verify._verify_rejection_sampling` and
per-row ``torch.multinomial(generator=...)`` calls in
:meth:`MtpHead.forward(... sampling_params=...)`. Both hot paths
route through one batched implementation that:

  1. Computes accept/reject decisions for ``(B, K)`` slots in a single
     ``torch.lt`` over a ``(B, K)`` uniform tensor;
  2. Samples recovery + bonus tokens via Gumbel-max
     (``argmax(log p - log(-log U))``) over ``(B, V)`` per slot —
     batched across the slate;
  3. Reads its randomness from a persistent ``int64 (1,)`` seed buffer
     advanced once per engine step. The seed buffer is the only stateful
     handle: tests fix it, captured graphs replay against it without
     baking ambient RNG state into the recorded launches, and sleep-
     mode lifecycle owns its zero/restore.

Why Gumbel-max + a persistent seed instead of ``torch.multinomial``
====================================================================

``torch.multinomial`` is not graph-capture-safe under cudagraph
recording: it consumes ambient ``torch.Generator`` state which the
captured launch records once, so every replay returns the same draw.
Gumbel-max with explicit uniforms read from a persistent buffer is
graph-safe by construction — the captured kernels read whatever
``copy_()`` wrote into the buffer before replay, so each step gets
fresh draws.

Determinism. Two engine runs that copy_() the same seed-counter
sequence into ``mtp_seed`` produce bit-identical accept/recovery
decisions. This is the contract the existing
``test_rejection_sampling_deterministic_under_fixed_seed`` test relies
on, generalised: same seed-counter sequence → same outputs across runs
AND across captured-graph replays.

The verify-pass and the drafter-side categorical share the same seed
buffer with deterministic offsets (different leading dims of the
generated uniform tensor) so their draws are independent without
needing two seed counters.

Noise source per device
=======================

The accept-test uniforms are ALWAYS a tiny CPU-generator draw (device-
independent, bit-identical for the same seed), so accept decisions never
depend on the recovery noise source. The full-vocab recovery/bonus
Gumbel noise has two sources:

  * **CUDA (default)** — generated INSIDE the single-pass Triton kernel
    from counter-based Philox keyed on ``(seed, row, vocab_index)``
    (:func:`arbi_serve.sampler.gumbel_argmax_triton.residual_bonus_gumbel_from_seed`).
    No materialized ``(K+1, B, V)`` noise tensors, no full-vocab argmax
    passes, and byte-identical across SPMD ranks for the same seed. The
    kernel reads the seed from a device pointer — no host readback.
  * **CPU, and CUDA under ``ARBI_MTP_PHILOX_VERIFY=0``** — materialized
    seeded uniforms + ``-log(-log u)`` feeding the pure-tensor op
    (``mtp_sample_residual``).

Both draw from the SAME lossless recovery/bonus distributions
(Leviathan-2023 residual + bonus); only the noise realization per seed
differs between them.

Batched, deterministic Leviathan-2023 / Chen-2023 rejection sampler.

No state — the seed lives on the engine's ``mtp_seed`` buffer, not
on the sampler instance. This is constructed once per engine and
invoked per verify call.

The sampler handles the full per-slot slate in three batched ops:

  1. **Accept test.** ``u = U_seeded((B, K))``;
     ``accept[b, k] = (u[b, k] * q_full[b, k, x_b_k]) <= p_target[b, k, x_b_k]``.
     Equivalent to ``u <= min(1, p/q)`` after clamp; both forms
     pass the q=0 edge by the ``q == 0 → accept iff p > 0`` rule
     the legacy code implemented (when q is zero the multiplied
     RHS is zero too, so accept iff p > 0).
  2. **Recovery sample (per slot, batched).** Build the
     residual distribution ``residual[b, k] = (p[b, k] - q[b, k])+``,
     normalise per row, Gumbel-max sample. Done over the whole
     ``(B, K, V)`` block in one kernel.
  3. **Bonus sample (per row, batched).** Gumbel-max sample
     ``p_target[K]`` per row.

First-mismatch reduction. ``num_accepted = first_false(accept_mask)``;
``argmax(~accept_mask, dim=1)`` gives the first reject slot,
falling back to ``K`` when every slot accepts. Identical to
:func:`_verify_greedy`'s reduction.

The output is consumed by :func:`_verify_rejection_sampling`'s
thin orchestrator which turns the per-row ``num_accepted`` /
``recovery_tokens`` / ``bonus_tokens`` into the per-row
:class:`VerifyResult`. There is no per-row Python loop in the
sampler itself.

Validate the :meth:`sample_batched` operands and return shape ints.

Checks the rank of every operand, cross-consistency of ``(K, B, V)``
across ``main_logits`` / ``draft_tokens`` / ``draft_probs``, the
``sampling_params`` length, and the ``seed_buf`` dtype/numel. A
:class:`PointMassQ` ``draft_probs`` (the greedy-draft one-hot marker)
has no tensor shape to check — the implied ``q`` is defined by
``draft_tokens``, whose shape IS validated. Returns
``(K_plus_1, K, B, V)``. Pure validation — raises on a mismatch, has
no side effects.

Build the ``(K+1, B, V)`` per-slot target distribution.

Applies the per-row SamplingParams chain (temperature / top_k /
softmax / top_p / min_p) to every verify slot — OUTSIDE the
``mtp_sample_residual`` op boundary so the op sees a pure-tensor
function. Uniform slates take the single fused
:func:`_apply_sampling_params_uniform_block`; heterogeneous slates
loop per slot through :func:`_apply_sampling_params_batched`. When
``scratch`` is supplied the result lands in its persistent
``p_target`` slice (validated here, fail-loud on a mis-sized scratch)
instead of a fresh allocation. Byte-identical across the three
branches for the same inputs.

``main_logits_d`` may be fp32 or the lm_head's native dtype; the
chain widens the block it is about to process. Feeding it native
keeps the widened copy chunk-sized instead of slate-sized.

Draw the seeded accept uniforms + recovery/bonus Gumbel noise.

The materialized-noise path (CPU, and CUDA under
``ARBI_MTP_PHILOX_VERIFY=0``). Both draws come from ``seed_buf``
via deterministic offsets so the op core stays a pure function of
its tensors (graph-capture-safe). Returns
``(u_accept (K, B), gumbel_noise (K+1, B, V))``. With ``scratch``
(whose ``u_gumbel``/``gumbel_noise`` slices exist) the full-vocab
draws land in the persistent slices and the Gumbel transform runs
in place — byte-identical to the fresh-allocation path for the
same seed. A Philox-path pool carries no such slices
(``u_gumbel is None``); the fallback then allocates fresh.

Draw the ``(K, B)`` accept-test uniforms from ``seed_buf``.

Always the CPU-generator branch of :func:`_seeded_uniforms` (the
draw is tiny), so accept decisions are bit-identical for the same
seed on every device AND across the Philox / materialized-noise
recovery paths — the accept rule never depends on which noise
source drives the recovery draws.

The tiny CPU→device copy always rides a pinned non_blocking
staging buffer (bit-neutral) so it never stalls the host behind a
pending verify forward, unconditionally on every verify path
(sync, async, SPMD).

Re-derive the legacy :class:`RejectionSampleResult` public layout.

The op's compact ``(K+1, B)`` ``accepted_tokens`` carries the recovery
slots (``[:K]``) and the bonus slot (``[K]``); split them into the
legacy ``(B, K)`` recovery + ``(B,)`` bonus tensors and rebuild the
per-row ``accept_mask`` as ``arange(K) < num_accepted``. Pure tensor
reshapes on ``device``; no host transfer here.

Run the full rejection-sampling slate in batched ops.

Pre-conditions: ``main_logits`` first dim is ``K+1``;
``draft_tokens`` first dim is ``K``; ``draft_probs`` first
dim is ``K``; ``B`` matches across all three; every row has
``temperature > 0``.

``draft_probs`` may be the
:data:`~arbi_serve.spec_decode.drafter.POINT_MASS_Q` marker (the
default greedy-draft proposal): the implied ``q[k, b, v] = 1``
iff ``v == draft_tokens[k, b]``. The ops consume it INDEX-FORM —
accept ``u <= p_target[draft]`` (``u * 1.0`` is exact) and a
residual that zeroes exactly the drafted token — BIT-IDENTICAL
to densifying the one-hot, with zero ``(K, B, V)`` q traffic.

``seed_buf`` is the engine-owned ``int64 (1,)`` step seed
(lazily allocated, copy_()'d each step). Two calls with the
same seed value produce identical outputs — the determinism
contract.

``scratch`` (optional) is the engine's persistent
:class:`~arbi_serve.spec_decode.verify_buffers.SamplerScratch`
sliced to this step's ``(K, B)``. When supplied the vocab-scale
``p_target`` / ``u_gumbel`` / ``gumbel_noise`` are written into the
persistent slices instead of fresh heap tensors — the served path
then allocates ZERO new vocab-scale device memory per verify step
(the "memory locked, can't OOM" invariant). The accept decisions
and recovery / bonus draws are BYTE-IDENTICAL to the ``scratch=None``
path for the same seed — a pure allocation change. ``None`` keeps
the fresh-allocation path (unit tests / non-engine callers).

Exact-match verification against the target's keyed draw.

``t*[d, b] = argmax(p_target[d, b] + Gumbel(u(seed[d*B+b], ·)))``
is the token the watermarked K=1 sampler WOULD emit at depth
``d``'s context — computed for all depths in one pass by the
residual/bonus kernel with ``q == 0`` (a ``-1`` draft-token
tensor synthesizes an all-zero one-hot, so every row draws from
the pure filtered target). A drafted prefix is accepted exactly
as far as it matches ``t*``; the committed frontier token IS
``t*`` at the first mismatch (or the bonus depth), so
``accepted_tokens = t*`` satisfies the public layout contract
directly.

Fused single-pass rejection sample over the scaled verify logits.

The ``ARBI_MTP_FUSED_REJECTION`` path for a uniform ``temperature
> 0`` slate (the eligibility is checked by the caller via
:func:`_fused_rejection_ineligible_reason`). Pre-op work is the
per-row temperature divide plus — for a FILTERED slate — ONE mask
launch that top_k/top_p/min_p-masks the scaled logits to ``-inf``
in place AND writes the masked-row logsumexp
(:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
with ``min_p`` / ``lse_out``). That launch REPLACES the
``torch.logsumexp`` the no-filter path issues, so the filtered
fused path has the SAME launch count as the no-filter one: NO
softmax, NO separate mask pre-pass, NO dense ``(K+1, B, V)``
``p_target``. Masking is exactly the kernel the materialized chain
uses (:func:`_topk_topp_probs`), so the kept set is identical;
the in-kernel ``exp(logit - lse)`` differs from the materialized
``softmax`` only by fp32 reduction order — the drift class the
no-filter fused lever already shipped with.

Accept uniforms are drawn EXACTLY as the materialized path draws
them (the same CPU-generator branch of :func:`_seeded_uniforms`
keyed on ``seed_buf`` + ``OFFSET_ACCEPT_U``), so accept decisions
stay deterministic and cross-rank identical. CUDA takes the Philox
kernel (in-kernel noise, no materialized ``(K+1, B, V)`` tensor);
CPU / ``ARBI_MTP_PHILOX_VERIFY=0`` takes the materialized-noise
reference op (byte-identical to the ``p_target`` path fed the same
masked logits), with the filters applied by
:func:`_mask_scaled_logits_reference_` (CPU) or the same Triton
mask kernel (CUDA rollback).

Drafter-side categorical sampler methods for the rejection sampler.

Holds :class:`_DrafterSamplingMixin`, the drafter-token draw path
(:meth:`sample_drafter_token`, :meth:`_sample_drafter_token_device`,
:meth:`sample_drafter_token_tensors`) that
:class:`~arbi_serve.spec_decode.rejection_sampler.GraphSafeRejectionSampler`
inherits. The mixin is a base of ``GraphSafeRejectionSampler`` so every
``sampler.sample_drafter_token(...)`` call site keeps resolving unchanged;
``self.OFFSET_DRAFTER`` (and the other ``self.<attr>`` seams) resolve on the
concrete class via the MRO.

The ``drafter_fused_draw`` flag-truth counter lives HERE — with its only
``fire()`` site (:meth:`sample_drafter_token_tensors`) — so the counter's
GRAPH_RECORD firing window and its record remain co-located.

Draft a single ``(B,)`` token + return ``(B, V)`` distribution.

Used by :meth:`MtpHead.forward` and :meth:`MtpHead.forward`'s
sampling path. Returns ``(tokens, probs_full)``:

  - ``tokens``: ``(B,) int64`` — sampled per-row draft picks.
  - ``probs_full``: ``(B, V) fp32`` — the per-row q_full
    distribution the verify-pass needs for the residual
    (Leviathan-2023 Theorem 1).

Greedy rows (``temperature == 0``) collapse to argmax + a
point-mass on the picked token, matching the legacy contract.
Non-greedy rows go through the same chain
(``_apply_sampling_params_batched``) verify uses for ``p_target``,
keeping ``q_full`` and ``p_target`` on the same processed-logits
manifold (the ``p / q`` ratio is unbiased only under that
invariant — see :func:`_apply_sampling_params_to_logits`).

``slot_offset`` lets the caller request distinct draws per
autoregressive step in a K-step drafter chain (each chained
call passes a different small int so the K calls' uniforms
don't alias).

CUDA (default): the whole draw stays ON DEVICE — the
SamplingParams chain runs where the logits live (the same
Triton-masked :func:`_apply_sampling_params_batched` the verify
pass uses for ``p_target``) and the token comes from the
single-pass counter-based Gumbel-max kernel keyed on
``(seed, row, OFFSET_DRAFTER + 17 * slot_offset)`` — no
``(B, V)`` D2H of the drafter logits, no per-row Python loops,
no host readback of a device seed (capture-safe, and byte-
identical across SPMD/TP ranks for the same seed). CPU (and
``ARBI_MTP_PHILOX_VERIFY=0`` rollback) keeps the legacy
CPU-generator path. Both are exact categorical draws from the
SAME ``q_full``; the per-seed noise realization differs.

``wm_row_seeds`` selects the KEYED-WATERMARK coupled draw: the
slot token is ``argmax(log q + Gumbel(seed_row, v))`` with the
zero-lane per-row Philox layout — the SAME noise the keyed verify
pass adds to ``log p`` at this depth
(:meth:`~arbi_serve.sampler.watermark.WatermarkContextCache.mtp_draft_seed_row`),
so the draft agrees with the target's keyed choice whenever the
two filtered distributions do. Still an exact categorical draw
from ``q_full`` marginally over context hashes; the step-seed
lanes are unused.

Keyed-watermark coupled slot draw (see :meth:`sample_drafter_token`).

``q_full`` comes from the same per-row SamplingParams chain the
verify pass applies to ``p_target``; the token is
``argmax(log q + Gumbel(seed_row, v))`` with zero counter lanes —
bit-matching the noise the keyed verify adds to ``log p`` at this
depth. Greedy rows come back from the chain as a point mass, so
the keyed argmax returns their argmax deterministically.

CUDA runs the per-row-seed kernel
(:func:`~arbi_serve.sampler.gumbel_argmax_triton.gumbel_argmax_from_seeds`);
CPU mirrors it with the host Philox reference (tests / rollback —
never a served path).

On-device :meth:`sample_drafter_token` (CUDA + Philox path).

``q_full`` comes from the SAME per-row SamplingParams chain
(:func:`_apply_sampling_params_batched` — temperature → top_k →
softmax → top_p → min_p, Triton-masked on CUDA) the verify pass
applies to build ``p_target``, so the drafter distribution and
the target live on the same processed-logits manifold (the
``p / q`` acceptance ratio is unbiased only under that
invariant). Greedy rows (``temperature <= 0``) come back from
the chain as a point mass on their argmax; ``log(one-hot)`` is
``0`` at the mass and ``-inf`` elsewhere, so the Gumbel-max
below returns their argmax deterministically (the Philox
Gumbel spread is < 27, far below the clamp gap) — no separate
greedy branch.

The token is one single-pass counter-Gumbel draw over
``log(q_full)`` keyed on ``(seed, row,
OFFSET_DRAFTER + 17 * slot_offset)`` — the drafter's lane-2
namespace next to the verify tail's ``OFFSET_RECOVERY`` (see
:func:`~arbi_serve.sampler.gumbel_argmax_triton.gumbel_argmax_from_seed_ptr`).
A masked-out token has ``q = 0 → log q = -inf`` and can never
win, so the drawn support is exactly ``q_full``'s. The chain
guarantees every row keeps at least one positive-probability
token (top_k/top_p keep >= 1; min_p keeps the max), so the
kernel's degenerate all ``-inf`` row cannot occur.

No D2H of the logits, no host readback of the seed
(``_seed_to_device`` stages a host seed through the pinned
twin), no per-row Python loops. Deterministic in
``(seed, slot_offset)`` and byte-identical across SPMD/TP ranks
for rank-identical logits + the rank-agreed seed — the basis of
the broadcast-free TP true-stochastic drafter.

Fused draw (default, ``ARBI_DRAFTER_FUSED_DRAW`` on): the tail —
softmax, min_p, ``log``, and the counter-Gumbel draw + ``q``
write-back — collapses into ONE multi-SM op after the Qrita mask,
via the SAME tensor-driven path the captured chain uses (so the
live fallback and the captured replay draw byte-identically). The
drawn token + kept support are bit-identical to the unfused chain;
``q`` differs from :func:`_apply_sampling_params_batched` by a few
fp32 ULPs (softmax reduction order), lossless for the residual
sampler (Thm 1). The unfused rollback keeps ``q`` bit-identical to
``_apply_sampling_params_batched``.

Graph-CAPTURABLE :meth:`sample_drafter_token` (tensor-driven params).

The SAME fp32 upcast → per-row temperature divide → fused
top_k/top_p Triton mask → softmax → counter-Gumbel draw over
``log(q)`` keyed on ``(seed, row, OFFSET_DRAFTER + 17 *
slot_offset)`` the live chain runs. The difference is purely WHERE
the parameters come from: every knob is read from a device TENSOR
(:class:`DrafterSamplingTensors`) instead of Python
``SamplingParams`` values, so a captured graph replays with
whatever ``copy_()`` last wrote into the buffers — no values baked
at record time, no host reads, no fresh H2D uploads inside the
captured region.

Relation to :meth:`_sample_drafter_token_device`, which decides
whether a caller may substitute this method for it:

  - Fused draw ON (the default): that method IS a wrapper around
    this one — it encodes the slate via
    :meth:`DrafterSamplingTensors.from_params` and delegates. The
    two are then the same call, so the substitution is an identity
    for EVERY slate, mixed greedy and live ``min_p`` included.
  - Fused draw OFF: that method takes the eager
    :func:`_apply_sampling_params_batched` chain, which applies
    ``min_p`` and writes greedy rows as an explicit one-hot, while
    the unfused tail below does neither. The substitution is
    byte-identical there only for an all-stochastic,
    ``min_p == 0`` slate.

Contract deltas vs the live path (enforced by the ROUTING layer,
:meth:`MtpDriver._draft_captured`, never silently here):

  - ``min_p`` is applied ONLY on the fused-draw path
    (``ARBI_DRAFTER_FUSED_DRAW`` on, ``tensors.min_p`` supplied);
    on the unfused-tail rollback the caller must route ``min_p > 0``
    rows to the live chain.
  - Greedy rows (``temperature <= 0``) must be encoded by the
    caller as ``temp=1.0, top_k=1, top_p=1.0`` — a point mass on
    the row argmax via the mask chain (equal to the live path's
    explicit one-hot up to exact-fp logit ties, where the mass
    splits across the tied maxima).

Fused path (default): the whole tail — softmax, optional min_p,
``log``, and the counter-Gumbel draw + ``q`` write-back — collapses
into ONE multi-SM op after the (unchanged) Qrita mask, byte-matching
the unfused chain on the drawn TOKEN + kept SUPPORT (q within a few
fp32 ULPs; lossless for the residual sampler). See
:func:`~arbi_serve.sampler.gumbel_argmax_triton.fused_softmax_minp_gumbel_from_masked`.

``wm_row_seeds`` selects the KEYED-WATERMARK coupled draw: the token
is drawn against per-row Philox seeds with zero counter lanes — the
noise the keyed verify adds to ``log p`` at this depth — instead of
the step seed's ``(row, draw_offset)`` lanes. ONLY the noise source
changes: ``q`` is the same filtered distribution the unkeyed draw
reports, off the same mask. That is what lets the keyed draft ride
this tensor-driven path at all — before it, an armed watermark sent
the whole walk to the eager ``SamplingParams`` chain, which rebuilds
its per-row device tensors from Python lists at EVERY position.

Stateless helpers + value objects for the graph-safe rejection sampler.

Module-level, stateless helper functions and the two value-object
dataclasses (:class:`DrafterSamplingTensors`, :class:`RejectionSampleResult`)
used by
:class:`~arbi_serve.spec_decode.rejection_sampler.GraphSafeRejectionSampler`
and by the MTP verify / drafter paths. Every symbol here is re-exported
from ``arbi_serve.spec_decode.rejection_sampler`` so
``from ...rejection_sampler import <name>`` imports resolve.

True iff the counter-based (Philox) verify tail serves ``device``.

The single-pass Triton kernel needs CUDA + an importable ``triton``
and honours the ``ARBI_MTP_PHILOX_VERIFY`` rollback flag. Consulted by
both the sampler dispatch (:meth:`GraphSafeRejectionSampler.sample_batched`)
and the persistent-scratch builder
(:class:`~arbi_serve.spec_decode.verify_buffers.VerifyBuffers`) so the
materialized-noise scratch is allocated exactly when the fallback path
can run — one source of truth for the path selection.

True iff the fused drafter filter+draw tail serves ``device``.

The multi-SM fused tail
(:func:`~arbi_serve.sampler.gumbel_argmax_triton.fused_softmax_minp_gumbel_from_masked`)
needs CUDA + Triton and honours the ``ARBI_DRAFTER_FUSED_DRAW``
rollback flag. Consulted by the drafter draw dispatch (both the live
device path and the graph-capturable tensor path) AND the driver's
min_p routing gate (a ``min_p > 0`` slate rides the captured chain only
when the fused tail — which applies min_p — is active). One source of
truth for the path selection.

Candidate capacity the fanned top-C mask may use on ``device``; ``0`` = off.

The fanned mask (:func:`~arbi_serve.sampler.topc_mask.apply_topc_mask_`)
needs CUDA + Triton and honours the ``ARBI_TOPC_MASK`` rollback. One
source of truth for the path selection, asked by both spec-decode call
sites; the capacity it returns is also what the caller must feed
:func:`~arbi_serve.sampler.topc_mask.topc_mask_eligible`, because a
slate the capacity cannot represent keeps the pivot kernel.

Capacities are rounded DOWN to a power of two
(:func:`~arbi_serve.sampler.topc_mask.topc_capacity_floor`).

Does this slate present a row a candidate set cannot bound on its own?

``top_p`` active with ``top_k`` DISABLED is that row: the nucleus
renormalises over the whole vocabulary and can run past any capacity, so
the mask needs its escalation launch and a candidate set deep enough to
hold a nucleus rather than a served ``top_k``.

Asked of the SLATE and nothing else, because both of those are launch
structure: widening the candidates and adding a launch are what a
``top_k``-capped slate must not be charged for, and what a captured chain
cannot acquire by ``copy_()``.

Which mask kernel serves a slate, and at what capacity.

The dispatch :func:`mask_spec_logits_` performs, split out as a value so
a caller that masks the SAME slate many times — the DFlash block walk
masks it once per block position — resolves it once instead of
re-deciding from host tuples at every launch. The decision is a pure
function of the slate's host ``top_k``/``top_p`` values, the vocabulary,
the logits' dtype and the device's flag state; none of those move inside
a walk, so re-asking is pure host cost.

Resolving it as a value is also what makes the choice VISIBLE: the
fanned mask bakes its capacity as a launch constant, so a captured chain
must be able to name the capacity it recorded and refuse a slate that
capacity cannot represent (see
:meth:`~arbi_serve.runtime.capture.drafter.DrafterChainGraph.replay`).
A decision re-taken inside the loop has no name to record.

Apply the spec-decode slate's top_k/top_p/min_p mask in place.

ONE dispatch point for both speculative masks — the ``(K+1) * B`` verify
slate and the ``B = 1`` drafter slate — so the eligibility predicate is
asked in exactly one place and the two paths cannot drift apart.

Takes the fanned :func:`~arbi_serve.sampler.topc_mask.apply_topc_mask_`
when the slate fits its capacity, and the vendored pivot kernel
(:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`)
otherwise. The two produce the same masked tensor and the same
``lse_out`` bit for bit on LM-shaped rows, so this is a cost choice, not
a behaviour choice: the caller cannot observe which ran.

``top_ks`` / ``top_ps`` are the HOST-side values the device tensors were
built from — any iterable covering every distinct ``(top_k, top_p)``
pair in the slate is enough (a uniform slate passes one pair). They are
read only to answer :func:`~arbi_serve.sampler.topc_mask.topc_mask_eligible`;
the mask itself never reads a host value, so nothing here forces a
device readback.

Resolve-and-apply in one call, for a caller that masks a slate ONCE. A
caller that masks the same slate repeatedly holds a
:class:`SpecMaskPlan` instead — same rule, resolved once.

Return ``seed_buf``'s value as an ``int64 (1,)`` tensor on ``device``.

A CUDA-resident ``seed_buf`` is returned as-is (already keyable by the
kernel). A host buffer is staged pinned → device ``non_blocking`` into
a persistent per-device twin; the copy is ordered on the current
stream before the kernel launch that reads it. No ``.item()`` anywhere
— the value never round-trips through a device→host sync.

Device-tensor encoding of a slate's drafter sampling parameters.

The graph-capturable drafter sampling chain
(:meth:`GraphSafeRejectionSampler.sample_drafter_token_tensors`) reads
every per-row knob from these tensors instead of Python
``SamplingParams`` values, so a captured chain replays with fresh
parameters via ``copy_()`` — nothing is baked at record time.

Encoding conventions (mirroring the fused Triton kernel's disable
rules): ``top_k`` disable = ``V`` (any value ``>= V``); ``top_p``
disable = ``1.0``; ``min_p`` disable = ``0.0``; a GREEDY row
(``temperature <= 0``) is encoded as ``temp=1.0, top_k=1, top_p=1.0,
min_p=0.0`` (point mass on the row argmax).

``min_p`` (optional, ``None`` = the whole slate is ``min_p == 0``) is
only honoured by the fused-draw path
(:func:`~arbi_serve.sampler.gumbel_argmax_triton.fused_softmax_minp_gumbel_from_masked`,
``ARBI_DRAFTER_FUSED_DRAW`` on); the unfused softmax→Gumbel tail never
applied min_p, so the routing layer sent ``min_p > 0`` slates to the
eager chain. When the fused path is active the captured chain carries
min_p and those slates ride it too.

Apply per-row top_k + top_p and return ``(N, V)`` softmax probs.

On CUDA this reuses the main sampler's fused Triton kernel
(:func:`apply_top_k_top_p_triton`) — one launch, no full-vocab sort —
masking the logits then softmaxing. On CPU (no Triton) it falls back
to the legacy reference: top_k pre-softmax, softmax, then the
sort-based :func:`top_p_mask`. Both branches keep the SAME kept-token
set and renormalised distribution; the kernel just avoids the
``torch.sort`` over V the verify-step host hot spot was paying for.

``top_k``/``top_p`` are per-row tensors with the kernel's disable
convention: a row with ``top_k >= V`` or ``top_p >= 1`` is a no-op.
``top_ks``/``top_ps`` are the HOST values those tensors were built
from, forwarded to :func:`mask_spec_logits_` so a representable slate
can take the fanned mask. Empty (the default) keeps the pivot kernel —
this is the DEFAULT-BOOT verify mask, so a caller that cannot answer
the eligibility question must not be guessed at.

Softmax temperature-scaled logits, applying top_k/top_p first if set.

The shared tail of the per-row and uniform-block SamplingParams chains:
when neither threshold tensor is supplied this is a plain ``softmax``;
otherwise it routes through :func:`_topk_topp_probs`. Byte-identical to
inlining the same branch — a pure dedup of the two call sites.

Generate fresh uniforms ``(0, 1]`` shape ``shape`` deterministic in ``seed_buf``.

``seed_buf`` is an ``int64 (1,)`` tensor (CPU or device, doesn't
matter — we read its value once via ``.item()``). ``offset`` is a
small integer added to the seed before using it, so different
callers (verify accept-test vs verify residual sample vs drafter
sample) get independent draws even when sharing the same seed
buffer.

By default the implementation goes through a CPU :class:`torch.Generator`
so the drawn values are bit-identical on CPU and CUDA: the uniforms are
generated once on the CPU generator, then copied to ``device``. This
keeps the determinism contract device-independent — a CUDA verify slate
and a CPU equivalence test that copy_() the same seed counter draw the
exact same uniforms.

``device_rng=True`` opts a LARGE draw out of that cross-device
bit-identity into a per-device generator seeded from the SAME counter.
The CPU branch is unchanged (a CPU generator), so a CPU equivalence
test stays bit-identical; only CUDA inputs switch to a CUDA generator,
which skips the single-threaded CPU ``torch.rand`` over the full vocab
+ the H2D copy. The CUDA draw is still deterministic in the same
seed counter (so two CUDA runs with the same seed sequence match), and
the noise feeds only the residual-recovery / bonus Gumbel-max — whose
cross-device value-identity was already only float-reduction-order
equivalent, never an accept-rule guarantee. Used for the verify
recovery noise; the tiny accept-test uniforms keep the CPU path.

Why ``.item()`` is OK here. ``seed_buf`` is a small scalar that the
engine zeros / refills synchronously before each verify call; the
one-shot read does not block any kernel-launch pipeline. Pulling it
inside a captured graph would not work — but this helper is invoked
OUTSIDE the captured drafter chain (post-replay sampling) and
OUTSIDE the verify forward (the verify pass itself runs on the
main-stream after the forward completes).

``out`` is an optional pre-allocated tensor to write into — the
persistent-scratch path threads a slice of
:class:`VerifyBuffers`'s ``sc_u_gumbel`` here so the draw lands in the
frozen serving layout instead of a fresh heap allocation. It must
match ``shape`` / ``dtype`` / ``device``. ``torch.rand`` writes the
same generator stream into ``out`` as it would return fresh, so the
drawn values are byte-identical to the ``out=None`` path for the same
seed — a pure allocation change, not a numeric one.

Batched Gumbel-max sample over the last dim of ``probs``.

``argmax_v(log p_v - log(-log U_v))`` where ``U_v ~ Uniform(0, 1)``;
distributionally equivalent to ``multinomial(p, 1)`` but
graph-capture-safe (no ambient RNG dependency) and parallel across
the leading dims.

Degenerate-row handling. When a row's probability sums to zero
(``residual = (p - q)+`` collapsed for instance), the log is
``-inf`` everywhere and the argmax is implementation-defined. We
fall back to argmax over ``probs`` itself in that case (the
point-mass on the most-likely token), matching the legacy
``_sample_from_probs`` fallback.

Convert ``(B, V)`` logits into per-row probability vectors under
each request's sampling-params chain — temperature → top_k →
softmax → top_p → min_p.

Implementation note. ``temperature``, ``top_p``, ``min_p``, and
``top_k`` vary per request, so a fully vectorised path would need
per-row mask construction with bounded V. The Python loop over B
only issues kernel launches (one ``topk``/``where`` per row that
sets a finite top_k, one per-row ``top_p``/``min_p`` when the
thresholds differ); the heavy ``softmax`` is one batched ``(B, V)``
launch. Runs on whichever device ``logits`` lives — the verify
slate keeps it on the GPU. Greedy rows
(``temperature == 0``) collapse to a point mass on argmax; that
branch happens at the verify-pass level (greedy slates take
``_verify_greedy``), so this function REQUIRES every row to have
``temperature > 0``.

Returns ``(B, V)`` fp32 probabilities; each row sums to 1 (or 0
if the truncation collapsed every token, which the caller treats
as the degenerate fallback in :func:`_gumbel_max_sample`).

Mixed-temperature slates. When ANY row in an MTP slate is stochastic
the verify caller routes the WHOLE slate (greedy rows included)
through rejection sampling — the scheduler does not split a slate by
temperature. A greedy (``temperature == 0``) row is handled here as a
point mass on its argmax: rejection sampling against a one-hot target
accepts the draft iff it equals the argmax and otherwise recovers the
argmax — distributionally identical to the dedicated greedy verify
path. (Dividing the logits by 0 would produce NaNs, so the greedy
rows take a placeholder temperature of 1.0 for the chain and get
their one-hot target written back afterwards.)

True iff every row shares the same temperature/top_k/top_p/min_p.

The verify pass applies the SamplingParams chain to each of the K+1
flat verify slots with the SAME per-row ``sampling_params``. When the
rows are uniform (the c=1 single-request case, and any homogeneous
slate) the chain is identical across slots AND across rows, so the
whole ``(K+1, B, V)`` block can take ONE batched chain (one temps H2D,
one softmax, one masking pass) instead of ``K+1`` separate ``(B, V)``
calls — the per-slot loop was the verify-step host hot spot at c=1.

``None`` iff the slate can take the fused single-pass rejection sample,
else a short machine-readable reason for the fallback (logged once +
counted — never silent).

Eligible: a UNIFORM slate (:func:`_sampling_params_uniform`) at
``temperature > 0`` — INCLUDING active top_k / top_p / min_p filters.
A filtered slate is served fused by masking the scaled logits to
``-inf`` before the in-kernel softmax (the same
:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
kernel the materialized chain uses, with the min-p + LSE epilogue
folded into that same launch), so the production regime — Qwen
``generation_config`` backfill ``temp 1.0 / top_p 0.95 / top_k 20`` —
fires the fused path.

Ineligible (falls back to the materialized ``p_target`` path):

  * ``empty_slate`` — no rows (defensive).
  * ``heterogeneous_slate`` — per-row params differ; the fused kernel
    chain is a uniform-block chain (one scalar mask broadcast).
  * ``greedy_temperature`` — ``temperature <= 0``; an all-greedy
    uniform slate is routed to the dedicated greedy verify path
    upstream and never reaches here (defensive).

Reference (CPU / non-Triton) logit-space filter mask for the fused path.

Writes ``-inf`` at every token the uniform ``top_k → top_p → min_p``
chain drops, IN PLACE, and returns ``flat``. ``softmax(masked)`` then
equals the truncated-and-renormalised target distribution — the same
semantics the CUDA production chain has ALWAYS used for top_k/top_p
(:func:`_topk_topp_probs` masks logits with
:func:`~arbi_serve.sampler.topk_topp_triton.apply_top_k_top_p_triton`
and softmaxes), now shared by the fused-from-logits ops.

Kept-set semantics per stage (each matching the eager reference the
chain already ships, so a parity test can compare sets exactly):

  * top_k — keep ``logit >= k-th largest`` (ties at the threshold
    kept; :func:`_topk_topp_probs`'s CPU branch rule).
  * top_p — on ``softmax`` of the top_k-masked row, keep the smallest
    prefix of the descending sort whose cumulative mass reaches
    ``top_p`` (:func:`~arbi_serve.sampler.processors.top_p_mask`'s
    ``cumprobs - p_v < top_p`` rule, identical kept set).
  * min_p — keep ``logit >= max_logit + log(min_p)``
    (:func:`~arbi_serve.sampler.sampler._vectorized_min_p_mask`'s
    logit-space rewrite of ``prob >= min_p * max_prob``; the ratio is
    renormalisation-invariant so applying it on the masked row equals
    applying it on the renormalised probs).

This is the correctness REFERENCE + CPU/rollback path; the served CUDA
path applies the same chain inside the single Qrita mask launch.

Apply ONE uniform SamplingParams chain over a ``(S, B, V)`` block.

Fast path for :meth:`GraphSafeRejectionSampler.sample_batched` when
every row shares the same chain (checked by
:func:`_sampling_params_uniform`). Bit-identical to calling
:func:`_apply_sampling_params_batched` per ``(B, V)`` slot and
stacking — the per-row chain is the same scalar temperature / top_k /
top_p / min_p broadcast over the whole block — but issues the temps
H2D, the softmax, and the top_p/min_p masks ONCE over ``(S*B, V)``
instead of ``S`` times. Caller guarantees uniformity; this reads only
``sampling_params[0]``.

Greedy rows (``temperature <= 0``) are still overwritten with a point
mass on their argmax — a uniform slate is either all-greedy (handled
by the dedicated greedy verify path, never reaching here) or all the
same positive temperature, so the greedy overwrite is a no-op on the
fast path but kept for safety.

``out`` is an optional persistent ``(S, B, V)`` destination (a slice of
:class:`VerifyBuffers`'s ``sc_p_target``). When supplied the result is
written into it via ``copy_`` — byte-identical to the fresh-return
path, a pure allocation change. The chain's intermediates still use
transient storage (the allocator reuses it across steps), so only the
final co-resident ``p_target`` lands in the frozen layout.

Peak-bounded SLOT CHUNKING. The chain allocates several fresh ``(rows, V)``
fp32 intermediates (``scaled`` / softmax ``probs`` / top_p / min_p masks).
At high concurrency × K the whole ``(S*B, V)`` block is hundreds of MiB
and — co-resident with the DFlash draft transient + the verify ``lm_head``
output — can OOM the razor-thin post-grow free pool. When a persistent
``out`` dest exists AND the block exceeds
:data:`_VERIFY_SAMPLER_CHUNK_ENGAGE_ROWS`, the leading ``S`` axis is
processed in slices of ``~_VERIFY_SAMPLER_CHUNK_STEP_ROWS`` rows written
straight into ``out`` — capping the fresh intermediate at ``(step, V)``.
Every op in the chain is PER-ROW (temperature /
top_k / top_p / min_p broadcast independently over rows), so chunking is
BIT-IDENTICAL to the single-shot block. The gate (``out`` present + ``N >
engage``) leaves the c1/c2 fast path (``N <= engage``) byte-for-byte
unchanged — no host-loop regression there.

One uniform SamplingParams chain over a flat ``(rows, V)`` block.

The per-row core of :func:`_apply_sampling_params_uniform_block` — factored
so the whole-block and the peak-bounded row-chunked drivers share ONE chain
(they can never diverge). Returns a fresh ``(rows, V)`` probs tensor.

Output of :meth:`GraphSafeRejectionSampler.sample_batched`.

Shape contract (uniform K over the slate; verify-pass post-R1
invariant):
  - ``accept_mask``: ``(B, K)`` bool — ``True`` iff the ``k``-th
    draft accepted on row ``b``.
  - ``num_accepted``: ``(B,)`` int64 — first-mismatch index per
    row, in ``[0, K]``. ``K`` means full accept.
  - ``recovery_tokens``: ``(B, K)`` int64 — sample from
    ``max(0, p_target[k] - q_full[k])+`` for each ``(b, k)``. The
    caller picks ``recovery_tokens[b, num_accepted[b]]`` when the
    row partial-rejected, else ignored.
  - ``bonus_tokens``: ``(B,)`` int64 — sample from ``p_target[K]``
    per row. Used only when the row full-accepted.

Encode ``sampling_params`` into the per-row host value lists.

Returns ``(temps, top_k, top_p, min_p)`` under the conventions
above; the caller writes them into (pinned) staging tensors. Greedy
rows come back as ``(1.0, 1, 1.0, 0.0)``.

``ARBI_DRAFTER_LOGIT_TEMP`` scales a stochastic row's temperature
here, which is the one place every drafting path encodes its knobs.
It moves the proposal only: the accept rule redraws a rejected slot
from the target, so the emitted distribution is unchanged and only
the accept rate moves. A greedy row is left alone — its ``top_k=1``
makes the draw an argmax, which no temperature scale moves.

Build a slate's drafter param tensors from Python ``SamplingParams``.

The single construction seam for the LIVE (uncaptured) fused draw:
:meth:`~arbi_serve.spec_decode.rejection_sampler_drafter._DrafterSamplingMixin._sample_drafter_token_device`
builds one per call, and the DFlash block walk builds one per BLOCK
and reuses it across the block's ``S`` positions (only ``slot_offset``
varies between them). Both must encode identically or the walk's ``q``
would drift from the single-slot draw's, so the encoding lives here
rather than at either call site.

``seed_dev`` is supplied by the caller (via :func:`_seed_to_device`)
because the walk stages the step seed once for the whole block.

SpecMode IntEnum + engine-side strategy interface for speculative decoding.

The :class:`SpecMode` IntEnum is the dispatch key. The enum carries
`@property` predicates that let the run loop branch without
isinstance / class hierarchy:

    plan.spec_mode.is_none          # the K=1 path
    plan.spec_mode.is_mtp           # the spec-decode verify path

A new strategy is one new enum value plus one strategy class; not
a new class hierarchy.

The :class:`SpecDecodeStrategy` Protocol stays as the run-loop seam.
Concrete implementations:

  - :class:`NoSpecDecode` — default. ``SpecMode.NONE``. Dispatches to
    the legacy K=1 forward in :mod:`arbi_serve.engine.run_step`.
  - :class:`MtpStrategy` — the spec-decode verify path. ``SpecMode.MTP``.
    Lives in :mod:`arbi_serve.spec_decode.mtp`. The bundled MTP head,
    the external draft model
    (:class:`~arbi_serve.spec_decode.external_drafter.ExternalModelDrafter`),
    and DFlash all run under this one strategy — they vary the
    *drafter* (chosen in
    :func:`~arbi_serve.engine.mtp_attach.build_mtp_driver`), not the
    verify logic.

A second verify strategy (EAGLE tree-verify, n-gram) is one new
``SpecMode`` value plus one strategy class.

The strategy interface is intentionally thin (``run_step`` +
``warmup`` + ``stats_for``).

Speculative-decoding mode dispatch key.

An IntEnum with ``@property`` predicates the run loop branches on,
instead of a class hierarchy with isinstance checks. Adding a new
mode is one new enum value + one strategy class; the dispatch
stays flat.

The bundled MTP head, the external draft model, and DFlash all run
under ``MTP`` — they differ by *drafter*
(:func:`~arbi_serve.engine.mtp_attach.build_mtp_driver` picks one),
not by verify strategy, so they share one mode.

Stable integer ordering (wire-compatible with metric labels —
don't renumber existing values):

  0 — NONE — no spec decode (K=1 path)
  1 — MTP  — bundled-head / external / DFlash drafters, shared verify

Materialize + commit any in-flight deferred spec-decode step.

The non-speculative output path calls this before its own commit
so per-request token order is preserved across a spec→K=1
transition. A strategy that defers no work makes it a no-op, so
the K=1 path drains polymorphically instead of reaching for an
MTP-specific free function.

Run the slate through the engine's legacy K=1 (non-speculative) forward.

Routes through the event-loop-friendly :func:`step_async`, which
offloads the blocking GPU forward to a thread executor so the API
event loop stays responsive (interactive admission near idle
under batch load). ``step_async`` falls back to the synchronous
inline path when offloading is disabled.

Wiring the lossless multi-candidate accept into the served verify step.

The algebra is :func:`arbi_serve.spec_decode.tree_spec.accept_stochastic`
and its device expression is ``arbi_serve::mtp_verify_stochastic_tree``.
This is the seam between them and the engine: it resolves the slate's
target distribution through the SAME per-row SamplingParams chain the
chain's rejection sampler uses, draws the step's randomness from the
engine-owned seed, and hands the op the cached geometry.

The drafter's proposal is a POINT MASS per candidate — the tree drafter
offers each node's token deterministically. So this path needs no
``q`` at all: the accept rule is inverse-CDF over the candidate set of
the TARGET's own distribution, and a slate whose drafts carry no
``q_full`` yet is as serviceable as one that does. That is why there is
no ``have_probs`` gate here, unlike the chain.

Device-resident lossless accept for a SAMPLED tree slate.

Returns ``(n_accepted (B,), committed (depth+1, B), path_rows
(depth+1, B))`` — the greedy tree op's contract, every tensor on the
verify device, so the async path's optimistic advance and deferred
commit stay temperature-blind.

Nothing here reads a tensor value on the host. The one host read is
the step seed's scalar, which the chain's sampler already performs at
the same point in the step (:func:`_seeded_uniforms`) and which is
written synchronously before the verify call rather than produced by
it.

At ``widths=(1,) * K`` every operand reduces to the chain's: the
accept uniforms are the same ``(K, B)`` draw at the same offset, the
target rows are the same ``p_target``, and the frontier draw reads
the same Gumbel row the chain's recovery / bonus draw reads. That is
what makes the reduction checkable rather than merely plausible.

What a tree needs from the rest of the stack, checked ONCE at boot.

Every condition here has the same failure mode if it is not checked: the
tree serves, acceptance looks healthy, tok/s looks healthy, and the
tokens are wrong. None of them raises on its own. So they are checked at
boot, by name, with the flag that caused each one — never by falling
back to a chain, which would leave a boot claiming ``ARBI_MTP_TREE`` and
serving something else.

The load-bearing one is the attention mask. A tree's siblings are
alternatives for ONE position, so a node that can see its sibling is
conditioning on a token outside its own hypothesis; the verify then
accepts a token the model never predicted from that prefix. Both verify
routes on this stack bake a linear per-row causal bound, which IS a
chain, so the mask is a kernel capability and not something arbi-serve
can arrange for itself.

Whether the installed tkv's bypass ``attend`` forwards a tree mask.

Probed off the signature rather than a version pin: the parameter is
the capability, and a version number is a proxy for it that goes
stale in both directions. ``lru_cache`` because the answer is fixed
for a process and this is read on the verify hot path.

False on any tkv where the import fails at all — a stack that cannot
even name the attend cannot be relied on to mask.

Whether the installed tkv's codec ``TKVCore`` forwards a tree mask.

The compressed-codec route reaches the kernel through
``TKVCore.forward`` and the bf16 bypass through
``BypassBf16Attend.attend``. They are two call sites into the same
Turbo prefill mainloop, so a mask threaded to one does not travel to the other
and each needs its own probe.

Probed off the signature for the same reason as the bypass one: the
parameter is the capability, and a version number is a proxy for it
that goes stale in both directions.

Ask THIS route's probe whether the installed tkv forwards the mask.

Resolved from the backend name rather than a table of callables so the
probe is looked up when the gate runs, not when this module is
imported — the two probes are the substitutable seam a test drives.

Refuse at boot on any config a tree cannot be served correctly on.

No-op with ``ARBI_MTP_TREE`` unset — this runs on every boot.

A width-1 tree is the chain written as geometry (its accept walk is
token-for-token the chain's), so it is exempt from the mask
requirement: with one candidate per depth no node has a sibling to
see. It is the lane that proves the tree machinery reproduces the
chain, and refusing it would remove the only equivalence check
available without the kernel.

Refuse a tree on a drafter that proposes a LINEAR run of tokens.

Only the bundled-head chain drafter expands a :class:`TreeSpec` into
one token per node (``_MtpTreeDraftMixin``, which declares
``expands_tree_nodes``). Every other speculation source on this stack
— DFlash, DSpark, an external draft model — proposes ``k`` tokens in
sequence and has never been handed the geometry.

The tree lives on the VERIFY side of a seam the drafter does not
cross, so a boot naming both takes neither path cleanly. The step's
draft width comes from ``tree_step_k``, so the drafter is asked for
``num_nodes`` tokens and answers with a straight run of them; the
verify plan then accepts that run on a length check alone, positions
token ``i`` at ``depths[i]`` — where siblings SHARE a position — under
the ancestors-only mask, and the accept walk compares each node
against its PARENT's row. A chain gets verified as a tree: wrong
positions, wrong mask, wrong parent.

Which of those symptoms surfaces depends on an arithmetic coincidence
between the drafter's ``k`` and the tree's node count, which is not a
property anyone chose. Refuse on the capability instead, at boot,
while the config is still a decision.

Read off a declared attribute rather than an ``isinstance``, for the
same reason the mask is read off a signature: the capability is the
thing, and a class check would admit a future subclass that dropped
the mixin and refuse a future drafter that grew it.

Refuse unless THIS boot's active paged-KV route can carry the mask.

Which route is live decides which signature is the capability. The
bf16 bypass passes the mask from ``TkvBypassAttnOp`` into
``BypassBf16Attend.attend``; the compressed codec passes it from
``TkvAttnOp`` into ``TKVCore.forward``. Probing the wrong one reads a
class this boot never calls, which is an answer about some other
stack.

A paged-KV backend outside :data:`_MASKED_PAGED_ROUTES` has no seam
to pass a mask through at all, so it is refused BY NAME: a route that
would silently drop the mask must stay a boot failure, since a block
row that still sees a sibling produces finite, plausible logits and
nothing downstream raises.

Read off ``active_backends`` and not ``cfg.attention_backends``: the
latter is the REGISTRY of what this boot can hot-swap to, so a boot
that lists both routes would be judged on whichever is named first
rather than on the one it serves. A hot-swap to another route is the
swap's problem to refuse, not this gate's — it happens long after
boot. No paged-KV backend at all means no attention route to mask.

Refuse a sibling-bearing tree on a recurrent config it is not wired for.

The GDN half IS wired: ``_GDNFLAVerifyTreeMixin`` runs the verify
block through the shipped ``fused_sigmoid_gating_delta_rule_update``
over a path-packed arrangement, with the conv window gathered along
each node's own ancestor chain, and reconciles the accepted path at
accept time out of node-order staging. What stays refused is
everything that wiring does NOT cover, and each of these is silent
rather than loud if it is allowed through:

  * **a non-GDN recurrent pool.** ``Mamba2Block`` and the DSv4
    compressor both replay a PREFIX of block order, which for a tree
    names a different set of tokens than the accepted path. Their
    rollbacks refuse a ``path_rows`` by name; this refuses earlier,
    at boot, where there is still a decision to make.
  * **recompute-mode rollback.** The tree forward's packed launch has
    one row per leaf against one slab row, so it commits no state at
    all — the accept-time masked replay is what makes it right, and
    recompute mode has no replay buffers to gather the path out of.
A width-1 tree needs none of it: its accepted path is ``0, 1, ..,
n-1`` at block rows ``1..n``, a prefix, and no node has a sibling to
absorb. It stays on the chain scan and remains the equivalence lane.

Refuse a tree whose verify block would land on the split-K kernel.

The ancestors-only mask exists on the Turbo prefill verify route only. The
split-K register kernel takes no mask argument at all — so the mask
kwarg the backend builds is simply not consulted and every sibling
becomes visible to every later block row. Nothing raises and
acceptance looks healthy, which is why this is a boot refusal.

Three ways a layer declines the Turbo prefill verify route, and all three
are checked here rather than at the layer, because the layer only finds
out mid-step:

  * ``TKV_MTP_PREFILL_SPLIT=0`` — the global rollback switch;
  * a SLIDING-WINDOW layer — split-K serves a local verify correctly
    today and the Turbo prefill kernel is not validated for it;
  * an ATTENTION-SINK layer — arbi-serve's sink-carrying subclass owns
    ``mtp_verify``, so ``prefill_verify_enabled`` keeps its own route
    rather than have a per-head sink silently counted once per split.

Retire a verify step's draft slots when the accept was SCATTERED.

``free_draft_slots`` is a pure suffix trim, and for a chain that is
exactly right: a chain accepts draft slots ``0..n-1``, a prefix, so
dropping the tail leaves the accepted KV where the next step expects it.

A tree's accepted path is not a prefix. Accepting nodes ``1, 4, 11`` of
14 and trimming to length 3 keeps slots ``0, 1, 2`` — three tokens the
model never committed. Nothing raises: the next step reads plausible KV
for a prefix that does not exist, acceptance stays healthy, and the
output is quietly wrong. That is the failure mode this module exists to
remove, and the only end-to-end gate that catches it is byte-identity
against non-spec greedy decode.

Both verify paths retire their slots through :func:`retire_draft_slots`
so the compaction cannot be present on one and missing on the other.

Compact the accepted path down, then trim the rejected tail.

``path[d]`` is the draft-slot ORDINAL accepted at depth ``d``; it has
``n_accepted`` entries. ``None`` — or a path that is already
``range(n_accepted)`` — is the chain, and reduces to the bare trim
this replaces, byte-for-byte.

Order matters and is not interchangeable: the moves are resolved
against the RoPE is not
re-applied — a tree's depth-``d`` node was written at absolute
position ``P + d``, which is the position its destination ordinal
occupies, so the bytes are already correct and rotating them again
would corrupt them.

Tree speculation: geometry, candidate expansion, and acceptance.

A K-chain proposes ONE token per slot — the drafter's rank-0 pick — and
throws away everything the drafter ranked 2nd or 4th. Measured on the 27B
(``ARBI_MTP_Q_RANK_PROBE``, see G9 in ``docs/decode-stall-open-gaps.md``):
the target's argmax is the drafter's top-1 only 63.2% of the time in
thinking, but is inside its top-4 **86.2%** of the time. A tree spends the
same verify forward checking several of those candidates instead of one.

Why that can pay: at B=1 the verify forward is memory-bound — its cost is
dominated by reading 27B of weights out of VRAM, and a chain spends that
read checking K+1 positions. Checking N+1 costs barely more of the term
that dominates. It is NOT free, and the shape a tree may take here is
constrained by what the attention kernels can express, which is narrower
than this module's geometry: see ``docs/tree-speculation-feasibility.md``
for the constraint, the cost budget, and the rank table above.

This module is the part of tree speculation that owns no memory and calls
no kernel: the shape, the ancestor relation, the candidate expansion, and
the accept walk. It is deliberately separable so the geometry and the
accept algebra can be tested exhaustively on CPU before any attention
mask or captured graph depends on them.

The chain is a tree of width 1. ``TreeSpec.chain(K)`` produces geometry
whose accept walk is token-for-token identical to the chain verify, which
is the property ``tests/test_tree_spec.py`` pins — a tree implementation
that cannot reproduce the chain it replaces is not trustworthy.

A static draft tree: per-depth branching, in one of two topologies.

``tail is None`` — the PRODUCT tree. ``widths[d]`` is how many
children EVERY node at depth ``d`` spawns, so the tree is complete.

``tail = t`` — the SPINE tree. ``widths[d]`` applies to the argmax
path only; a node that has already deviated from it spawns ONE child
while it is within ``t`` depths of its branch point, and none after.
The spine keeps a sibling at every depth its width asks for, and the
subtrees hanging off those siblings are truncated chains rather than
replicas of the whole remaining tree.

Which matters because a product tree REPLICATES every later width
across every earlier sibling, and those replicas are the worst rows
in the tree: they are consulted only on a step that has already
deviated. Measured on the 27B q-rank joint,

Either way the shape is a compile-time constant, which is what lets a
captured graph exist for it at all. A ragged / confidence-pruned tree
(EAGLE-2 style) is a later step and needs a shape bucket per
topology; this deliberately does not model one.

Node ids are assigned breadth-first, so all of depth ``d`` precedes
depth ``d+1`` and a node's children are contiguous. Both facts are
relied on by the mask build and by the accept walk. The spine node at
each depth carries the LOWEST id there — it is child 0 of the lowest
id at the depth above — which is what lets the drafter's row 0 stay
the spine row.

``ARBI_MTP_TREE`` resolved to geometry, or None for the K-chain.

OFF costs one attribute load: the string is only parsed when it is
non-empty, so an unset flag never builds a :class:`TreeSpec` and never
touches the node tables.

A step's speculation width: the node count, or the chain's ``K``.

Every consumer that sizes something per step — the drafter's bucket,
the verify plan's ``step_k``, the captured verify shape ``S = K + 1``,
the capture preflight's ``K`` bound — has to agree on ONE number, and
they agree by resolving it here rather than each deriving a value that
happens to match. They came from ``driver.max_k`` before, which is the
chain's depth and is not the tree's width.

A drafter still emitting a chain's worth of drafts fails the plan's
``len(cached) >= step_k`` and the step collapses to K=0 plain decode:
slow, correct, and visible as ``mtp_tree_verify`` staying at zero.

OFF costs one attribute load.

Rows the per-token MTP snapshot staging holds for ONE verify step.

A TREE stages one row per BLOCK row, not per chain step: its accepted
path is gathered out of node order at accept time, so the staging has
to hold the whole block. Plus one for the committed token the step
starts from.

Every route that ends with a drafter attached sizes this buffer, and
the buffer is engine-lifetime storage the partial-accept rollback
replays out of — so a route sizing it to the chain's depth while the
verify side lays out a tree's nodes under-allocates it. Resolving the
count here is what keeps a route's ORDERING relative to the boot
refusals out of the arithmetic: a refusal that happens to run before
the sizing is not what makes the sizing right.

The chain's own ``K + 1`` with no tree configured.

Widest verify slate ONE step can present, in draft rows.

Everything sized per verify row — the flat batch, the draft tensor,
the rejection sampler's vocab-scale scratch, the logits gather, and
the budget line that reserves for them — has to agree on ONE number.
A tree proposes one draft per NODE, so a route keyed off the chain's
depth is short by the node count under ``ARBI_MTP_TREE``.

``max`` of the two rather than the tree's count alone: a boot may
still serve chain rows whose per-request ``mtp_k`` exceeds a shallow
tree's node count, and a buffer sized below either shape is the one
that cannot be recovered from at step time.

``chain_k <= 0`` is "no speculative step this boot", and it stays
itself: ``tree_step_k`` answers with the node count whatever it is
handed, so widening a K=0 probe to a tree's block would reserve for a
verify step that does not run — the over-book that makes a depth
ladder price K=0 as if it speculated.

The chain's own ``K`` with no tree configured.

Attention rows ONE request's drafter runs at its widest depth.

A chain step runs one query row per request. A tree's depth ``d`` runs
one per depth-``d-1`` node THAT HAS A CHILD, so its widest depth is
the widest :attr:`TreeSpec.draft_rows` entry — the deepest level never
expands, and a spine tree's truncated tails stop expanding earlier
still. The drafter's page-metadata buffers are sized once at boot for
a fixed row count, so that factor has to be known there rather than
discovered when a slate arrives.

This is the axis a tree shares with the BATCH: the drafter runs
``B * this`` attention rows, and the exl3 GEMM behind them is
``TILESIZE_M=16`` strip-serial, so the product crossing 16 buys a
second full weight decode.

``1`` with no tree and for a width-1 tree, which is the chain's own
number and leaves a chain boot's buffers exactly where they were.

``(parent_row, depth_of)`` for the device greedy tree accept op.

``parent_row[i]`` is node ``i``'s parent's ROW in the verify block —
0 for a depth-0 node, since the block's row 0 is the committed token
every hypothesis descends from, and ``p + 1`` otherwise. That +1 is
the whole reason this lives here rather than at the call site: a
tree's accept compares node ``i``'s proposal against the argmax at
its PARENT's row, and off-by-one there is a silent accept against the
wrong distribution.

Cached per ``(widths, device)`` alongside the mask, for the same
reason: a fresh allocation per verify step would move a pointer a
captured graph baked, and these are pure functions of a
process-global flag.

``(child_nodes, child_valid)`` — the candidate slate at every block row.

``child_nodes[r, c]`` is the node id of block row ``r``'s ``c``-th
child and ``child_valid[r, c]`` says whether that slot is a real
child. Block row 0 is the committed token, whose children are the
depth-0 nodes; block row ``i + 1`` is node ``i``.

The stochastic accept needs the candidate SET at the row it currently
stands on, and that row is data-dependent — a different branch per
batch row, resolved by an index_select rather than by Python. The
greedy walk needs the inverse relation (:func:`verify_block_tables`,
parent-of-node) because it tests each node against its parent's
argmax; the stochastic walk tests one uniform against the cumulative
mass of a whole sibling set, so it addresses children-of-row.

Columns are padded to ``max(widths)`` so the tables are one fixed
shape for the whole tree rather than one per depth. Padded slots
carry node id 0 — a real index, so the gather that reads them stays
in bounds — and are excluded by ``child_valid``; a LEAF row is all
padding, which is what makes its accept slate empty and its draw the
unmodified bonus.

Cached per ``(widths, device)`` for the same reason as the mask and
the block tables: a fresh allocation per verify step would move a
pointer a captured graph baked, and these are pure functions of a
process-global flag.

The verify-block ancestor mask for the active tree, or None.

Cached per ``(widths, device)``: the mask is a pure function of the
geometry, which is a process-global flag, so it is built once and the
same tensor is handed to every layer of every step. That stability is
also what makes it safe to reference from inside a captured graph —
a fresh allocation per step would move the pointer the graph baked.

``{"ancestor_mask": ...}`` for a tree verify block, else empty.

The mask says which block rows a row may attend to. Without it a
kernel's linear per-row causal bound lets block row ``i`` see rows
``0..i`` — which for a chain IS the ancestor set and for a tree
includes SIBLINGS, alternatives for the same position rather than
predecessors. A node conditioned on its sibling is conditioned on a
token outside its own hypothesis, and the verify then accepts a token
the model never predicted there.

One definition for every attention route that can carry the mask, so
the four properties below cannot drift apart between them.

Gated on the call's SHAPE being the tree's verify block, not on the
flag alone and not on ``mtp_block_m`` alone. Decode and prefill come
through the same op, and ``mtp_block_m`` is a per-op attribute that
outlives the verify step that set it — so a prefill landing between
two verify steps would read a stale width and be handed a mask for
rows it does not have. The ``batch_size * block_m == num_tokens``
test is the same predicate the attend itself routes on.

Empty dict, not ``ancestor_mask=None``, so a tkv that does not take
the parameter is not handed it — boot already refused a
sibling-bearing tree on such a stack
(:func:`~arbi_serve.spec_decode.tree_boot.assert_tree_supported`),
and the width-1 tree it does admit has no sibling to mask.
``attend_takes_mask`` is that per-route capability probe's answer.

The tensor is the cached, address-stable one from
:func:`active_tree_mask`: a fresh allocation per step would move a
pointer a captured graph baked.

Parse ``ARBI_MTP_TREE`` — ``"4x2x2"``, ``"spine:2x2x2x1x1:tail2"``,
or empty/``"off"`` for None.

A bare width list is the PRODUCT tree and means exactly what it has
always meant — ``2x2x1x1`` is the same 14 nodes it was before the
spine form existed, which a measured A/B depends on. The ``spine:``
prefix selects the other topology and the trailing ``tailN`` says how
many depths a deviated path runs past its branch point.

Fails LOUD on a malformed value rather than silently serving a chain:
a typo'd tree spec that quietly disables the feature is exactly the
phantom-config failure ``flag_truth`` exists to prevent. That includes
a ``spine:`` with no tail — a spine tree without one is a shape whose
node count differs from what the operator wrote, so it is refused
rather than defaulted.

Flatten per-parent top-k picks into ``(N,)`` node token ids.

``node_topk_ids[p, c]`` is the drafter's rank-``c`` token for parent
slot ``p``, where slot 0 is the ROOT and slot ``i + 1`` is node ``i``.
Only the first ``widths[d + 1]`` columns of a depth-``d`` node are
read, so one ``topk`` at the widest width feeds every level.

Greedy tree accept. Returns ``(committed_ids, accepted_node_path)``.

``argmax_at[0]`` is the target's prediction at the committed token
(the root); ``argmax_at[i + 1]`` is its prediction at node ``i``.

Walks from the root: among the children of the current node, accept
the one whose proposed token equals the target's argmax there. On a
miss — or at the last depth — the target's own argmax is committed as
the recovery / bonus token, so a step ALWAYS commits at least one
token and can never stall.

At ``widths=(1,) * K`` this is token-for-token the chain rule: the
single child either equals the argmax (accept, descend) or does not
(commit the argmax and stop).

The distribution left after every candidate at a slot is rejected.

``slot_probs`` with all of ``candidate_ids`` zeroed and the remainder
renormalized. This is the tree analogue of the chain's
``(p - q).clamp_min(0)`` residual: with a point-mass proposal the
chain zeroes the ONE token it offered, and a tree slot zeroes the
``W`` it offered.

A slot whose candidates carry the entire mass leaves nothing to
renormalize; the caller is handed the unmodified ``slot_probs`` back
so a draw from it still terminates on a real token. That branch is
unreachable whenever some candidate was accepted, which is exactly
the case where the total is 1.

Lossless tree accept. Returns ``(committed_ids, accepted_node_path)``.

THE RULE. At a slot with candidates ``x_1..x_W`` (the drafter's top-W,
a POINT-MASS proposal per candidate — the drafter offers each token
deterministically, it does not sample it), accept ``x_i`` iff::

    cum[i - 1] < u <= cum[i],   cum[i] = sum(p(x_j) for j <= i)

for a SINGLE uniform ``u`` shared by all candidates at that slot, and
otherwise commit a draw from :func:`sibling_residual`.

WHY THAT IS LOSSLESS. The rule is inverse-CDF sampling restricted to
the candidate set, so it lands on ``x_i`` with probability exactly
``p(x_i)``, and reaches the residual with probability exactly
``1 - sum(p(x_j))`` — which is the residual's normalizer, so a token
outside the candidate set is emitted with probability exactly
``p(y)``. Every token therefore keeps its target probability and the
step's output distribution IS the target's. Induction over depth
extends that to the whole tree: the slot reached at depth ``d`` is
conditioned on a prefix that was itself emitted from the target.

ONE uniform, not ``W``, is what makes it lossless. Testing each
candidate against its own independent coin accepts ``x_1`` and
``x_2`` with probability ``p(x_1) * p(x_2)`` of BOTH firing, and the
tie-break between them is not a distribution anyone chose — that is
the silent way a multi-candidate verify stops being lossless.

DUPLICATE candidates are skipped rather than refused: a repeated
token adds no new mass, so its slice of ``u`` is empty and it can
never be accepted. Without the skip its probability would be counted
twice and the token would be over-emitted.

``draw(probs, slot)`` performs the categorical draw the caller's
sampler owns (the engine's Gumbel-max, a test's inverse-CDF). The
algebra here decides WHICH candidate and WHICH distribution to draw
from; it deliberately does not own the draw, so the tree can reuse
the verify path's existing sampler bit-for-bit.

At ``widths=(1,) * K`` this consumes ``accept_uniforms[d]`` at depth
``d`` and accepts iff ``u <= p(x)`` — the chain's own test, on the
chain's own uniform, in the chain's own order. The comparison is
``<=`` (not ``<``) for that reason.

``(parents, depths, level_sizes)`` — one build for both topologies.

Walks depth by depth, assigning ids in creation order, which IS
breadth-first order with a node's children contiguous. Each
frontier entry carries the depth at which its path left the
spine (``-1`` while still on it), because that is the only state
a tail truncation needs and it is what makes the spine form fall
out of the same loop as the product form.

The widest sibling set anywhere in the tree.

``1`` exactly for a chain — the width-1 lane every gate that asks
"does this tree have siblings at all" keys on. A tail only ever
adds single-child nodes, so this is ``max(widths)`` for both
topologies; it is a property rather than that expression so a
third topology cannot make the expression quietly wrong.

``(N, N)`` bool: can query node ``i`` attend to key node ``j``?

True iff ``j`` is ``i`` itself or one of its ancestors. Every node
additionally attends to the whole cached prefix, which is outside
this block and needs no mask.

This is the ONE piece of geometry the attention backend must
honour, and getting it wrong is silent: a node that can see a
sibling is conditioning on a token that is not in its own
hypothesis, so verify would accept a token the model never
actually predicted from that prefix.

No attention kernel on this stack takes a mask — both verify
routes bake a linear per-row causal bound, which is a chain by
construction (``docs/tree-speculation-feasibility.md``). So this
is the SPECIFICATION of the relation and the reference oracle a
masked implementation would be tested against, not something the
serving path can hand to a kernel today.

``(N+1, N+1)`` int32 ancestor mask for the verify block.

The verify block is ``[last_committed, node_0 .. node_{N-1}]`` —
the same ``K+1``-shaped batch a chain builds, widened. Block row
0 is the committed token the step starts from; block row
``i + 1`` is node ``i``. Entry ``[i, j]`` is 1 iff block row
``i`` may attend to block key ``j``.

Every node attends to block key 0: the committed token is on
EVERY hypothesis, so it is an ancestor of all of them. Beyond
that a node sees only its own ancestors and itself, which is
:meth:`attention_mask` shifted by one.

This is the tensor the attention kernel consumes, and it is int32
rather than bool because that is what the kernel's aux-tensor
slot takes. The shared PREFIX needs no representation here — it
is visible to every row and lives outside the block.

``(N,)`` int32 position offset of each node past the committed end.

Depth ``d`` sits at ``+d`` — siblings SHARE a position, which is
the whole point: they are alternatives for the same slot, not a
longer sequence.

Per depth: the node ids expanded to produce that depth's nodes.

``(ROOT,)`` at depth 0. Deeper, the depth-``d-1`` nodes that
actually HAVE a child, in id order — so row 0 is the spine's.
A product tree's entry is its whole previous level, unchanged.

Per depth: the widest ``top_w`` any parent row takes there.

One ``topk`` at this width covers every row of the depth; rows
that want fewer children drop the extra ranks through
:attr:`draft_take`. A per-row width would be a data-dependent
kernel shape, which a captured graph cannot have.

Per depth: each node's flat slot in that depth's ``(rows, W)`` topk.

Indexed by depth-``d`` node in id order, so selecting with it
lands the level in breadth-first node order without a sort. The
identity permutation for a product tree, where every row uses
every rank.

Per depth ``d``: what depth ``d+1`` inherits from depth ``d``.

``(slots, rows)``, one entry per depth-``d+1`` parent row in that
row order: ``slots`` is the node's flat slot in depth ``d``'s
``(rows, W)`` topk — its TOKEN — and ``rows`` is the depth-``d``
row whose output hidden it carries. Children differ by token, not
by state, which is why those are two different tables.

The last depth expands nothing and its entry is empty.

Compact human form, and it round-trips through the flag.

``2x2x1x1 (N=14, depth=4)`` for a product tree,
``spine:2x2x2x1x1:tail2 (N=14, depth=5)`` for a spine — the
leading token of each is exactly what ``ARBI_MTP_TREE`` takes, so
a refusal message names a value an operator can paste back.

Verify + accept loop for speculative decoding.

Two acceptance modes:

  - **Greedy** (``temperature == 0`` per request, ``greedy=True`` here):
    a draft token accepts iff it equals the main-model argmax at the
    same slot. Bit-identical to the K=1 decode path at ``temperature
    == 0``: same argmax, same logits, same numerical path. A
    quality-no-regression test asserts this argmax-stability.

  - **Rejection sampling** (``greedy=False``, Leviathan-2023 /
    Chen-2023): a draft token ``x_k`` sampled from drafter
    distribution ``q_draft(·)`` accepts with probability
    ``min(1, p_target(x_k) / q_draft(x_k))``, where ``p_target`` is
    the main model's post-sampling-params distribution at slot ``k``.
    On rejection the recovery token is sampled from the residual
    distribution ``max(0, p_target - q_draft)+`` (normalized). On
    full accept (all K) the bonus token is sampled from
    ``p_target`` at slot ``K``. The output distribution is
    statistically identical to non-spec sampling at the same
    ``temperature/top_p/top_k/min_p`` — that's the contract.

Implementation. The rejection-sampling slate runs in batched ops via
:class:`GraphSafeRejectionSampler`: a single accept-test over
``(B, K)``, one batched Gumbel-max recovery sample over ``(K, B, V)``,
one batched Gumbel-max bonus sample over ``(B, V)``. Randomness reads
from the engine's persistent ``mtp_seed`` buffer — graph-capture-safe
(replay reads whatever ``copy_()`` wrote) AND deterministic-on-seed
(two runs with the same seed sequence produce bit-identical accept
decisions).

Output shape — per request, a list of accepted token IDs of length
``1 .. K+1``:

  - ``[acc_0, acc_1, ..., acc_n]`` where ``n <= K`` is the longest
    prefix s.t. all draft tokens up through slot ``n-1`` accepted;
  - the trailing entry is the main model's bonus / recovery token at
    position ``n``: when the full draft accepts (``n == K``), the main
    model's slot-K logits produce a "free" bonus token; when an earlier
    draft rejected (``n < K``), the recovery token replaces the
    rejected ``draft[n]``. The output length is always ``n + 1``: at
    least 1, at most ``K + 1``.

One request's acceptance outcome.

``accepted_tokens``: 1..K+1 token IDs to commit to the request.
``num_drafted``: K (the proposal count this step).
``num_accepted``: 0..K — how many draft tokens were accepted
    (excludes the main-model bonus / recovery token, which is
    always 1 extra committed token regardless).
``accepted_path``: WHICH draft slots were accepted, one ordinal per
    depth. ``None`` — the chain, and every path today — means the
    accepted slots are ``0..num_accepted-1``, a prefix, which is
    what lets the KV retire as a pure suffix trim. A TREE's accepted
    slots are scattered and must be compacted down before that trim,
    so it carries the ordinals here.

Pull several device tensors to host with a SINGLE targeted wait.

Stages each with a non_blocking ``.to("cpu")`` (queues an async copy
on the current stream), then waits ONCE — replacing N serial blocking
``.cpu().tolist()`` round-trips with one. CPU-resident tensors pass
straight through (no wait). Bit-identical to calling
``.cpu().tolist()`` on each.

The wait is a TARGETED ``cudaEvent`` on the current (compute) stream,
NOT a global ``torch.cuda.synchronize()``. The verify→accept host pull
only needs the staged D2H copies (and their stream-ordered predecessor,
the verify forward) to complete — it does NOT need every other GPU
stream drained. A global ``synchronize()`` drains ALL streams (the
async-output copy stream, recurrent-savepoint D2H, etc.), stalling the
whole GPU on the single-stream critical path while the host runs the
python accept/commit loop. Recording an event on the current stream
after the copies and blocking on ``event.synchronize()`` waits only
for that stream's queued work up to the copies — the next phase's
launches on other streams are free to proceed. Byte-identical host
values (the copies completed); this is purely a scheduling change.

The current stream is read inside the CUDA branch so the CPU-only
test path (no copies queued) never touches a CUDA stream.

Convert ``main_logits`` for one (slot, sequence) into a probability
distribution under the request's sampling parameters.

Mirrors the chain in :class:`arbi_serve.sampler.Sampler.sample` for
the per-row, post-temperature path: temperature scale → top_k
pre-softmax mask → softmax → top_p / min_p mask + renormalize.

Caller guarantees ``sp.temperature > 0.0``; ``temperature == 0`` is
the greedy path and routed elsewhere.

Penalties / logit_bias / external (xgrammar) processors are NOT
applied here. The verify pass receives the main-model's logits
AFTER the engine's existing pre-softmax processing chain (those
happen in the sampler when running the K=1 path; for the
rejection-sampling path we accept logits as-is and only apply
the pure sampling-params distortions here, the same way the
drafter-side ``q_draft`` is computed). The two distributions
therefore live on the same processed-logits manifold and the
rejection-sampling math (``p / q`` ratio) is unbiased.

Run the verify pass and return per-request acceptance.

Shape contract:
  - ``main_logits`` carries the main model's logits at K+1
    positions per sequence: position ``k`` for ``k in [0, K)`` is
    the prediction conditioned on the (k-1)-th draft token having
    been appended; position ``K`` is the post-draft "bonus" slot
    (or recovery slot when an earlier draft rejected).
  - ``draft_tokens`` carries the draft head's K proposals.

Greedy mode (``greedy=True``): a draft token accepts iff it equals
the main-model argmax at the same slot. Bit-identical to the K=1
decode path under ``temperature == 0``.

Rejection-sampling mode (``greedy=False``): requires ``draft_probs``
— either the dense ``(K, num_seqs, V)`` drafter distribution per
slot (true-stochastic drafting), or the
:data:`~arbi_serve.spec_decode.drafter.POINT_MASS_Q` marker (the
default greedy-draft proposal: ``q[k, b, v] = 1`` iff
``v == draft_tokens[k, b]``, consumed by the sampler's index-form
fast path — BIT-IDENTICAL accept/recovery/bonus to densifying the
one-hot, without the ``(K, B, V)`` fp32 materialization) — and
``sampling_params`` of length ``num_seqs`` (one per sequence; mixed
temperatures across sequences in the same step are honoured
per-request). Implements Leviathan-2023 / Chen-2023 acceptance: per
slot, accept with probability
``min(1, p_target(x_k) / q_draft(x_k))``; on reject, sample recovery
from ``max(0, p_target - q_draft)+`` (normalized); on full accept,
sample bonus from ``p_target`` at slot ``K``. The output
distribution at the request's stochastic settings is provably
equal to non-spec sampling — that's the correctness contract this
function honours. ``draft_probs=None`` with ``greedy=False`` is
still a LOUD error (a missing carry must never silently degrade the
accept rule).

``seed_buf`` is the engine-owned persistent ``int64 (1,)`` step
seed (typically ``Engine.mtp_seed``). Two runs with the same seed
sequence produce bit-identical accept decisions and recovery /
bonus draws — the determinism contract reproducibility tests rely
on. When omitted (``None``), a one-shot CPU tensor seeded by
``torch.seed()`` is used; that's the test-only convenience path —
production callers always pass the engine seed buffer.

``scratch`` (optional) is the engine's persistent
:class:`~arbi_serve.spec_decode.verify_buffers.SamplerScratch` for this
step. When supplied the rejection sampler writes its vocab-scale
intermediates into the frozen serving layout instead of allocating
fresh per step — byte-identical accept/recovery decisions for a fixed
seed (a pure allocation change). ``None`` keeps the fresh path.

Greedy acceptance — bit-identical to the pre-rejection-sampling path.

Routes through :func:`torch.ops.arbi_serve.mtp_verify_greedy` so an
``@support_torch_compile`` region that includes the verify step
splits at this op boundary instead of trying to trace the per-row
Python commit loop. The op is a pure-tensor argmax-and-compare;
Inductor sees stable shapes via the registered ``register_fake``.

Leviathan-2023 / Chen-2023 rejection sampling, fully batched.

Thin orchestrator over :class:`GraphSafeRejectionSampler`. The
sampler returns the accept mask + first-mismatch indices +
recovery / bonus tokens in batched form; this function turns those
per-row outputs into the per-row :class:`VerifyResult` the engine
expects.

``draft_probs`` is either the dense ``(K, B, V)`` drafter
distribution or the :data:`POINT_MASS_Q` marker (greedy-draft
one-hot on ``draft_tokens``, resolved index-form inside the sampler
— the residual sampler still sees the full ``q`` semantics required
for distributional correctness, Leviathan-2023 Theorem 1, without
the dense materialization). The slate runs on the device the verify
logits live on; only the small per-row results (``num_accepted``
``(B,)``, ``recovery`` ``(B, K)``, ``bonus`` ``(B,)``) cross to
host below — the same single-D2H budget the greedy path uses.

Determinism. Identical ``seed_buf`` values → identical accept
decisions and recovery / bonus draws. The engine advances
``mtp_seed`` once per step so successive calls draw fresh
randomness while staying reproducible-on-seed for replay tests.

Persistent verify-pass buffers for cudagraph readiness.

Allocating fresh device tensors per step
via :func:`torch.tensor(flat_*, device=...)` — one allocation per
``input_ids`` / ``positions`` / ``slot_mapping`` / ``cu_q`` /
``seq_lens`` / ``block_table`` / ``cu_k`` / ``draft_tensor`` per
forward — does not work with captured cudagraphs: the kernel
launches bake in ``data_ptr``s, so the next step's fresh tensors
have different addresses and the replay reads from stale memory.

This module pre-allocates worst-case-sized device buffers once at engine
build (or MTP-driver attach time) and rebinds the verify-pass
builder to ``copy_()`` host-side data into stable slices of those
buffers each step. The captured graph then captures launches
against the persistent ``data_ptr``s and replays correctly across
steps as long as the per-step ``(B, K + 1)`` slice fits the
worst-case shape.

Lifecycle. The pool is built once when the MTP driver is attached
(:func:`build_for_engine`) and dropped by a model rebuild
(``reset_model_state_for_build``), which builds a new generation. Its
tensors are registered with :class:`SleepableTensorPool` so their VAs
survive release/resume without invalidating captured graphs.

R1 (uniform K per step) is a hard prerequisite: the per-step slice
of every flat dim is ``B * (step_k + 1)`` with one ``step_k`` for
the whole step, so worst-case sizing
``max_num_seqs * (max_k + 1)`` covers every reachable shape.

Worst-case-sized persistent buffers for the MTP verify pass.

Sizing. Every flat-token dim is sized at ``max_num_seqs *
(max_k + 1)`` so any uniform-K verify slate (R1) fits a
contiguous prefix slice. ``block_table`` is
``(max_num_seqs, max_pages)`` with ``max_pages`` taken from the
engine's ``per_seq_pages`` bound. ``draft_tensor`` is ``(max_k,
max_num_seqs)`` matching :func:`verify_and_accept`'s expected
layout.

Per-step the engine writes the active prefix in-place via
:meth:`load_step` and consumes the returned ``StepView`` (a
bundle of slices into the same persistent buffers). The slice
objects do not allocate; they just rebind ``shape`` / ``stride``
over the existing storage.

Dropped only by a model rebuild, which sets ``eng.verify_buffers =
None`` and lets GC reclaim before :func:`build_for_engine` makes the
next generation. Sleep keeps the pool: the run loop is quiesced while
``_memory_released`` is set, and the tensors' VAs are stable across
release/resume.

Per-step slices into :class:`VerifyBuffers`'s persistent rejection-
sampler scratch. Threaded into
:meth:`GraphSafeRejectionSampler.sample_batched` as ``scratch=`` so the
sampler writes ``draft_probs`` / ``p_target`` / ``u_gumbel`` /
``gumbel_noise`` into these views instead of allocating fresh.

Shapes (uniform K over the slate, R1): ``draft_probs`` ``(K, B, V)``;
the rest ``(K + 1, B, V)``. Slicing shares storage with the pool, so a
served stochastic verify step allocates ZERO new vocab-scale memory.

``u_gumbel`` / ``gumbel_noise`` are ``None`` on a Philox-path pool
(the counter-based verify tail generates its Gumbel noise in-kernel,
so there is no materialized-noise scratch to slice). ``draft_probs``
is ``None`` on a pool whose boot cannot carry a dense drafter ``q``,
which is the only thing that reads it.

Per-step slice view over :class:`VerifyBuffers`'s persistent
storage. Returned by :meth:`VerifyBuffers.load_step`; passed to
the verify-pass batch builder which threads the slices into
:class:`ScheduledBatch`.

Slicing a tensor doesn't allocate; the view's tensors share
storage with the pool, so captured graphs that recorded
launches against the pool's ``data_ptr``s still see the right
addresses on replay.

Build a :class:`VerifyBuffers` sized to the engine's bounds.

Worst-case dims:
  * ``max_num_seqs = cfg.batch.max_batch``
  * ``max_k = mtp_driver.max_k``
  * ``max_pages = max(ceil((max_context + max_k) / block_size), ceiling)``

The block table spans a row's committed context PLUS the speculative
draft tokens appended for verify, so a request whose committed context
reaches ``max_context`` needs page room for ``max_k`` more tokens — one
page beyond ``ceil(max_context / block_size)`` when the context ends on a
page boundary. It is sized at the pool page ceiling so a post-grow
max_context widen stays addressable without a re-capture.

Copy host-side step data into the persistent buffers and
return a :class:`StepView` of slices.

The view's tensors share storage with this pool — captured
graphs that recorded launches against
``self.input_ids.data_ptr()`` still see the right addresses
on replay, with the per-step content written in place.

Incremental host staging. The per-step flat content is written
directly into the pre-allocated pinned host mirrors
(``h_input_ids`` / ``h_positions`` / ...) via the zero-copy
``numpy()`` view, then one ``copy_(non_blocking=True)`` per dim
refreshes the device twin. This replaces the per-step
``torch.as_tensor(python_list, ...)`` host-tensor allocations
(6 fresh CPU tensors per verify step) — the MTP-verify analogue
of the decode path's persistent_input_batch pinned ring
buffer. The numpy assignment ``arr[:n] = list`` uses NumPy's
C-level coercion (a single C buffer + memcpy) rather than
building an autograd-aware Tensor each step.

Block-table skip. When ``bt_occupant_keys`` is supplied (one
``(req_id, n_pages)`` per row, from
:func:`_verify_block_table_rows`'s cache), a row whose key is
unchanged since the prior :meth:`load_step` already holds the
correct page ids in the pinned mirror, so its numpy rewrite is
skipped (append-only decode grows the chain only at a
``block_size`` boundary). The device H2D below always copies the
full active ``[:B, :rows_pages]`` prefix, so a skipped row's
device content is byte-identical to a full rewrite — the skip is
a host-side optimization only, never a content change. With
``bt_occupant_keys=None`` every row is rewritten (legacy /
from-scratch behaviour), so the device result is identical
either way.

Pre-conditions (R1 invariant + worst-case bounds):
  * ``len(cu_q) == B + 1``;
  * ``cu_q[-1] == B * (k + 1)`` for some uniform ``k``;
  * ``B <= max_num_seqs``;
  * ``k <= max_k``;
  * ``max(len(row) for row in block_table_host) <= max_pages``.

``cu_k_host`` is a CPU tensor (already includes the leading
zero entry). ``block_table_host`` is a Python list-of-lists
— we materialise it into the persistent ``block_table[:B]``
slice in one ``copy_()``.

Extend an already-staged ``B``-row, uniform-``S`` verify step
in-place to ``b_prime`` rows with benign padding, so a captured
``(b_prime, S)`` shared-buffer verify graph can replay over this
pool and produce correct logits for the real ``[0:B]`` rows.

Used only by the verify-forward pad-up
(:func:`arbi_serve.runtime.decode_pad.replay_padded_verify`) when
the live ``B`` falls between the captured verify-shape ladder
rungs (e.g. 5/6/7) and ``cfg.decode_pad_cudagraph`` is on. The
real ``[0:B]`` rows were already written by :meth:`load_step`;
this only fills the padding window ``[B:b_prime]`` (per-row dims)
and ``[B*S:b_prime*S]`` (flat dims) with values that the shared-
buffer captured graph reads but that cannot corrupt a real row:

  * flat ``input_ids`` / ``positions`` → 0.
  * flat ``slot_mapping`` → 0 (page-0 null slot; the KV scatter
    writes harmlessly into the reserved null page).
  * per-row ``seq_lens`` → 1 (each padding row reads a single
    null token of K context).
  * ``cu_seqlens_q`` → continue the uniform ``S`` stride so it
    reads ``[0, S, 2S, ..., b_prime*S]`` exactly as the captured
    graph baked it.
  * ``cu_seqlens_k`` → cumulative of the padded ``seq_lens``.
  * ``block_table[B:b_prime]`` → all-zero (page-0 null).

Per-row independence of the verify forward (attention is per-
sequence via ``seqused_k`` / ``cu_seqlens_k``; no cross-row
reduction) guarantees the real rows' math is unaffected by the
padding rows. ``b_prime <= max_num_seqs`` is the caller's
contract (the pad target never exceeds the pool capacity).

Copy per-row drafts into the persistent ``draft_tensor`` and
return the ``(k, B)`` view.

Rows whose drafts are ``None`` (cold-path-degraded) are not
expected to call this — under R1 every row is either active
at the same K or the whole step is K=0 (no draft tensor
needed). The caller is expected to enforce that.

The host staging goes through the persistent PINNED
``h_draft_tensor`` mirror (zero-copy numpy column writes, then one
``copy_(non_blocking=True)`` of the matching ``[:k, :B]`` slice) —
a fresh pageable CPU tensor per step made the H2D effectively
synchronous behind whatever is enqueued on the stream. Content is
byte-identical to the legacy fresh-tensor path.

Gather the ``(K+1, B_mtp, V)`` verify logits into persistent scratch.

``logits`` is the flat ``(N_tokens, V)`` per-token verify output;
each MTP row contributes ``K+1`` contiguous flat tokens at
``[offsets[i], offsets[i+1])``. The target slot ``(slot, col)`` reads
flat row ``offsets[mtp_rows[col]] + slot``. We materialise that gather
index into the pinned host mirror, copy it H2D once, and run ONE
``index_select`` into the persistent ``sc_verify_logits`` slice — no
full-vocab ``torch.stack`` + ``.permute`` + ``.contiguous`` copy and
no fresh per-step vocab-scale allocation. The returned slice shares
storage with the persistent buffer (stable VA, cudagraph-safe).
``None`` when no MTP row is present (cold step).

Returns ``None`` (caller falls back to the legacy stack path) when
the pool was built without the gather scratch (``vocab_size == 0``).

Return the persistent drafter-``q`` destination sliced to ``(k, b, vocab)``.

``None`` when this pool has none (``vocab_size == 0``, or a boot that
cannot carry a dense drafter ``q``), when the slate exceeds the
worst-case bounds, or when the drafter forms its proposal at a width
this pool was not sized for — the caller then allocates fresh.
Returning ``None`` rather than raising is deliberate: the drafter is
the one path whose failure trips the chain breaker and cold-paths a
request's speculation, and a shape this pool cannot serve is a reason
to stop POOLING, not a reason to stop drafting.

Slicing does not allocate. The view shares storage with the pool, so
the NEXT draft call overwrites what this one returned — every caller
takes its per-row ``(K, V)`` carry off it with
:func:`~arbi_serve.spec_decode.drafter.carry_draft_probs`, which
clones, before the next call runs.

Return the persistent rejection-sampler scratch sliced to ``(k, b)``.

``draft_probs`` is the ``(k, b, V)`` prefix of ``sc_draft_probs``;
``p_target`` / ``u_gumbel`` / ``gumbel_noise`` are the
``(k + 1, b, V)`` prefixes of their ``sc_*`` buffers. A member the
pool skipped rides as ``None`` (the materialized-noise scratch on a
Philox device; the dense-q scratch on a boot that cannot carry one)
and its consumer allocates fresh. Slicing does not allocate — the
views share storage with the persistent buffers, so a stochastic
verify step writes into them in place.

``sc_p_target`` is the member that decides whether this pool HAS
scratch, because it is the one no boot-fixed condition can retire:
it is written by the gathered verify path, and a live
``ARBI_SHARDED_VERIFY=0`` (``scope="runtime"`` — no rebuild) or a
single request with an unbounded ``top_k`` routes there with the
pool already frozen.

Raises if the pool was built without scratch (``vocab_size == 0``)
or the requested slate exceeds the worst-case bounds — fail loud,
never silently fall back to a fresh allocation here.

The constraint x draw-site matrix: what every verify draw does about
every per-token constraint the K=1 draw would have honoured.

A speculative verify draw must reach the accept test on the same processed-
logits manifold the K=1 draw would have used. Two independent things can go
wrong, and only one of them is visible to a per-guard test:

* a draw site does not apply a constraint it should — a guard that exists
  somewhere else did not run here;
* nobody ever wrote the guard, on any draw site, so there is nothing for a
  per-guard test to be missing.

The second is invisible to any mechanism keyed on a guard that exists. This
module closes it by making the ACCOUNTING, not the guard, the unit of
enforcement: :data:`CONSTRAINTS` x :data:`DRAW_SITES` is a full product, and
every cell of it must resolve to exactly one declared state.

``Applied``
    The draw site applies the constraint. Proven structurally: the applier
    must be reachable in the call graph from the draw site's entry
    (``tests/test_verify_constraint_matrix.py``), so a declaration that the
    code contradicts is a failure, not a comment.
``Refused``
    The draw site refuses a slate carrying the constraint, by name, at its
    gate. The declared ``reason`` must appear in the gate's source and
    ``proof`` must name a test that trips it.
``NotApplicable``
    The constraint and the draw site cannot co-occur. ``proof`` names the
    test that establishes that structurally; a reason with no proof is a
    claim, and claims are what this module exists to replace.

:data:`KNOWN_DROPS` carries the cells that are none of the three: the
constraint is silently dropped there. It is a shrink-only ratchet with a
hard ceiling — growing it requires editing the ceiling literal in the same
change, in review. A cell may leave the matrix only by being fixed.

Both axes are closed AGAINST THE CODE, not against prose: the draw sites are
derived by walking the call graph for callers of the accept primitives, and
the constraints are derived from the boot-time logits-processor chain plus
the verify-side appliers. A new verify path, or a new processor in the K=1
chain, therefore lands as an undeclared axis entry and fails the gate.

A per-token transform the K=1 draw applies before the sample.

``chain_source`` anchors the constraint to the code that owns it on the
K=1 path: the class appended to the sampler's logits-processor chain, or
the in-sampler transform function.

``dispatch`` says how the verify draw reaches it, because the two forms
are checked differently:

``"function"``
    A named verify-side applier the draw site must reach in the call
    graph. ``verify_applier`` names it.
``"processor"``
    Dispatched through
    :meth:`~arbi_serve.sampler.processors.LogitsProcessor.apply_to_verify_slate`
    off the registered chain. Reachability cannot attribute this: every
    processor shares the one method name, so a draw site reaching it
    proves the CHAIN runs, never that THIS constraint is in it. The
    evidence is the implementing class instead — whether
    ``chain_source`` is registered in the chain the seam iterates, and
    whether its definition of the method masks or refuses.
    ``verify_applier`` is ``None``.

One place a verify draw is taken.

A draw site is the innermost function that both calls an accept
primitive and can see the slate's per-request state — the level at which
a per-token constraint is still knowable. ``gate`` names the predicate
that admits a slate to this site, or ``None`` when the site is
unconditional for its configuration.

A refusal a draw site's gate carries that is NOT a per-token constraint.

The constraint x draw-site product cannot see these: they are properties
of the CONFIGURATION or of the request's SUPPORT, not transforms the K=1
draw applies. They were therefore invisible to every check in this
module — which is exactly the blind spot the module exists to close, one
axis over. A gate refusal reason that is neither a declared constraint
cell nor declared here fails
``test_every_gate_refusal_is_declared``.

``fragment`` must appear verbatim in the gate's source, so a reason that
is reworded or removed cannot keep a stale declaration alive.

The constraint is dropped at this draw site: nothing applies it and
no gate refuses it.

Every entry names the owner accountable for closing it. Entries leave
only by being fixed.

Repetition / frequency / presence penalties on the MTP verify slate.

The K=1 sampler applies the three penalties before temperature
(:func:`arbi_serve.sampler.processors.apply_penalties`). The verify slate
is the speculative twin of that draw, so it must apply them too or a
request's penalties are honoured on exactly the steps that did not
speculate — which, under MTP, is almost none of them.

What makes the speculative case different from the K=1 case is that the
slate's ``K + 1`` positions do NOT share a history: position ``p``'s
history is the committed prefix plus the first ``p`` drafted tokens. A
single ``(B, V)`` committed-prefix count broadcast across the slate would
under-count every token the step itself drafts, so the accepted prefix
would be drawn from a different distribution than the K=1 path would have
produced and speculation would stop being distribution-preserving.

The split here keeps both properties at the cost of one full-vocab pass:

* the committed-prefix term goes through the main sampler's own
  :func:`~arbi_serve.sampler.processors.apply_penalties`, so the resident
  accumulator / fused-kernel fast paths are shared rather than forked;
* the in-step term touches only the ``<= K`` distinct tokens the step
  actually drafted, recomputed EXACTLY from the raw logit at that entry.

The correction is exact rather than incremental because the repetition
knob is not additive in the count (it is a sign-dependent multiply gated
on ``count > 0``), so a delta applied on top of an already-penalised value
would not reproduce the K=1 formula.

Apply ``logit_bias`` to every position of the verify slate.

Unlike the penalties this needs no per-position history: a bias is a
constant offset per (row, token), so the K=1 chain's broadcast over the
leading dim is already exact for the slate.

``window`` narrows the write to one vocab-parallel rank's slice; the
arithmetic is the K=1 chain's, taking the window as a parameter.

Whether any row asks for a transform this module writes.

The single predicate both verify paths screen on, so "is there anything
to apply" cannot answer differently from "what gets applied".

``(B, K)`` committed-prefix occurrence counts for the drafted ids.

Read out of the sampler's OWN penalty accumulator whenever it is
resident. That buffer already holds the per-(row, token) occurrence
counts — it is what the K=1 draw's penalties are computed from — so the
committed term here is a gather of ``<= K`` entries per row and nothing
else. Deriving it instead from the padded history, as this used to,
stood up a SECOND history structure beside the accumulator's and paid
an O(context) broadcast compare per step to reach the same numbers the
accumulator was already maintaining.

MUST be called AFTER the committed-prefix pass, because that pass is
what folds this step's newly committed tokens into the accumulator.
Reading it before would be short by exactly the tokens the step just
committed, which are the ones most likely to repeat.

The padded-history path stays as the fallback for a caller with no
resident accumulator (CPU tests; a vocab/dtype the boot reserve was not
computed for), which is also what keeps it exercised.

``(drafts, ids)`` — the step's drafts, and an index-safe twin.

ONE H2D for the whole correction. ``drafts`` keeps the out-of-vocab pad
sink so "is this slot real" stays readable; ``ids`` substitutes every
slot the window cannot address — a pad slot, or (under a window) a real
draft another rank owns — with the row's FIRST substitutable drafted id,
so it can be handed straight to ``gather`` / ``scatter_``.

Substituting rather than masking is what makes the scatter safe without
a host-visible count of live entries. A substituted slot then duplicates
a real slot, and duplicates are harmless here BY CONSTRUCTION: two slots
carrying the same id see the same committed count and the same in-step
count, so they compute the same value and race to write identical bytes.
Masking instead would need ``nonzero()``, i.e. a device-to-host sync on
the hot path — the exact cost this rewrite exists to remove.

The substitution must replace the GLOBAL id, not just the column index:
every quantity the correction computes for a slot is derived from
``ids``, so two slots agree bit-for-bit exactly when their ``ids`` agree.
Rewriting only the column would leave the duplicate computing a
different value for the same column, and a ``scatter_`` with duplicate
indices and disagreeing values has no defined winner.

A row with NO substitutable draft falls back to the window's first id.
That id cannot be one the row drafted (or it would have been
substitutable), so its in-step count is zero and the correction leaves
the entry exactly as the committed-prefix pass wrote it.

Apply the three penalties to the verify slate, per speculative position.

No-op when no row asks for a penalty, which is the stock serving case
for a thinking-mode profile (Qwen ships ``presence_penalty=0.0``,
``repetition_penalty=1.0`` there) — the early-out keeps this off the
hot path entirely rather than paying a full-vocab pass to subtract zero.

Everything past that early-out is device work driven by ONE ``(B, K)``
upload. There is no host walk over positions x columns x tokens and no
per-entry coordinate list: the correction's shape is fixed by ``(K, B,
K)``, which is a few hundred elements at any served K, so deriving it on
device costs less than the Python that used to describe it.

``window`` narrows the WRITE to one vocab-parallel rank's vocab slice
while every count stays a global-vocab quantity. Both terms restrict
exactly: a token's committed count is a property of the token, not of
which rank holds it, and its in-step count is read off the whole
``(B, K)`` draft row, which every rank has identically. So each rank
reproduces the K=1 value for exactly the tokens it owns, and the ranks
together reproduce the whole row.

Shard-resident MTP verify sampling — candidate-compaction core.

Replaces the per-cycle full-vocab verify ``all_gather`` (the
:class:`~arbi_serve.models.linear.VocabParallelLMHead` gather over the
K+1 verify rows) with a tiny exchange of per-shard TOP-CANDIDATE
``(value, global_index)`` pairs, the
verify-sampling generalization of the draft path's distributed greedy
argmax (:meth:`Qwen3_5MtpHead._distributed_greedy_argmax`, which
exchanges one ``(B, 2)`` scalar pair).

Why a bounded candidate set is sufficient
=========================================

The stochastic verify sampler's ``p_target`` chain (temperature →
top_k → softmax → top_p → min_p, see
:meth:`GraphSafeRejectionSampler._build_p_target`) zeroes every token
outside the per-row top_k kept set. With a finite per-row ``top_k``
(the served stochastic preset: top_k=20), EVERY downstream consumer of
the full-vocab row reads only that kept set:

  * ``p_target(x_draft)`` for the accept test is either a kept token's
    probability or EXACTLY 0 (masked) — both derivable from candidates;
  * the recovery residual ``(p_target - q_draft)+`` is positive only
    where ``p_target > 0`` (support ⊆ kept set);
  * the bonus draw samples ``p_target[K]`` (support = kept set);
  * top_p / min_p / softmax renormalize WITHIN the top_k survivors.

So each rank contributes its local top-``n_cand`` (value, index) pairs
per verify row (``n_cand >= 2 * top_k`` — top_k winners plus equal-value
tie headroom), the tiny all_gather replicates them, and every rank runs
the IDENTICAL sampling chain over the compacted ``(K+1, B, C)`` rows —
rank-symmetric like the rest of the SPMD verify step (the symmetry is
the broadcast). The greedy verify (temperature 0) needs only the
per-slot argmax, i.e. ``n_cand = 1`` — byte-exact by the same
lowest-global-index tie-break argument as the draft-path primitive.

Numerics contract (what is and is not bit-exact)
================================================

Per-shard logits are the SAME ``invariant_linear`` shard matmul the
gathered path concatenates, so the compacted candidate VALUES are
byte-identical to the gathered row's entries. The chain reuses the SAME
sampling-params functions (:func:`_apply_sampling_params_uniform_block`
et al.) over ``C`` columns instead of ``V``; the accept uniforms are the
SAME device-independent CPU-generator draw; the recovery/bonus Gumbel
noise is the SAME Philox function keyed on ``(seed, row,
GLOBAL vocab index)`` — so noise per token is byte-identical to the
full-vocab kernel's. The residual divergence surface is exactly:

  1. the softmax normalizer / top_p cumulative sums reduce over ``C``
     elements instead of ``V`` (the excluded elements are exact zeros,
     but fp32 addition-tree GROUPING differs) → ~1-2 ULP drift in
     ``p_target``;
  2. a sampler decision (accept test / Gumbel argmax / top_p boundary)
     flips only when that ULP drift crosses an exact fp tie — far below
     the logit drift known to move accept_len. Drafter inputs (hidden
     states, committed-token embedding path) are untouched, so draft
     quality — the accept_len driver — is unchanged; a flipped decision
     is a resample from the same target distribution, not drift.

Greedy verify and the cold-row frontier argmax are BYTE-EXACT (index
selection only, no reductions over values).

Tie-overflow sentinel
=====================

The top_k mask keeps every token whose value ties the k-th largest, so
the kept set can exceed ``top_k`` on exact fp ties. A shard's gathered
``n_cand`` candidates cover the kept set unless MORE than ``n_cand`` of
its tokens clear the global threshold (>= ``n_cand - top_k`` extra
exact ties in one shard — bf16 logits make small ties common, deep ties
at one value rare). :func:`tie_overflow_rows` detects exactly this
(shard's last gathered value still >= the global k-th value) so the
integration layer can count / surface it; a missed deep tie perturbs
``p_target`` by at most the dropped tied token's renormalization share.

This module is the PURE core: it operates on already-merged candidate
tensors and takes all randomness as inputs, so the full compaction ↔
gathered-path equivalence is testable on CPU with no process group. The
TP wrapper (shard matmul + pair all_gather + engine wiring behind
``ARBI_SHARDED_VERIFY``) layers on top.

One rank's per-row top-``n_cand`` candidate pairs, globally indexed.

``local_logits`` is the rank's un-gathered lm_head shard output for
the flat verify rows (any leading shape; last dim = per-shard padded
vocab). ``n_valid`` counts the REAL vocab columns on this shard
(``min(per_shard, org_vocab - shard_start)`` — the padded tail of a
short last shard is excluded exactly like the gathered path's
``[..., :org_vocab]`` trim, so a zero-valued pad column can never
become a candidate). ``shard_start`` maps local to global indices.

Returns a ``(..., n_cand, 2)`` fp32 tensor packing ``(value,
global_index)`` per candidate — one tensor so the TP exchange is a
single small all_gather, exactly like the draft-path ``(B, 2)``
scalar-pair primitive. The fp32 index encoding is exact for vocab
sizes < 2^24 (248k ≪ 2^24). Rows are NOT ordered; callers merge
shards with :func:`merge_candidates`, which canonicalizes order.
Values are fp32 upcasts of the (bf16) logits — the same lossless
widening the gathered sampler applies, and order-preserving, so the
per-shard top-``n_cand`` SET equals the bf16 row's.

A shard with ``n_valid < n_cand`` pads with ``(-inf, PAD_IDX)``.

Merge per-shard candidate pairs into canonical per-row candidates.

Returns ``(cand_vals (..., C), cand_idx (..., C) int64)`` with
``C = n_shards * n_cand``, sorted ASCENDING by global index per row
(pads first). The canonical order makes every positional
first-occurrence tie-break downstream (``torch.argmax`` over
candidate columns, ``torch.sort``-based top_p) resolve to the LOWEST
GLOBAL VOCAB INDEX — the same winner the gathered full-row ops and
the chunked Triton kernels produce on exact fp ties.

Global argmax token id per row — lowest global index on ties.

Bit-exact to ``full_row.argmax(-1)`` over the gathered row whenever
the true argmax is among the candidates (guaranteed for any
``n_cand >= 1``: the global max is every shard's local top-1), with
the tie-break the draft-path primitive proved: the smallest global
vocab index among the max-valued tokens wins.

Greedy MTP verify over merged candidates.

Same contract as :func:`torch.ops.arbi_serve.mtp_verify_greedy` —
``(accepted_count (B,) int64, main_argmax (K+1, B) int64)`` — and
byte-exact to it (argmax + integer compare only; no value
reductions). ``n_cand = 1`` per shard suffices for this path.

The verify ``p_target`` chain over compacted candidate rows.

Routes the SAME sampling-params chain the gathered sampler applies
(:func:`_apply_sampling_params_uniform_block` for a uniform slate,
:func:`_apply_sampling_params_batched` otherwise — temperature →
top_k → softmax → top_p → min_p) over the ``C`` candidate columns.
Padded ``-inf`` columns carry zero probability through every stage.

Support-coverage precondition (fail-loud): every stochastic row
needs ``0 < top_k <= n_cand_shard`` — the per-shard candidate budget
must cover a shard holding the ENTIRE kept set. Unbounded-support
rows (``top_k`` disabled) cannot be compacted; the integration gate
routes those slates to the gathered path, and reaching here with one
is a routing bug.

Returns ``(K+1, B, C)`` fp32 probabilities over the candidate
columns (each row sums to ~1). Values match the gathered chain's at
the same tokens up to fp32 reduction-order ULPs (the module
docstring's divergence surface).

Read a per-row payload value at one global vocab id, else 0.

``p_target`` (and any candidate-aligned payload) is EXACTLY zero at
every non-candidate token (masked by top_k), so the
match-or-zero contract reproduces the gathered path's
``p_target.gather(-1, token)`` for any token in or out of the kept
set. A ``PAD_IDX`` column can never match a real token id.

Rejection-sampling accept + recovery/bonus draws over candidates.

The compacted twin of
:func:`torch.ops.arbi_serve.mtp_sample_residual_philox` — same
output contract ``(accepted_count (B,), accepted_tokens (K+1, B)
int64 GLOBAL ids)`` and the same math:

  * accept test ``u * q(x) <= p(x)`` per (slot, row) with the SAME
    uniforms — identical decisions given identical ``p(x)``;
  * recovery: Gumbel-max over ``log((p - q)+)`` per slot (the
    residual's ``1/Z`` drops out of the argmax, exactly like the
    Philox kernel), support restricted to positive-residual
    candidates — which IS the full-row support, since the residual
    is positive only where ``p > 0``;
  * bonus: Gumbel-max over ``log(p[K])``;
  * degenerate row (empty residual): argmax of ``p`` — the kernels'
    fallback;
  * every argmax tie-breaks to the lowest global vocab index (the
    canonical ascending candidate order + first-occurrence /
    min-index reductions), matching the chunked Triton kernel's
    strictly-greater running max and ``torch.argmax``.

``gumbel_cand`` must be the noise evaluated at ``cand_idx`` under
whichever noise source the run uses — the counter-based Philox
function of ``(seed, row, global index)`` on the served path (byte-
identical per token to the full-vocab kernel's in-kernel draw), or a
materialized full-vocab tensor indexed at the candidates (the
CPU/rollback path). Randomness stays an INPUT so the selection math
is a pure function.

Detect rows whose kept set may exceed the gathered candidates.

The top_k mask keeps every token whose value >= the global k-th
largest, so ties AT the threshold can push the kept set past
``top_k``. A shard's candidates cover its share of the kept set iff
its LAST (worst) gathered value falls strictly BELOW the global
threshold — otherwise deeper, un-gathered tokens on that shard could
tie the threshold and be silently dropped from ``p_target``'s
renormalization. Returns a ``(K+1, B)`` bool mask of the affected
rows (expected ~never at ``n_cand >= 2 * top_k``; the integration
layer counts and surfaces them — the distributional error is bounded
by one tied boundary token's share). A row whose threshold is ``-inf``
is never reported: see below.

The numerics condition every data-dependent verify WIDTH shares.

More than one lever decides, per step and from the data, how many rows the
verify block carries. The per-slot depth gate
(:mod:`arbi_serve.spec_decode.depth_gate`) stops the drafter chain short;
``ARBI_DFLASH_DYNAMIC_K`` truncates the slate's uniform K to a confidence
head's confident prefix. They share the hazard, so they share the refusal:
a second copy of it is how one lever comes to be guarded and its
neighbour, arriving later, not.

Under default numerics the verify block's row count selects the
accept-critical GEMM's reduction order, and wherever the next draft is
seeded by that same forward — a recurrent target's chunked verify scan, a
DSpark drafter's residual tap — it also selects the width that seeds it.
The served text is then a function of the widths the lever chose, so the
lever perturbs the trajectory it is scored on and nothing downstream
raises. ``ARBI_ACCEPT_INVARIANT=1`` pins the row count and routes the
recurrent verify to a per-token recurrence; both are width-free, which is
what turns a varying width back into a throughput decision.

There is no fallback that restores this, so it is refused at boot by name
rather than served.

Tenant abstractions.

Public surface:

* :class:`TenantContext` — frozen, attached to every :class:`Request`.
* :class:`QuotaBudget` — per-tenant rate / VRAM ceilings.
* :data:`ANONYMOUS_TENANT` — sentinel default for unauthenticated paths.
* :class:`QuotaTracker` — per-tenant runtime consumption tracker
  (rate-limit + concurrency + cumulative tokens), wired into the
  scheduler's admission filter.

The :class:`TenantContext` / :class:`QuotaBudget` interfaces live in
:mod:`context`; :mod:`quota` ships the consumption tracker.

TenantContext + QuotaBudget.

The :class:`TenantContext` is the per-request tenant identity shipped
through every layer that cares: scheduler (fairness + quota), pool
registry (admission), cache key (radix tree namespace), and the
deployment-level A/B routing rule (``backend_overrides``).

Per-tenant rate / VRAM ceiling.

All fields are upper bounds; ``None`` means unlimited (the
:data:`UNLIMITED` instance below is the all-``None`` sentinel).

These are consumed in scheduler admission and pool-registry
fail-loud refusal.

First-class tenant identity attached to every :class:`Request`.

The ``backend_overrides`` mapping carries deployment-level A/B
routing decisions: a routing rule at request entry populates the
tenant context with the cohort's desired backend per
:class:`StateKind`. Per-request ``?attention_backend=...`` overrides
still win on top.

``cohort`` is an opaque label (``"shadow"`` / ``"canary-7"`` / etc.)
surfaced into per-tenant metrics labels so dashboards split traffic
by cohort without engine changes.

Frozen so it can be hashed / used as a cache-key component without
surprise-mutation between scheduler and pool-admission paths.

Per-tenant rate-limit + cumulative-consumption tracker.

The :class:`QuotaBudget` dataclass itself lives in
:mod:`arbi_serve.tenant.context`. This module ships the
*runtime* tracker that the scheduler consults during admission:
:class:`QuotaTracker` keeps a per-tenant token-bucket for
``tokens_per_second``, an in-flight gauge for
``max_concurrent_requests``, and a cumulative VRAM-bytes gauge.

The tracker is intentionally in-memory and best-effort. Production
deployments back the same surface with Redis / SQL — the
:meth:`QuotaTracker.try_admit` API stays identical. The
:class:`arbi_serve.server.auth.AuthManager` already persists daily
token quotas; that path is orthogonal (admin-quota; this is
scheduler-quota).

Threading model: single-process scheduler-side enforcement. The
scheduler calls :meth:`try_admit` from one event loop; no async lock
is needed (asyncio is cooperative). Tests that spin up multiple
schedulers must construct one tracker per scheduler.

Token-bucket meter — capacity == rate (1-second burst window).

Mirrors :class:`arbi_serve.server.auth.TokenBucket` shape but is
sync-only: scheduler admission runs inside the engine's run loop
and consumes ``time.monotonic()`` directly. Kept private because
the auth path's bucket carries different semantics (daily quota,
user-id keying, async lock); the scheduler's bucket only needs
rate enforcement.

Result of :meth:`QuotaTracker.try_admit`.

``admitted=True`` means the scheduler may emit this request to the
active batch; ``admitted=False`` means defer (or reject at the
server boundary with HTTP 429).

``reason`` is a stable string for metrics / logs:
  * ``"ok"`` — admitted.
  * ``"rate_limited_tps"`` — tokens-per-second bucket empty.
  * ``"concurrent_cap"`` — in-flight at ``max_concurrent_requests``.
  * ``"prompt_too_long"`` — single prompt exceeds ``max_prompt_tokens``.

Per-tenant rate / concurrency / cumulative-tokens tracker.

The scheduler consults this once per admission decision. A
:class:`TenantContext` carries the *budget* (immutable, frozen);
the *consumption* lives here so two requests of the same tenant
share one bucket / one in-flight gauge.

Lookup is by ``tenant_id`` so both the ``Request.tenant_id``
field and the ``Request.tenant.tenant_id`` field route to
the same per-tenant state.

All operations are O(1).

Decide whether ``tenant`` may be admitted right now.

``prompt_tokens`` is checked against the budget's
``max_prompt_tokens`` (a per-request size cap). ``cost_tokens``
is the load this admission imposes on the tps bucket — usually
the number of new tokens the request will produce this step
(1 for decode, ``want`` for prefill chunks).

On admission, the tps bucket is debited and the in-flight
counter is incremented. The caller must call :meth:`release`
when the request finishes so the in-flight gauge stays
accurate.

Output-quality checks shared by benches and live-GPU gates.

Degenerate repetition (a generation collapsing into a short loop —
``\n\n\n…``, ``I\n\nI\n\n…``, a phrase echoed forever) silently
corrupts two things at once: spec-decode accept rates soar on the loop
(a drafter trivially predicts it), inflating throughput numbers, and
the served output is garbage. Every perf claim and every greedy-output
gate must run :func:`is_degenerate` so a loop can never pass as a
healthy measurement.

Shortest period ``p`` such that the tail is ``>= 3`` consecutive
repeats of the last ``p`` elements covering at least
``_LOOP_TAIL_FRACTION`` of the inspected window. ``None`` if the
tail is not a loop. Works on any sliceable sequence — ``str``
(characters) or a token-id list.

Image-build bake and source fingerprint for tkv's JIT CUDA extensions.

tkv builds its hand-written CUDA kernels through
``torch.utils.cpp_extension.load_inline``, which resolves under
``$TORCH_EXTENSIONS_DIR`` and is keyed by the extension NAME alone. The
image bakes those kernels at build time into ``/opt/cache-baked`` and the
boot seeds them across (``cli.bootstrap._seed_prebaked_torch_extensions``),
so the served ``.so`` is chosen by name and nothing on that path ever looks
at the sources it was compiled from. That is the silent-stale hazard
``docs/boot-cache-keys.md`` catalogues: under a bind-mounted or re-pinned
tkv, a name that still matches serves machine code built from other source.
``load_inline``'s own ninja rebuild does not save it — the seeded ``.so``
lands in the build dir before ninja ever considers it, and a production
``TKV_NO_JIT=1`` boot refuses to build at all.

So each baked extension dir carries a sidecar naming the tkv sources it was
built from, and the seed refuses a dir whose recorded digest disagrees with
the installed tkv. Refusing to seed is not a failure: the runtime then
builds from what is actually installed, which is the correct answer.

**What the digest covers.** Every nvcc-reachable source under
``tkv/kernels`` AND the Python beside it — the launchers that compose the
``defines`` and the ``_jit_common`` module that composes the nvcc flag
list. Folding the build recipe's own source in is what makes a FLAG change
move the digest without arbi-serve restating tkv's flag list, which would
be a second expression of a rule tkv owns. Test trees are excluded: they
reach no build.

**What it deliberately does not gate on.** The target architecture. tkv's
``resolve_archs`` narrows a live-GPU build to that one device while a
device-less image bake emits the whole ``TORCH_CUDA_ARCH_LIST`` fat binary,
so the two never agree by construction and gating on them would refuse
every seed forever — a key that moves too easily leaves what it guards
permanently inert. The bake records its arch set for diagnosis; an arch a
fat binary does not carry fails LOUDLY at first launch ("no kernel image is
available"), which is the failure mode that needs no silent guard.

Runnable as a script so the image builder can bake before ``arbi_serve``
is installed; it therefore imports nothing from ``arbi_serve`` at module
scope.

Directory of the installed ``tkv.kernels`` package, or ``None``.

Resolved through ``find_spec`` rather than an import: the boot seed runs
before torch is imported, and importing ``tkv.kernels`` would drag torch
in with it.

Digest of the installed tkv kernel sources, or ``None`` when absent.

Names and bytes are both length-prefixed so no two distinct trees can
hash equal by concatenation. ``None`` means "nothing to compare", which
every caller must treat as unverifiable rather than as a mismatch — an
image whose turbo-attn predates the native int8 kernel has no tkv
kernel tree to digest and must still boot.

``(verdict, detail)`` for one baked extension dir.

:data:`STALE` is returned only on PROOF — a recorded digest that differs
from the installed one. No record, or no installed digest to compare
against, is :data:`UNVERIFIABLE`. An extension outside the tkv family
(exl3, xgrammar, AWQ-Marlin — each keyed on its OWN sources by
:mod:`arbi_serve._prebake_loader`) is :data:`NOT_TKV`: nothing here
judges it either way.

The ``(k_bits, v_bits)`` pairs the turbo_attn_int8 bake must cover.

Each uniform width is run through
``tkv.kernels._precompile_combos.resolve_bw_combos`` — the same resolver
``tkv.kernels.precompile`` walks — so a mounted calibration bundle's
per-layer (possibly asymmetric) allocation is picked up rather than
guessed at, and an image build with no bundle collapses to the diagonal
pair. The codec FAMILY (scalar vs vq2) is not a build axis for this
kernel: it changes the decode table's contents, not the module, so the
pairs are deduped across families.

Every compile-time tile shape the launcher serves, newest contract first.

The tile is a build axis, not a runtime argument: a served call picks one
of these by its own rule (a short-context split-KV verify launch takes a
different build from a prefill launch), and each is a separate ``.so``.
Baking one of two leaves the other cold for exactly the step it was added
to serve, so the set is READ from the launcher rather than assumed.

``SERVED_TILES`` is the launcher's own enumeration. A turbo-attn that
predates it serves one tile, and naming ``DEFAULT_TILE`` for that case is
not a guess — it is that build's only shape.

A short label for one tile, by what it CHANGES from the default.

Derived rather than hardcoded, so a third served tile names itself in the
build log instead of being reported as one of the two that existed when
this was written.

Build every turbo_attn_int8 prefill module a served call can request.

The variant axes are the kernel's own: ``HEAD_DIMS`` x
:func:`turbo_attn_int8_widths` x :func:`served_tiles`. Each point is asked for
through the launcher's ``_module``, the same call a serving dispatch
makes, which is what keeps the bake from drifting from what a boot loads
and why the cache name is never spelled here.

Returns the extension dir names that appeared under
``TORCH_EXTENSIONS_DIR``, so a caller can report exactly what this bake
produced.

Record the bake provenance on every tkv extension dir under ``root``.

Stamps the whole tkv family, not just this bake's dirs: the decode /
MTP-verify kernels baked earlier in the image are seeded by the same
boot path and carry the same hazard, and one guard covering all of them
is one seam rather than two.

Image-builder entry: bake the native int8 prefill kernel, then stamp.

Never fatal. An absent kernel (the turbo-attn pin predates it) and a
failed compile both leave exit 0 with a printed reason, worded so the
two can never be read as each other; the stamp runs regardless, because
it guards the kernels that ARE baked. A bake this loses costs back only
the JIT it was removing, on one default-off route, and the boot that
pays it names the kernel on its own timeline — a failed image build
would be the larger outage.

Thin wrapper around the Rust ``tokenizers`` library + tokenizer_config.json.

We deliberately do not use the HF modeling library (hard project rule).
The HF Rust BPE tokenizer is a separate package (``tokenizers``) and
covers everything Qwen3 needs: encode prompts, decode outputs,
incremental detokenization for SSE streaming, and chat-template
rendering via the bundled Jinja template.

Encode / decode + chat-template rendering. Self-contained.

Loads:
  - ``tokenizer.json`` — Rust-side BPE / decoder / pre-tokenizer.
  - ``tokenizer_config.json`` — eos_token_id, bos_token_id, the
    Jinja chat_template string, and the added_tokens_decoder map
    for special-token IDs.

The ``system`` turns' text, joined — the default ``system_message``.

Templates that take the system prompt as a separate variable (only
NemotronVoiceChat's, today) get it from the standard OpenAI message
list when the caller did not pass one explicitly. Several system turns
join with a blank line, matching how a template that looped over
``messages`` would render them in order.

Build the ``{% generation %}`` block-tag Jinja extension.

HF transformers compiles chat templates with an extension that
registers the ``generation`` tag (used to extract assistant-token
masks during training). Newer instruct / reasoning templates — e.g.
LFM2.5-8B-A1B — wrap the assistant turn in ``{% generation %} ...
{% endgeneration %}``. A bare Jinja2 environment rejects that with
``TemplateSyntaxError: Encountered unknown tag 'generation'``.

For *serving* we don't need the mask indices, only the rendered text,
so the tag is a transparent passthrough: it parses the block body and
emits it unchanged. Defined behind a factory so ``jinja2.ext`` is
imported lazily.

Number of tokens in the tokenizer's vocabulary, added tokens INCLUDED.

The whole id space: ``max(id) + 1`` for a dense vocabulary, which is
what a caller sizing a head or bounding an id needs. It is NOT the
base-vocabulary count some tokenizer libraries publish under this
name — :attr:`base_vocab_size` is that number, and the two differ by
exactly the added tokens.

Number of BPE ids, added tokens EXCLUDED.

The boundary a caller needs to ask whether an id is an added one.
Answered by the backend that owns the vocabulary rather than derived
from :attr:`vocab_size`, which counts the added tokens in.

Ids this checkpoint ADDED above its BPE vocabulary, ascending.

The backend's own added-token decoder is the answer; the
``tokenizer_config.json`` map this object already carries is consulted
only when the backend reports none, because a checkpoint may declare
an added token in the config that its ``tokenizer.json`` does not
repeat.

Pull a special-token string form from the config (eos / bos / pad).

Some chat templates substitute ``{{ bos_token }}`` or
``{{ eos_token }}`` literally — passing ``None`` yields ``"None"``
in the rendered prompt. The HF tokenizer config carries either a
plain string or an ``AddedToken`` dict with a ``content`` field.

Encode many texts in one call, in input order.

The Rust ``tokenizers`` library parallelises ``encode_batch``
internally (rayon), so a single call is far cheaper than N
per-string ``encode`` calls — and, crucially for the embedding
endpoint, it collapses N thread-pool dispatches into one. Returns
one id list per input text.

Return the new text fragment produced by appending the latest token.

``state`` is mutated in-place and carries one of two per-request
algorithms, latched on the FIRST decode (so a mid-request flag
flip never mixes them):

  * **DecodeStream (default, ``ARBI_DECODE_STREAM=1``)** — the
    Rust ``tokenizers.decoders.DecodeStream``: one C-side
    ``step(token_id)`` per fresh token (``state["stream"]`` +
    ``state["consumed"]``). Cheaper per token than the window
    algorithm on the engine thread; identical accumulated text
    (fuzz-verified incl. multibyte + special-token boundaries —
    the stream buffers incomplete UTF-8 in Rust exactly where
    the window path holds back a ``�`` tail).
  * **Two-decode window (fallback / ``ARBI_DECODE_STREAM=0``)** —
    ``state["prefix_offset"]`` / Also the
    automatic path when the installed ``tokenizers`` predates
    ``DecodeStream``.

DecodeStream-backed incremental decode (see :meth:`incremental_decode`).

Feeds every not-yet-consumed token (usually exactly one; the
worker-detok path can catch up several queued tokens in one call)
through the per-request Rust stream. ``skip_special_tokens=True``
matches the window path: the matched stop token stays out of the
user-visible text.

Names the model's chat template reads out of the render context.

Jinja's static analysis of the parsed template; conservative (a
name assigned inside a branch or loop body can also appear). Used
to refuse a client-supplied template variable the loaded model
would ignore — a knob that renders nothing must fail loud, not
silently serve a different prompt than the client asked for.

Empty when the model ships no chat template.

Refuse template variables the loaded model cannot honor.

Two failure modes, both :class:`ValueError` (the routes map it to
400):

* the template never reads the name — the variable would render
  nothing at all;
* the template reads it but rejects the VALUE. The accepted set is
  the model's contract, not ours, so this asks the template
  (:meth:`_probe_template_var`) instead of hardcoding a value list
  that would go stale against the checkpoint.

Render a minimal conversation to ask the template about a value.

Returns the template's own rejection message, or ``None`` when the
value is accepted. DIFFERENTIAL: a probe that also fails WITHOUT
the variable is failing for an unrelated reason (a template that
demands tools, a specific turn order, …) and is not attributed to
the value. Memoised per ``(template, name, value)`` so a request
never pays for more than the first probe of a given value; the
store is an idempotent same-value write, so the tokenizer-pool
worker threads share it without a lock.

Reconcile OpenAI's reasoning-effort ladder with this model's.

``/v1/chat/completions`` is a COMPATIBILITY surface: a client that
speaks OpenAI's ladder (``none`` / ``minimal`` / ``low`` /
``medium`` / ``high``) must not hard-fail on a checkpoint that
spells the same rung differently. No checkpoint serves the whole
ladder and none of them agree on it, so the reconciliation is
LITERAL-FIRST and asks the loaded template at every step rather
than hardcoding one model's vocabulary:

1. the template accepts the client's own spelling — send it
   verbatim, and never consult the table (gpt-oss interpolates the
   value straight into its system prompt, so ``high`` must stay
   ``high`` there, and translating it eagerly would write an
   off-vocabulary rung into the prompt);
2. the template refused it and the word means "do not reason at
   all" — that names no rung, so it moves to the ``enable_thinking``
   channel, which is where a template expresses the same thing;
3. the template refused it and a cross-checkpoint synonym IS
   accepted — send the synonym;
4. otherwise leave it alone, so :meth:`check_chat_template_vars`
   refuses with the template's OWN message.

Only requests that would already be a 400 can change outcome.

Render OpenAI-format messages via the model's chat template.

Mirrors the HF ``PreTrainedTokenizer.apply_chat_template`` surface
we cannot import (no transformers dep — project rule). The Rust
``tokenizers`` library does not ship a chat-template runner of
its own, so we drive Jinja2 directly against the template string
loaded from ``tokenizer_config.json``.

Raises :class:`ValueError` when the model ships no
``chat_template`` — same behavior as vLLM and SGLang. A silent
ChatML fallback would hide model-specific prefixes (e.g.
Qwen3's ``<think>\n\n</think>`` thinking-disabled marker)
and produce subtly wrong output for any model whose template
doesn't happen to be plain ChatML. ``tools`` and ``thinking``
are passed through to the Jinja env so tool-calling-aware
templates (Qwen3) can use them. ``template_vars`` carries any
further client-chosen template variables and is vetted by
:meth:`check_chat_template_vars` before it reaches the render.
``system_message`` is for NemotronVoiceChat's override template
(``arbi_serve.chat_templates.NEMOTRON_VOICECHAT_CHAT_TEMPLATE``),
which takes the system prompt as its OWN variable rather than
reading it out of ``messages`` — the realtime path has no message
list to read it from. Every other template ignores the kwarg
(Jinja does not error on an unused context variable).

An omitted ``system_message`` is resolved from ``messages``' own
``system`` turns (joined, in order) rather than left at ``None``:
a Jinja ``{{ system_message }}`` renders ``None`` LITERALLY, so
the ``/v1/chat/completions`` path — which passes messages and
never this kwarg — used to prefix every NemotronVoiceChat prompt
with the four characters ``None``. Falls back to the empty string
when there is no system turn either, which the template already
special-cases (no blank line before the tool block).

Weight-quantization backend registry + generic loader.

Distinct from KV-cache quantization (which lives under
:mod:`arbi_serve.cache` and the TKV / bf16 attention backends).
This package is about replacing dense projection ``Linear`` layers
with their quantized mirrors at load time — the rest of the engine,
including KV codec choice, is unchanged.

Public entry points (each arch's ``from_safetensors`` uses two):

  - :func:`prepare_dims_for_quant` — detect backend + apply lm_head
    untie before model construction.
  - :func:`apply_quant_if_present` — swap dense Linears to their
    quantized mirrors after construction, before dense weight load.

Per-arch ``weight_map`` routes through :func:`filter_weight_map_for_quant`
to drop ``.weight`` entries for swapped linears.

Backends register themselves via :mod:`arbi_serve.weight_quant.registry`.
Bundled: EXL3 (trellis), AWQ (INT4 packed), NVFP4, FP8 (per-tensor +
per-channel float8_e4m3fn / float8_e5m2). A new backend registers itself
under ``arbi_serve/weight_quant/<name>/``; no arch code changes are
needed for it to start activating.

AWQ (Activation-aware Weight Quantization) backend for arbi-serve.

Importing this subpackage registers the backend with the global
quant-backend registry. The Triton dequant kernel is JIT-compiled on
the first forward; the loader path itself is pure-PyTorch and runs at
import time without GPU. A vanilla bf16 run never imports any of this
module's compute paths.

The backend handles INT4 packed AWQ checkpoints in the canonical AWQ
on-disk shape::

    qweight  int32   [in_features, out_features // 8]   (8 weights/int32)
    qzeros   int32   [in_features // group_size, out_features // 8]
    scales   bf16/fp16 [in_features // group_size, out_features]

Detection key is ``*.qweight``. The same shape is shared by
casperhansen-style ``quant_method: awq`` checkpoints and
``llm-compressor``-style ``quant_method: compressed-tensors`` packs
when ``format: pack-quantized`` and ``num_bits: 4`` (the on-disk byte
layout is identical; only the metadata wrapping differs).

Base class for the AWQ weight-quantization Linears.

:class:`AWQLinearBase` is the common surface (weight buffers, Marlin
repack, forward, dequant, fusion, readout) shared by the parallel
variants in :mod:`arbi_serve.weight_quant.awq._parallel`. The Marlin
kernel-dispatch predicates live in
:mod:`arbi_serve.weight_quant.awq._marlin_dispatch`.

``linear.group_size`` where it can be derived, else ``None``.

The property computes the group size FROM the bound scales and raises on a
linear that has not been bound yet. The warm-reload a8 check runs before
that is guaranteed, and it must not turn a provenance question into an
AttributeError/RuntimeError. Module-level and duck-typed on purpose: the
check is exercised against lightweight doubles that carry the a8 marker
without the whole scale machinery, and requiring them to grow a method
would be the test tail wagging the dog. An undeterminable group size means
the gs16 gate cannot apply — the same answer the pre-gate code gave.

Per-token e4m3 quant for the W4A8 GEMMs: returns ``(x_fp8,
a_scales)`` with ``a_scales`` shaped ``(M, 1)`` float32, as
``marlin_gemm`` requires for 8-bit activations.

``pre_scale`` (``(K,)`` or None) is the SmoothQuant per-input-channel
``act_pre_scale = 1/s``: for the no-norm sensitive projections
(``down_proj`` / ``o_proj`` / ``out_proj``) the smoothed weight
``W·diag(s)`` is served, so the activation must be multiplied by ``1/s``
per channel BEFORE the per-token quant to keep the product exact. The
multiply flattens the per-channel outliers the per-token scale otherwise
burns its range on — that is the whole point of the recipe.

Plain tensor ops on purpose — the AWQ projections run inside
``torch.compile`` regions (GDN ``in_proj``, attention q/k/v), where
Inductor fuses this (the per-channel multiply included) into a couple of
kernels ahead of the opaque Marlin op; it is cudagraph-safe (no host sync,
no allocation surprises, ``pre_scale`` is a resident broadcast buffer).
Mirrors ``Fp8LinearBase._fp8_forward_triton_fused``.

Whether W4A8 is ON *by default* (no explicit ``ARBI_SERVE_AWQ_NO_A8``)
on an fp8-MMA-capable arch. Pure over the compute capability so it is
unit-testable without CUDA.

``True`` wherever the fp8-activation Marlin kernels exist: sm89 (Ada /
RTX 4090) and the sm12x family.

sm89 was ``False`` from 2026-07-19 to 2026-09-01. That default came from a
real collapse (accept 1.10, incoherent output) which was misattributed to
the per-token e4m3 format floor. It was root-caused five days later to a
single component — fp8 activations on the GDN ``in_proj_qkv`` projection
corrupting the fp32 SSM recurrent state — and fixed by
``loader._A8_GDN_EXCLUDE_DEFAULT`` ("proven minimal fix, bisected on
Qwen3.6-27B"), which has shipped ever since. With that exclusion a8 was
measured lossless on the very model that collapsed, twice: accept 4.30 vs
4.28 (W4A16), and 4.47/4.43 vs 4.48/4.46 with prefill-KL <= 7.9e-3 against
a 0.0 two-boot null floor.

The flip resolves issue #1113 UNCONDITIONALLY rather than
concurrency-gating it. Measured trade: prefill TTFT -12% (c1) / -29% (c8)
and c8 decode aggregate +12.6%, against ONE regression — the c1
short-decode corner at -4.9%, where fp8 activations do not cut the INT4
weight bytes and the GEMM is memory-bound on a skinny M. Gating that
corner was the alternative; it was rejected because it buys ~5% in one
corner at the cost of another dispatch boundary, and dispatch boundaries
in this codebase have repeatedly cost more than they returned.

``ARBI_SERVE_AWQ_NO_A8=1`` restores W4A16 everywhere if the c1 corner
matters for a given deployment.

The W4A8 decision for this module on this device, under this boot's flags.

One resolver for both seams that need the answer: the cold repack, which
BAKES it into the weight bytes (the a8 nibble remap folds the zero-points
in and the group scales are pre-multiplied by 512), and the warm reload's
:meth:`AWQLinearBase.rebind_after_compaction`, which has to decide whether
the bytes a dump restored were baked under the same decision this boot
would make. A second copy of the rule is how a warm boot comes to serve
W4A8 bytes through the W4A16 kernel.

``ARBI_SERVE_AWQ_NO_A8`` is tri-state: unset ⇒ the arch default, ``1``
forces W4A16, ``0`` forces W4A8.

``group_size`` gates the one combination that has no kernel. The a8 matmul
consumes 2 k-blocks per step and rejects ``group_blocks < 2``, so there is
no ``kFE4M3fn`` instantiation at ``group_blocks == 1`` — of 192 a8 selector
entries the group_blocks histogram is {-1: 48, 2: 48, 4: 48, 8: 48}, with
NOTHING at 1. A group-16 pack therefore repacks a8 at load and then dies in
kernel selection at the first forward: fail-loud, but deferred from boot to
serve time. Route it to W4A16 here instead, which is what
``_marlin_supports`` already documents should happen ("other fp8-MMA archs
must not route a gs16 pack through W4A8") without anything enforcing it.
Pass ``None`` only where the group size genuinely is not yet known.

Single Marlin GEMM entry for a 2-D ``(M, K)`` input — the one place
the W4A16 and W4A8 dispatch lives, shared by the per-linear forward
and the fused-projection path.

W4A8 quantizes activations per token, passes ``a_scales`` and forces
``use_atomic_add=False`` (the atomicAdd reduce is not wired for the
fp8 epilogue), and casts the 16-bit-scale output back to ``out_dtype``.
``a8_pre_scale`` is the optional SmoothQuant per-channel ``act_pre_scale``
for a weight-fold projection. W4A16 passes activations through and consults
the atomicAdd gate.

Common surface for the three AWQ parallel variants.

Carries the weight buffers (persistent — captured in the warm-reload
flat dump), the Marlin-laid-out repacked weight + scales + workspace,
and a one-bit ``_marlin_active`` flag indicating whether forward goes
through Marlin (default) or the torch fallback.

The :attr:`_format` flag records the on-disk format the linear was
bound to:

  - ``"legacy"`` — populated by :meth:`awq_load`, repacks via
    ``awq_marlin_repack`` directly (legacy qweight layout matches
    what that kernel expects).
  - ``"pack_quantized"`` — populated by :meth:`pack_quantized_load`,
    repacks via ``pack_quantized_to_marlin`` (W4A16, chunked over the
    output axis) or ``pack_quantized_to_standard`` →
    ``gptq_marlin_repack`` (W4A8, whole-matrix).

The Marlin python flags + workspace are DERIVED state (not buffers);
:meth:`rebind_after_compaction` rebuilds them after a warm reload
restores the persistent buffers.

Set the python ``_format`` string AND mirror it (plus the W4A8
bit) into the persistent ``_format_marker`` buffer so a warm
flat-dump reload can recover both (the freed raw buffers can't
disambiguate legacy from pack_quantized, and the Marlin buffers
can't reveal whether the a8 nibble remap was baked in — see
:meth:`rebind_after_compaction`).

Refuse a restored W4A8 decision this machine/boot cannot honour.

``_marlin_is_a8`` comes back from the dump's ``_format_marker``, but it
is a CAPABILITY decision, not a weight: W4A8 exists only where the
fp8-activation ``ARBI_SERVE_AWQ_NO_A8`` / ``--activation-quant`` asks for it.
The repack BAKED that decision into the bytes — the a8 path folds the
zero-points into the nibbles and pre-multiplies the group scales by
512 — so a dump and a boot that disagree cannot be reconciled by
dispatching differently; the bytes are simply for the other kernel.

A flat dump is therefore NOT architecture-agnostic for W4A16 weights.
Same silence in the other direction, on ONE arch: an operator asking
for bf16 activations warm-loads an a8 dump and serves fp8.

Rebuild the Marlin DERIVED state after a warm flat-dump reload.

A warm reload restores the persistent buffers (raw + Marlin +
``_format_marker``) but leaves the python flags + workspace at
their ``__init__`` defaults (the meta graph never ran ``awq_load``
/ ``_try_marlin_repack``). Reconstruct them from buffer presence so
the forward path is byte-correct WITHOUT re-running the repack:

  * ``_format`` ← the persisted marker (unambiguous; buffer
    presence alone can't tell legacy from pack_quantized after the
    raw buffers were freed on the Marlin path),
  * ``_marlin_active`` ← whether ``_marlin_qweight`` is populated
    (assigned only on a successful repack, never re-emptied),
  * ``_marlin_has_zp`` ← whether ``_marlin_zeros`` is populated
    (the no-zp repack leaves it empty),
  * ``_marlin_workspace`` ← created ONLY when absent/wrong (SM-count
    zeros; cheap, non-persistent) — an EXISTING valid workspace is kept
    in place (see below),
  * ``retain_dense_for_readout`` ← whether the readout weight is
    present.

Also fired post-compaction on the COLD path, where every value is
already correct so this is an idempotent no-op. A no-op while the
buffers are still empty placeholders (marker == unbound).

Repack the per-format raw buffers into Marlin's layout.

Returns True if Marlin is now wired up and forward should route
through ``marlin_gemm``; False if the shape falls outside
Marlin's tile constraints (caller routes to dequant fallback).

Bind the compressed-tensors pack-quantized tensors.

Subclasses override for TP slicing; at TP=1 every subclass calls
straight through.

On-disk shapes (sliced for TP at the subclass layer before
getting here):
  - weight_packed: ``(out_features, in_features // 8)``, int32
  - weight_scale: ``(out_features, in_features // group_size)``, fp
  - weight_zero_point (optional, asymmetric): ``(out_features // 8,
    in_features // group_size)``, int32 (packed along OUTPUT axis)

Bind the SmoothQuant per-input-channel ``act_pre_scale = 1/s`` for a
weight-fold projection (called by the loader after the weight load).

Validates the ``(in_features,)`` shape and fails loud if this linear is
served on the W4A16 Marlin kernel — the pack's weight is ``W·diag(s)``,
so serving it without multiplying the activation by ``1/s`` (which only
the a8 path does) would be silently wrong. An ``--a8-smooth`` pack MUST
be served with ``ARBI_SERVE_AWQ_NO_A8=0``. On the CPU/dequant fallback
(Marlin inactive) the pre-scale is applied in :meth:`_awq_forward`.

True if ``self`` and ``other`` can be fused into one column-
concatenated Marlin GEMM.

Two AWQ linears that share the same K (in_features), group_size,
has_zp flag and are both Marlin-active can be fused by
concatenating their already-repacked weight/scale/zero tensors
along the OUTPUT (N) axis. This is bit-exact: Marlin's repacked
column layout is N-tile-local (tiles of 64), so cat along the
column axis equals repacking the concatenated standard weight
(verified numerically). Cuts the per-decode kernel + launch count
for adjacent column-parallel projections (e.g. GDN
``in_proj_qkv`` + ``in_proj_z``).

Concatenate a list of Marlin-active AWQ linears into one fused
GEMM descriptor, or ``None`` if they are not all fusable.

Returns a dict carrying the concatenated Marlin tensors and the
total output width, consumed by :meth:`marlin_gemm_fused`. The
concatenation is along the OUTPUT (N) axis and is bit-exact (see
:meth:`marlin_fusable_with`).

Return the dense ``(out, in)`` weight for non-GEMM consumers.

For MLA ``kv_b_proj`` the backend needs the raw matrix, not the
Marlin GEMM. On the Marlin path the dense weight was captured at
load time (see :meth:`_try_marlin_repack`) into
``_readout_dense_weight``; on the non-Marlin fallback path it is
dequantized on demand from the still-present raw buffers. Raises
if :attr:`retain_dense_for_readout` was not set before load (the
Marlin path freed the int4 buffers, so the weight is
unrecoverable).

Dense ``(out, in)`` weight via the pure-torch dequant (no Triton).

Same numerics as :meth:`dequantized_weight` but pins the torch
reference dequant — avoids the Triton-JIT deadlock the fallback
path hit at first launch. Only valid on a non-Marlin linear (the
raw int4 buffers are still present).

Dequantize to the dense ``(out_features, in_features)`` weight in
``F.linear`` layout (row-major over output).

Same math as the dequant-fallback arm of :meth:`_awq_forward`; that
arm calls through here so the two stay byte-identical. Used to build
fused dense ``[B | A]`` weights for sub-tile GDN projections that
route to the dequant path (out width below Marlin's tile bound) —
there is no Marlin descriptor to concatenate, so we fuse the dense
form.

Raises (never silently degrades) when the linear is Marlin-active —
a Marlin-repacked linear has freed its dense buffers, so the caller
must use the Marlin fusion path (:meth:`build_fused_marlin`) instead.

Marlin kernel-dispatch helpers for the AWQ weight-quant linears.

The Marlin tile-constraint predicate and the ``use_atomic_add`` cross-CTA
reduce gate (with its flag-truth counter). The AWQ ``out_proj`` /
``down_proj`` forwards consult :func:`_should_use_atomic_add_reduce` on
every GEMM, and :func:`_marlin_supports` gates the one-time repack.

The ``ARBI_MARLIN_USE_ATOMIC_ADD`` flag, read as a per-process constant.

Isolated into a 0-arg ``assume_constant_result`` helper so a Dynamo trace
that reaches :func:`_should_use_atomic_add_reduce` — the AWQ Marlin
``out_proj`` runs inside compiled model forwards (e.g. the GDN prefill
path) — bakes the flag instead of tracing ``RuntimeFlags.from_env``
(``dataclasses.fields`` is untraceable on torch-2.12 when the runtime-flags
cache is disabled).

Whether the Marlin cross-CTA reduce should use atomicAdd.

Returns False unless ``n < 2048``, ``k >= 2048``, the device is CUDA,
and NOT support atomicAdd + bfloat16, so it is force-disabled there for
bf16 even when the flag is set.

True if the (in, out, group_size) shape satisfies Marlin's tile
constraints. AWQ-marlin requires:
  - in_features divisible by ``GPTQ_MARLIN_MIN_THREAD_K = 128``
  - out_features divisible by ``GPTQ_MARLIN_MIN_THREAD_N = 64`` AND
    by ``tile_size = 16``
  - group_size in {-1, 16, 32, 64, 128}
Anything else routes to the dequant fallback.

group_size 16 (group_blocks=1) serves on the W4A16 kernels only; the
W4A8 (fp8-activation) matmul consumes 2 k-blocks per step and rejects
group_blocks<2. sm89 (RTX 4090) already defaults to W4A16, so this is
the served path there; other fp8-MMA archs must not route a gs16 pack
through W4A8.

AWQ equivalent of :class:`ReplicatedLinear`.

Used for the lm_head when an AWQ export quantizes it (most do —
AWQ's quantizer treats lm_head as a regular Linear unless told to
skip via ``modules_to_not_convert``). Most public AWQ packs leave
lm_head DENSE; in that case the EXL3-style untie isn't required and
the registered backend's ``lm_head_untie_required()`` returns False.

``__init__`` + ``forward`` come from :class:`ReplicatedLinearMixin`.

AWQ equivalent of :class:`ColumnParallelLinear`.

Output dim is sharded across the TP group. ``out_features`` is the
GLOBAL output size; this class stores the LOCAL slice and slices
the AWQ tensors along the output axis inside the per-format load.

Legacy AWQ TP slicing math (qweight is ``(in, out)``-major):
  - ``qweight: (in_features, out_features // 8)`` →
    slice dim 1: ``[:, rank * (local_out // 8) : (rank+1) * (local_out // 8)]``
  - ``qzeros:  (in // group_size, out_features // 8)`` →
    slice dim 1: same as qweight
  - ``scales:  (in // group_size, out_features)`` →
    slice dim 1: ``[:, rank * local_out : (rank+1) * local_out]``

Pack-quantized TP slicing math (weight_packed is ``(out, in)``-major):
  - ``weight_packed: (out_features, in_features // 8)`` →
    slice dim 0: ``[rank * local_out : (rank+1) * local_out, :]``
  - ``weight_scale: (out_features, in_features // group_size)`` →
    slice dim 0: same as weight_packed
  - ``weight_zero_point: (out_features // 8, in // group_size)`` →
    slice dim 0 in pack-factor units (rank * (local_out // 8))

Constraint: ``out_features % (tp_size * 8) == 0`` so each rank's
local out aligns to a full int32 packed nibble boundary
(:attr:`tp_block_align` = 8). ``__init__`` + ``forward`` come from
:class:`ColumnParallelMixin`.

Multi-GPU end-to-end is unverified on this build (single-GPU box);
the slicing math is unit-tested.

AWQ equivalent of :class:`~arbi_serve.models.linear.VocabParallelLMHead`.

Vocab-aware column-parallel head: ``out_features`` is the PADDED
vocab (a tp_size multiple — same convention as the dense head it
mirrors), ``org_vocab_size`` is the real vocab, and :meth:`forward`
trims the padded logit columns after the gather so they never leak
into sampling. For a tp-divisible vocab (padded == org, e.g.
Qwen3.6-27B's 248320 at TP2) the trim is a no-op.

Two ways this class comes to exist:

  - checkpoint-shipped quantized lm_head: the swap loop maps the
    dense ``VocabParallelLMHead`` here via
    :meth:`AWQBackend.quant_class_for` and binds the on-disk GLOBAL
    tensors through the inherited TP-slicing load methods;
  - load-time RTN head quantization
    (:mod:`arbi_serve.weight_quant.head_quant`): each rank quantizes
    its ALREADY-LOCAL dense vocab shard and binds it via
    :meth:`pack_quantized_load_local` (bit-identical to slicing a
    globally quantized tensor — RTN groups along the input axis, so
    rows are independent).

The drafter's distributed greedy argmax
(:meth:`Qwen3_5MtpHead._distributed_greedy_argmax`) consumes
:meth:`local_logits` — the rank-local shard row WITHOUT the gather —
and reconstructs the global winner from ``(max, idx)`` pairs; that
path duck-types on this class's ``local_logits`` +
``org_vocab_size`` surface.

AWQ equivalent of :class:`RowParallelLinear`.

Input dim is sharded across the TP group. ``in_features`` is the
GLOBAL input size; this class stores the LOCAL slice and slices
along the input axis inside the per-format load method.

Legacy AWQ TP slicing math (qweight is ``(in, out)``-major):
  - ``qweight: (in_features, out_features // 8)`` →
    slice dim 0: ``[rank * local_in : (rank+1) * local_in, :]``
  - ``qzeros:  (in // group_size, out_features // 8)`` →
    slice dim 0 along group rows: each rank gets local_in/group_size groups
  - ``scales:  (in // group_size, out_features)`` → same row-slice as qzeros

Pack-quantized TP slicing math (weight_packed is ``(out, in)``-major):
  - ``weight_packed: (out_features, in_features // 8)`` →
    slice dim 1 in pack-factor units (8 nibbles per int32):
    ``[:, rank * (local_in // 8) : (rank+1) * (local_in // 8)]``
  - ``weight_scale: (out_features, in_features // group_size)`` →
    slice dim 1 in groups: ``[:, rank * (local_in // gs) : (rank+1) * (local_in // gs)]``
  - ``weight_zero_point: (out_features // 8, in // group_size)`` →
    slice dim 1 same as weight_scale

Constraint: ``in_features % (tp_size * group_size) == 0`` so each
rank's local in aligns to whole groups (per-group scale validity)
AND aligns to pack-factor=8 (int32 boundary). Since ``group_size``
is unknown at construction time, we re-check the looser
``__init__`` +
``forward`` come from :class:`RowParallelMixin`.

AWQ equivalent of :class:`MergedColumnParallelLinear`.

The dense merged-linear stacks multiple column-parallel
sub-projections along the OUTPUT axis with potentially-distinct
widths (e.g. GDN ``in_proj_qkv``: ``[Q(key_dim) | K(key_dim) |
V(value_dim)]``; gate_up_proj: ``[gate(intermediate) |
up(intermediate)]``). The safetensors stores ONE fused tensor
(e.g. ``in_proj_qkv.weight_packed``) whose output dim is
``sum(_shard_out)``.

At TP=1 the existing :class:`AWQColumnParallelLinear` handled this
cleanly because the fused tensor maps to a single column-parallel
AWQ Linear with ``out_features = sum(_shard_out)``. At TP>1 a
naive contiguous output-dim cut splits mid-sub-block (rank 0
receives ``[full Q + half K]`` at TP=2 with QKV widths
``[key_dim, key_dim, value_dim]``), corrupting the per-shard
semantics consumers downstream depend on (the GDNBlock splits
the output of ``in_proj_qkv`` by ``[key_dim_local,
key_dim_local, value_dim_local]``).

The fix: shard EACH sub-block independently across TP ranks,
then concatenate the per-rank slices in declared order. This
matches how :class:`MergedColumnParallelLinear`'s dense
weight_loader stacks per-rank sub-shards.

Per-sub-block constraint: each sub_dim must be divisible by
``tp_size * 8`` (the AWQ INT4 pack boundary). The constructor
asserts this so we fail loudly on a bad model shape rather
than silently truncating mid-int32 nibble.

The fused output is ``out_features = sum(of // tp_size for of in
shard_out_features)``; downstream consumers' ``.split([Q_local,
K_local, V_local])`` produces the right per-shard slices because
the local fused dim mirrors the dense merged-linear's local
weight layout.

Rank-local ``(B, per_shard)`` logit shard — no gather, no trim.

Exactly the tensor the gather path concatenates (same kernel,
same input), so each row is bit-identical to its slice of the
gathered full-vocab row. Padded-tail exclusion is the CALLER's
job (via ``org_vocab_size`` — see the drafter's distributed
greedy argmax).

Bind ALREADY rank-local pack-quantized tensors.

The inherited :meth:`pack_quantized_load` slices GLOBAL
checkpoint tensors by rank; the load-time RTN path quantizes the
rank's local dense shard directly, so its tensors must bind
WITHOUT re-slicing. Delegates straight to the base-class binder
(shape-validated against the LOCAL ``out_features``).

Slice the global SmoothQuant ``act_pre_scale`` (in_features,) to this
rank's input shard before binding — the weight-fold projections
(down/o/out_proj) are row-parallel, so their input channels are split
across ranks exactly like ``weight_packed``'s input axis above.

Slice ``tensor`` for one sub-block at this rank.

``unit`` rescales the offsets / widths into the on-disk
granularity of ``tensor`` along ``axis``: 1 for un-packed
tensors (scales / weight_packed in the column dim of packed
layouts), :data:`_AWQ_PACK_FACTOR` for tensors packed along
``axis`` (qweight / qzeros / weight_zero_point packed
directions). The caller passes the un-packed offset / size in
OUTPUT units; this helper divides by ``unit`` to land on the
int32-aligned slice.

Per-sub-block + per-rank slice of pack-quantized tensors
along the OUTPUT axis.

Pack-quantized layout: ``weight_packed: (out, in//8)``,
``weight_scale: (out, num_groups)``, ``weight_zero_point:
(out//8, num_groups)`` (packed along output axis 0). All three
slice on the OUTPUT dim (axis 0 for ``weight_packed`` /
``weight_scale`` / ``weight_zero_point``).

Activation-weighted MSE-optimal INT4 clip search (calibrated head quant).

The data-free RTN packer (:func:`rtn_pack_quantized_int4`) sets each
group's scale from the raw min/max — accept-unsafe on a large vocab head
(the resulting rel-L2 error flips argmax outcomes, dropping drafter
accept-length). This module
computes an ACTIVATION-AWARE per-group scale instead, keeping the EXACT
pack-quantized / marlin vehicle (same layout, same kernel, same runtime
— NO activation fold, NO extra GEMM), so the bandwidth win is unchanged
and only the scale/clip selection improves.

Method (why it is the right activation-aware objective for a weight-only
INT4 emit that must not touch the runtime):

  The output perturbation from quantizing ``W`` is ``ΔW · x``. For an
  output row ``i``, ``E[(ΔW_i · x)^2] = Σ_j ΔW_ij^2 · E[x_j^2]`` when the
  input channels are ~uncorrelated — i.e. the quantity to minimize is the
  weight-MSE WEIGHTED by the per-input-channel mean-square activation
  ``s2_j = E[x_j^2]`` captured over the calibration basket. Channels the
  real hidden state actually excites are protected; dead channels are
  spent. This matches the checkpoint trunk's own ``observer: "mse"``
  (config_groups.group_0.weights.observer) — the trunk was MSE-clipped,
  so the head now is too, just activation-weighted.

  Per group, per output row, we grid-search a symmetric shrink ``c`` of
  the zero-inclusive ``(wmin, wmax)`` range and pick the ``c`` minimizing
  the activation-weighted squared dequant error. ``c = 1`` recovers RTN,
  so the search can only improve (or tie) the weighted objective.

``act_importance`` is the ``(in_features,)`` per-input-channel weight
(``E[x_j^2]`` — or any non-negative importance); ``None`` falls back to
uniform (plain MSE clip, still better than raw min/max). The search runs
group-by-group to bound memory and is GPU-friendly (vectorized over
rows × grid within a group).

Activation-weighted squared dequant error per (row, candidate).

Quantizes ``w_g`` under each candidate's ``(wmin_c, wmax_c)`` using the
STORED-dtype scale/zp (bit-identical to the emit path), dequants, and
returns ``Σ_j imp_g[j] · (deq - w)^2`` — shape ``(out, n_grid)``.

Activation-weighted MSE-optimal INT4 pack-quantized emit.

Returns the same ``(weight_packed, weight_scale, weight_zero_point)``
triple as :func:`rtn_pack_quantized_int4` (asymmetric, group along the
input axis) — drop-in for :meth:`AWQLinearBase.pack_quantized_load`.

``act_importance``: ``(in_features,)`` non-negative per-input-channel
weight (``E[x_j^2]`` from the calibration capture). ``None`` → uniform
(plain MSE clip). Rows are quantized independently, so a per-rank vocab
shard may be quantized locally (bit-identical to slicing the global).

``n_grid`` clip candidates are searched per group over shrink factors
``linspace(min_shrink, 1.0, n_grid)`` applied to the zero-inclusive
range; ``c = 1`` == RTN, so the picked scale never loses to RTN on the
weighted objective.

Real AWQ per-input-channel scaling + fold, on the standard pack layout.

The shipped :func:`arbi_serve.weight_quant.awq.act_scale.awq_pack_quantized_int4`
is a SIMPLIFIED AWQ: an activation-weighted clip search with NO per-channel
scaling. This module adds the scaling step AWQ's accuracy comes from.

For a linear ``y = W x`` a per-input-channel scale ``s`` is chosen so that
quantizing ``W diag(s)`` (columns scaled up on salient channels) and dividing
the input by ``s`` leaves the product unchanged while spending fewer bits on
the channels the calibration activations actually excite. ``1/s`` folds into
the preceding RMSNorm weight (Gemma ``input_layernorm`` → q/k/v,
``pre_feedforward_layernorm`` → gate/up), so the served pack stays a standard
compressed-tensors W4 group-quant (``weight_packed`` int32 / ``weight_scale``
bf16 / ``weight_zero_point`` int32) — portable to any AWQ-capable engine.

``s`` is searched on the SUM of activation-weighted output MSE over the group
of linears that share one input tensor (q/k/v share ``input_layernorm``; gate/up
share ``pre_feedforward_layernorm``), so one folded ``s`` serves the whole group.

Per-group clipped ``(wmin, wmax)`` for the pack emit.

``n_grid <= 1`` → raw zero-inclusive min/max (RTN). Otherwise a shrink
search: symmetric shrinks both ends by one factor; ``asymmetric`` searches
the ``wmin`` and ``wmax`` shrink factors independently (coordinate order:
max then min). ``act_importance`` ((in,) non-negative) weights the
per-channel error; ``None`` → uniform (plain MSE clip).

AWQ per-input-channel scale ``s`` (in,) for a shared-input weight group.

Grid-searches ``alpha in [0, 1]`` (autoawq-style), building
``s = act^alpha / w^(1-alpha)`` normalized around 1, and returns the ``s``
minimizing the SUM over the group of the activation-weighted output MSE
``|x (Wq/s)^T - x W^T|^2`` — the real AWQ objective (RTN rounding inside the
search; clip is applied afterward at emit).

``group_weights``: list of ``(out_i, in)`` fp32 weights sharing input
``x_sample`` ((N, in) fp32 calibration inputs). ``act_importance``: (in,)
per-channel activation magnitude (E[|x|] or sqrt(E[x^2])).

Emit the standard pack-quantized triple for ``W diag(s)``.

The caller folds ``1/scale_s`` into the preceding RMSNorm. The clip search
runs on the SCALED weight with the SCALED activation importance
(``act_importance / s^2`` — ``E[(x/s)^2]``), so scaling and clip compose.
Output layout is byte-contract-identical to the shipped RTN/MSE/AWQ pack.

AWQ :class:`QuantBackend` implementation.

The class implements the protocol — detection from the safetensors keys,
mapping HF-canonical safetensors paths to the AWQLinear mirror class,
and reading the per-linear tensors.

Two on-disk layouts are accepted:

  1. **Legacy AWQ** (``casperhansen/AutoAWQ`` and older
     ``llm-compressor`` exports tagged ``quant_method: awq``):
     ``{qweight (i32), qzeros (i32), scales (fp16/bf16)}`` — qweight is
     ``[in, out//8]`` packed along the OUTPUT axis with the
     ``[0, 4, 1, 5, 2, 6, 3, 7]`` reverse-order permutation. Detection
     key: any ``*.qweight`` paired with matching ``*.qzeros`` + ``*.scales``.

  2. **Compressed-tensors pack-quantized** (the format every modern AWQ
     checkpoint uses — Qwen 3.5 / 3.6, Llama 3 AWQ, Mistral AWQ, etc.;
     produced by ``llm-compressor`` with ``format: pack-quantized``):
     ``{weight_packed (i32), weight_scale (fp16/bf16), weight_shape (i64),
     weight_zero_point (i32, asymmetric only)}`` — weight_packed is
     ``[out, in//8]`` packed along the INPUT axis sequentially.
     Detection key: any ``*.weight_packed`` paired with matching
     ``*.weight_scale`` + ``*.weight_shape``.

The two layouts are mutually exclusive (different key suffixes) so the
backend simply tries pack-quantized detection first, then legacy. Mixed
checkpoints are unheard of in the wild.

Per-linear contract:
  - lm_head DOES NOT need to be untied — most AWQ packs leave lm_head
    dense (it lives outside ``modules_to_convert``). When a pack does
    quantize lm_head, the swap picks it up via the same detection
    without any special-case wiring.
  - Output-axis pack boundary is 8 nibbles per int32. Group-axis stride
    along the input dim is read at load time from the scales shape.

AWQ INT4 weight-quantization backend (group-quantized, packed).

Handles two on-disk layouts: legacy AWQ and compressed-tensors
pack-quantized. A given checkpoint always uses one or the other;
detection runs the pack-quantized check first since that's the
format every modern public AWQ checkpoint ships in.

Read this backend's tensors and bind them onto ``lin`` via the
appropriate format-specific load method.

Per-linear format detection: a linear is pack-quantized if its
``<path>.weight_packed`` exists; otherwise legacy AWQ. Mixing
formats within a single checkpoint is unsupported (and unheard
of in the wild) — each linear is one or the other.

Bind legacy AWQ ``{qweight, qzeros, scales}``.

``dtype=None`` is critical for qweight + qzeros (int32 packed
nibbles). ``scales`` is read at ``compute_dtype``, NOT at the
export's dtype: the Marlin GEMM reinterprets the group scales
through the activation dtype's pointer type, so an fp16 export
served by a bf16 engine (or the reverse) is silently wrong.

Bind compressed-tensors pack-quantized
``{weight_packed, weight_scale, [weight_zero_point]}``.

``weight_scale`` is read at ``compute_dtype`` for the same reason
as the legacy path (see :meth:`_bind_legacy_awq`).

``weight_shape`` is metadata used to verify in/out at load time;
we already track these on the Linear instance so we don't need
to re-read it (but the validator inside
:meth:`AWQLinearBase.pack_quantized_load` cross-checks shapes).

``int4_w4a16`` + the pack's group size when EVERY routed-expert
projection is legacy AWQ with an explicit zero-point, else ``None``.

The fused kernel reads one uint8 expert stack per stage, so the
whole expert set must share one layout and one group. Anything
else — a compressed-tensors ``weight_packed`` expert, a symmetric
pack with no ``qzeros``, a W4A8 projection carrying
``act_pre_scale`` (whose activation is e4m3, which
``fused_moe_kernel_gptq_awq`` does not consume), a ragged group —
keeps the per-expert modules.

AWQ INT4 dequantize — Triton kernel + torch reference.

Two paths share one numerical contract::

    W[in, out] = (uint4(qweight[in, out//8] >> shift[out%8])
                  - uint4(qzeros[in//gs, out//8] >> shift[out%8]))
                 * scales[in//gs, out]

The "shift" lookup is the AWQ-specific reverse-order permutation
``[0, 4, 1, 5, 2, 6, 3, 7]`` (for the eight 4-bit lanes packed into one
int32). This is the load-bearing piece — getting the order wrong gives
non-trivial output that looks plausible but is silently scrambled.

Triton path: a port of vLLM's AWQ Triton dequant kernel — used when the
Marlin GEMM isn't available.

Torch path: pure-tensor dequant for unit tests + CPU smoke. Slow but
device-agnostic; matches the Triton output bit-for-bit on integer
arithmetic (tested in ``test_weight_quant_awq.py``).

Unpack an AWQ int32 grid ``(rows, cols // 8)`` to ``(rows, cols)``.

Returns int32 codes in ``[0, 15]``. The i-th output lane lives at bits
``4*REVERSE_AWQ_ORDER[i] .. +3`` of the int32, so the shift table
undoes the AWQ interleave as it unpacks.

Dequantize AWQ INT4 packed weights using pure torch ops.

Returns a dense matrix of shape ``(in_features, out_features)`` in
``scales.dtype``. Slow vs the Triton kernel but bit-identical and
device-agnostic — used on CPU + as the parity oracle in tests.

Args:
  qweight: int32 packed weights, shape (in_features, out_features // 8).
  qzeros:  int32 packed zeros,   shape (in_features // group_size, out_features // 8).
  scales:  fp16/bf16 scales,     shape (in_features // group_size, out_features).

Layout invariants (raise on mismatch):
  - ``qweight.shape[0]`` divides ``scales.shape[0]`` exactly (group rows).
  - ``qzeros.shape == (scales.shape[0], scales.shape[1] // 8)``.
  - All three on the same device + qweight/qzeros are int32.

Dequantize AWQ INT4 weights to a dense (in, out) matrix.

Picks the Triton kernel when running on CUDA + Triton is available,
otherwise falls back to :func:`awq_dequantize_torch`. Both paths
are bit-identical on integer arithmetic; the float multiply by
``scales`` matches up to the dtype's rounding.

Compile / capture safety: the Triton path defines an ``@triton.jit``
kernel whose decorator reads Triton knobs at call time — that config
read is an untraceable host op under Dynamo and aborts a CUDA-graph
capture (``cudaErrorStreamCaptureInvalidated``). When the AWQ
dequant-fallback linear (a shape outside Marlin's tile bounds, e.g.
a DeepSeek dense ``down_proj`` with ``in % 128 != 0``) is reached
inside a ``torch.compile`` trace or a live capture, route through the
pure-torch reference instead — it is bit-identical and fully
traceable / capturable.

AWQ weight-quantization Linear classes — concrete impl of the
:mod:`arbi_serve.weight_quant.base.QuantLinearBase` contract.

Two on-disk formats are supported:

  - **Legacy AWQ** — ``{qweight (i32), qzeros (i32), scales (fp16/bf16)}``;
    qweight is ``[in, out//8]`` (8 nibbles per int32 along output axis,
    using the AWQ-specific ``[0, 4, 1, 5, 2, 6, 3, 7]`` reverse-order
    permutation). Bound via :meth:`awq_load`.
  - **Compressed-tensors pack-quantized** —
    ``{weight_packed (i32), weight_scale (bf16/fp16), weight_shape (i64),
    weight_zero_point (i32, optional)}``; weight_packed is
    ``[out, in//8]`` (8 nibbles per int32 along INPUT axis, sequential
    order — no AWQ permutation). Bound via :meth:`pack_quantized_load`.
    This is the format every modern AWQ checkpoint uses (Qwen 3.5 / 3.6,
    Llama 3 AWQ, Mistral AWQ); the legacy layout is preserved here for
    older ``casperhansen/AutoAWQ``-produced packs.

Forward path
------------
On each load, the per-rank AWQ tensors are repacked (once) into Marlin's
interleaved INT4 layout — see
:mod:`arbi_serve.weight_quant.awq.marlin` for the JIT-compiled CUDA
kernel. Forward then calls ``marlin_gemm`` directly: a fused
dequant+matmul that never materializes the dense bf16 weight, so a 9B
card.

A torch-fallback path (the original dequant + ``F.linear``) is preserved
for shapes that violate Marlin's tile constraints
(``out % 256 == 0`` and ``in % 128 == 0``); this is rare on Qwen3.5 /
Llama AWQ exports but exists for completeness.

dtype handling. Activations stay at their incoming dtype (bf16/fp16);
``scales`` is read at its native dtype, repacked into Marlin's permuted
layout in the same dtype, and the GEMM output dtype follows the input.

TP. Pack-quantized lays out the weight as ``(out, in)`` so:
  - Column-parallel slices the OUTPUT axis (weight_packed dim 0,
    weight_scale dim 0, weight_zero_point along its OUTPUT axis).
  - Row-parallel slices the INPUT axis (weight_packed dim 1 in groups
    of pack-factor=8 nibbles per int32, weight_scale dim 1 in groups
    of group_size; weight_zero_point dim 1 in groups of group_size).

Legacy AWQ is the symmetric of that: qweight is ``(in, out)``-major
and TP slicing flips axes. Both classes retain the same TP signature;
the slicing math runs inside the per-format load method.

LoRA. AWQ + LoRA composes correctly — :class:`QuantLinearBase`
inherits :meth:`LinearBase._maybe_apply_lora`, which dispatches to the
BGMV kernel; the kernel accumulates the correction on top of the
Marlin GEMM output regardless of how it was computed.

Module layout (this module re-exports the public API):

  - :mod:`arbi_serve.weight_quant.awq._marlin_dispatch` — Marlin tile
    predicate + the ``use_atomic_add`` reduce gate.
  - :mod:`arbi_serve.weight_quant.awq._linear_base` —
    :class:`AWQLinearBase`.
  - :mod:`arbi_serve.weight_quant.awq._parallel` — the TP variants.

Vendored Marlin INT4 GEMM for AWQ — Python entry points.

Four torch ops are exposed:

  - :func:`awq_marlin_repack` — legacy AWQ qweight (size_k, size_n // pack)
    packed along the OUTPUT axis with the AWQ ``[0,4,1,5,2,6,3,7]``
    interleave → Marlin's interleaved layout.
  - :func:`gptq_marlin_repack` — standard-order qweight
    (size_k // pack, size_n) packed along the INPUT axis in sequential
    bit order → Marlin's interleaved layout. Used for the
    pack-quantized → Marlin layout transform after
    :func:`pack_quantized_to_standard`.
  - :func:`marlin_gemm` — fused dequant+matmul on Marlin-laid-out int4
    weights. Output dtype follows ``a.dtype`` for fp16/bf16 activations
    (W4A16) and ``b_scales.dtype`` for fp8 activations (W4A8).
  - :func:`marlin_int4_fp8_preprocess` — load-time nibble remap that
    folds the W4A8 kernels' sign-decode (and, for asymmetric AWQ
    checkpoints, the zero-points) into the stored weight.

All three ops register into the JIT-extension's private op namespace
(``arbi_serve_awq_marlin::*``) on the first call to any of them; the
backing ``.so`` is JIT-compiled with :func:`torch.utils.cpp_extension.load`
and cached at ``~/.cache/torch_extensions/...``. Vanilla bf16 runs never
touch this path.

Return the ``torch.ops.<ext>`` accessor, JIT-compiling on first hit.

After the first call the namespace handle is cached, so callers inside
a traced region hit a plain module-global read (Dynamo-safe) rather
than the un-traceable :func:`_ensure_loaded` filesystem path.

Repack a legacy-AWQ-format int32 qweight into Marlin's layout.

Input shape: ``(size_k, size_n // pack_factor)``, packed along the
OUTPUT axis with the AWQ ``[0, 4, 1, 5, 2, 6, 3, 7]`` reverse-order
permutation.

Output shape: ``(size_k // 16, size_n * 16 // pack_factor)``, int32.

Constraints:
  - ``size_k`` divisible by ``marlin::tile_k_size = 16``;
  - ``size_n`` divisible by ``marlin::tile_n_size = 64``;
  - ``num_bits ∈ {4, 8}`` (we only use 4).

Repack a standard-order int32 qweight into Marlin's layout.

Input shape: ``(size_k // pack_factor, size_n)``, packed along the
INPUT axis in sequential bit order ``[0, 1, 2, 3, 4, 5, 6, 7]``.

``perm``: int32 permutation of length ``size_k`` for GPTQ act-order
checkpoints, or an empty tensor (``torch.empty(0, dtype=torch.int32,
device=b_q_weight.device)``) for AWQ / no-act-order. Pass ``None``
and we'll synthesize an empty tensor.

Output shape: ``(size_k // 16, size_n * 16 // pack_factor)``, int32.

Enforce the Marlin GEMM's 16-bit dtype contract.

The kernel templates on the ACTIVATION dtype and reads ``b_scales`` /
``bias`` through that same ``scalar_t`` pointer with no conversion, so
a 16-bit tensor whose dtype differs from the template is reinterpreted
bit-for-bit and corrupts the output silently. Raise instead.

W4A16: ``a`` is fp16/bf16 and ``b_scales`` (and ``bias``) must match it
exactly. W4A8: ``a`` is float8_e4m3fn with ``a_scales``; the output
follows ``b_scales``, which must be fp16/bf16 and must match ``bias``.

Run Marlin's fused dequant + GEMM on AWQ-INT4 weights.

Args:
    a: (M, K) fp16 / bf16 activations (W4A16), or float8_e4m3fn
        activations (W4A8 — requires ``a_scales`` and a weight that
        went through the ``is_a_8bit=True`` repack +
        :func:`marlin_int4_fp8_preprocess`).
    b_q_weight: Marlin-laid-out int32 weight (output of
        :func:`awq_marlin_repack` or :func:`gptq_marlin_repack`).
    b_scales: ``(num_groups, N)`` fp16 / bf16 group scales — must be
        in Marlin's permuted layout (call
        :func:`marlin_permute_scales` before this; the W4A8 path
        additionally pre-multiplies by 512, the fp8-exponent trick).
        On the W4A16 path its dtype MUST equal ``a.dtype`` (see
        :func:`_check_gemm_dtypes`).
    workspace: int32 scratch buffer with at least
        ``num_compute_units(device)`` ints — call
        :func:`marlin_make_workspace` once at load time per device.
    a_scales: per-token activation scales, ``(M, 1)`` float32 —
        required iff ``a`` is 8-bit (the kernel folds them into the
        epilogue). ``None`` for fp16/bf16 activations.
    b_zeros: optional ``(num_groups, N // pack_factor)`` int32
        zero-points in Marlin's permuted layout (asymmetric AWQ
        checkpoints; symmetric ones pass ``None``). On the W4A8
        path the zero-points are ALSO folded into the weight by
        :func:`marlin_int4_fp8_preprocess`, but the kernel still
        needs the tensor to complete the negative-branch decode.
    g_idx, perm: GPTQ act-order params; pass ``None`` for AWQ.
    bias: optional ``(N,)`` fp16/bf16 bias in Marlin's permuted
        layout.
    size_m / size_n / size_k: GEMM shape (M is inferred from a but
        we pass it explicitly so the kernel can reshape).
    is_k_full: True for AWQ (no act-order); see
        :func:`marlin_is_k_full`.
    use_fp32_reduce: True for accuracy-critical paths (default).
    use_atomic_add: cross-CTA reduce via atomic_add — faster on
        small N×M but slightly less deterministic. Default False;
        arbi-serve doesn't toggle it.

Returns:
    (M, N) tensor in ``a.dtype`` for fp16/bf16 activations, or
    ``b_scales.dtype`` for 8-bit activations (the W4A8 output
    follows the scales' 16-bit dtype).

Fold the W4A8 nibble remap into an int4 weight (load-time only).

The fp8-activation kernels decode weights with a sign-remap that
must be pre-applied to the stored nibbles
(``v >= zp ? v - zp : 15 - v`` per element, ``zp = 8`` for the
symmetric u4b8 layout):

  - **Symmetric (u4b8, pack-quantized)**: call on the
    MARLIN-laid-out weight (AFTER the repack), ``qzeros=None``.
  - **Asymmetric (u4 + zero-points, AWQ)**: call on the AWQ-layout
    ``(size_k, size_n // 8)`` weight with the AWQ-layout ``qzeros``
    BEFORE :func:`awq_marlin_repack` — the fold needs the row →
    group association, which the Marlin interleave destroys.

Returns the transformed tensor (a copy unless ``inplace=True``).

Allocate Marlin's per-device int32 workspace.

Sized at the SM count on the device (Marlin uses one int32 lock
slot per CTA, capped at one CTA per SM in our default config).

Marlin's scale-permutation tables (fp16/bf16 weights).

Mirrors vLLM ``marlin_utils.get_scale_perms`` exactly. ``scale_perm``
is the per-group permutation (group_size > 0); ``scale_perm_single``
is the per-channel permutation (group_size == -1, or for the
"single" zero-points layout). Both are length-64.

Permute a ``(num_groups, N)`` scale tensor into Marlin's layout.

Mirrors vLLM ``marlin_utils.marlin_permute_scales``. The 8-bit-
activation kernels use the "single" permutation even for grouped
scales (their MMA tile maps scales differently), so ``is_a_8bit``
forces the ``scale_perm_single`` branch.

Permute a ``(num_groups, N)`` UNSIGNED int zero-point tensor (int32-
valued, range [0, 15]) into Marlin's layout, then re-pack to int32.

Output shape: ``(num_groups, N // pack_factor)`` int32.
Mirrors vLLM ``marlin_utils.marlin_zero_points`` for num_bits=4
(``size_k`` is omitted — vLLM's signature included it for symmetry
with ``marlin_permute_scales``, but the kernel-side layout is fully
determined by the reshaped row count). The 8-bit-activation kernels
skip the column interleave (their dequantize order differs), so
``is_a_8bit`` bypasses that step.

Convert a compressed-tensors pack-quantized weight to GPTQ-standard
layout.

Input: ``(out_features, in_features // 8)``, int32, packed sequentially
along the INPUT axis.

Output: ``(in_features // 8, out_features)``, int32, packed
sequentially along the INPUT axis (axis 0 here). This is what
:func:`gptq_marlin_repack` consumes (size_k = in_features along axis 0).

The two layouts pack the SAME axis (input) with the SAME sequential
shift order ``[0, 4, 8, …, 28]``, so input element ``c`` of output row
``r`` sits at word ``c // 8``, shift ``4 * (c % 8)`` in BOTH — the
layouts differ only in which axis indexes the word. The transform is
therefore an axis swap of the already-packed int32 words, not a
re-pack: nothing is unpacked, and peak scratch is the one output copy
(``out × in / 2`` bytes) instead of two dense ``out × in`` int32
nibble grids (8× larger).

Pack-quantized signed-int4 values are stored unsigned-shifted
(``v + 8`` ∈ [0, 15]); the axis swap moves words, so that encoding —
like every other bit-level detail — passes through untouched.

Output columns per :func:`pack_quantized_to_marlin` chunk.

The largest :data:`_MARLIN_TILE_N`-aligned block whose two per-chunk
transients (the standard-layout slice and the Marlin-layout slice, each
``block * in_features / 2`` bytes) fit :data:`_MARLIN_REPACK_CHUNK_BYTES`;
never below one tile.

Repack a pack-quantized int4 weight into Marlin's W4A16 layout with a
peak bounded by one output-column block.

Input ``(out_features, in_features // 8)`` int32; output
``(in_features // 16, out_features * 16 // 8)`` int32 — identical to
:func:`pack_quantized_to_standard` followed by :func:`gptq_marlin_repack`,
but the full ``(in_features // 8, out_features)`` standard-layout copy is
never materialised: only the resident result plus one block's worth of
slices is live at any point.

**Chunk axis and its granularity.** The chunk axis is the OUTPUT axis
(N, the vocab dimension of an LM head). ``gptq_marlin_repack`` writes
``out[(k_tile * n_tiles + n_tile) * tile_k * tile_n / pack_factor]``
with ``tile_k = 16``, ``tile_n =`` :data:`_MARLIN_TILE_N`, so one
declared output ROW is exactly one k-tile and the n-tiles sit side by
side inside it: output columns ``[n0 * 2, (n0 + block) * 2)`` hold
output-feature columns ``[n0, n0 + block)`` and nothing else. Reads are
equally n-tile-local (``b_q_weight[k * size_n + n_tile * tile_n + …]``).
A block that is a multiple of ``tile_n`` is therefore a self-contained
unit of the layout and the chunked result is bit-identical to the
whole-matrix one. A block that is NOT a multiple of ``tile_n`` would
straddle a tile, so it is refused rather than silently mis-laid.

W4A16 only. The W4A8 repack halves ``tile_n`` and doubles ``tile_k``,
which puts TWO declared output rows in one k-tile — an output-column
block is then not a tile unit, so that layout must be repacked whole
(callers pass ``is_a_8bit=True`` to :func:`gptq_marlin_repack` directly).

Output rows per :func:`pack_quantized_unpack_weight` chunk, bounded by
:data:`_UNPACK_GRID_CHUNK_BYTES` (one chunk is ``rows × in_features``
int32).

The division is :func:`~arbi_serve.runtime.row_tiling.rows_per_tile`; the
policy here is the per-row size (one int32 grid row) and that a weight
already inside the budget is unpacked whole.

Unpack a pack-quantized weight to the dense ``(in, out)`` int32
nibble grid (values in [0, 15], unsigned-shifted ``v + 8``).

Load-time only — feeds :func:`pack_awq_layout` on the W4A8
zero-point-folding route, which needs the AWQ layout for
:func:`marlin_int4_fp8_preprocess`.

The grid is filled one output-row block at a time, transposed straight
into the result, so the transient is one block rather than a second
full grid. Output rows carry no packing in ``weight_packed`` (the pack
factor runs along the INPUT axis), so a block boundary can fall
anywhere and the result is bit-identical.

Pack a dense ``(rows, size_n)`` int grid in [0, 15] into the
legacy-AWQ int32 layout ``(rows, size_n // 8)`` with the AWQ nibble
interleave.

Load-time only — builds the AWQ-layout weight + zero-point pair that
:func:`marlin_int4_fp8_preprocess` consumes on the asymmetric W4A8
path (the fold needs the row → group association, so it must run in
a K-major layout BEFORE the Marlin repack).

Convert pack-quantized ``weight_zero_point`` (packed along OUTPUT
axis, int32) into Marlin's int32 zero-point layout.

Input shape: ``(out_features // 8, num_groups)``, packed along axis 0.
Output shape: ``(num_groups, out_features // 8)``, packed along axis 1
(Marlin layout).

Pack-quantized stores zero-points as (signed_zp + 8) just like the
weights; Marlin expects unsigned-shifted [0, 15] zeros too — so no
arithmetic adjustment is needed, only a layout permutation.

JIT loader for arbi-serve's vendored Marlin INT4 GEMM.

The Marlin .cu sources live under :mod:`arbi_serve.weight_quant.awq.marlin.csrc`.
On first call to :func:`_ensure_loaded` we hand the file list to
:func:`torch.utils.cpp_extension.load`, which spawns nvcc to compile the
Marlin kernel templates, produces a ``.so``, and registers four Torch
ops into the private ``arbi_serve_awq_marlin`` library (stable-ABI
static registrations — the .so has no PyInit, so it is loaded with
``is_python_module=False`` via ``torch.ops.load_library``):

  - ``arbi_serve_awq_marlin::awq_marlin_repack``
  - ``arbi_serve_awq_marlin::gptq_marlin_repack``
  - ``arbi_serve_awq_marlin::marlin_gemm``
  - ``arbi_serve_awq_marlin::marlin_int4_fp8_preprocess`` (W4A8
    zero-point folding, load-time only)

Subsequent calls hit the torch-extensions cache
(``~/.cache/torch_extensions/...``) and load instantly.

Memory hygiene
--------------
The module is never imported at boot. ``AWQLinearBase.__init__`` doesn't
call into here; only ``pack_quantized_load`` / ``awq_load`` do, and only
when the host has CUDA. CPU-only smoke + import-time tests stay fast.

Return the .cu / .cpp sources that go into the JIT compile.

Order is irrelevant (the compiler links them all into one .so), but
we list them deterministically so cpp_extension's hash of the source
set is stable across runs.

JIT-compile (or pull from cache) the Marlin extension.

The ops register into the global ``torch.ops.arbi_serve_awq_marlin``
namespace as a dlopen side effect (``is_python_module=False`` — the
stable-ABI .so has no PyInit); the return value is None.

Cached via :func:`functools.lru_cache` so subsequent calls return
the already-loaded handle without re-running cpp_extension's hash
check.

Side effect: once the extension is loaded, this function attaches
Python-side ``register_fake`` impls to the three Marlin ops via
:func:`arbi_serve._custom_ops.register_marlin_fakes_if_loaded` so
that ``torch.compile`` / Inductor can plan allocations through the
Marlin GEMM as a black box.

Legacy-AWQ tensors → the fused-MoE Triton kernel's expert layout.

The AWQ on-disk triple is K-major and packed along the OUTPUT axis::

    qweight  int32  (K, N // 8)          8 nibbles/int32, AWQ interleave
    qzeros   int32  (K // gs, N // 8)    same packing
    scales   16-bit (K // gs, N)

``fused_moe_kernel_gptq_awq`` (``use_int4_w4a16``) reads an N-major
expert stack instead, and packs the weight along K but the zero-point
along N::

    B        uint8  (E, N, K // 2)       b_shifter = (k % 2) * 4
    B_scale  16-bit (E, N, K // gs)
    B_zp     uint8  (E, N // 2, K // gs)  b_zp_shifter = (n % 2) * 4

so the low nibble of a weight byte is the EVEN k and the low nibble of
a zero-point byte is the EVEN n. Each function here converts one
checkpoint tensor into one expert slot's slice of those buffers; the
values are the raw uint4 codes, unchanged — only the container moves.

Compressed-tensors ``pack-quantized`` INT4 dequantize.

The compressed-tensors ``pack-quantized`` layout is the canonical
on-disk format produced by ``llm-compressor`` for AWQ-style INT4
weight-quantized checkpoints (Qwen 3.5 / 3.6, Llama 3 AWQ,
Mistral AWQ, etc.). It is what every modern AWQ checkpoint uses;
the legacy AWQ layout (``.qweight`` / ``.qzeros`` /

Layout (per linear, see ``compressed_tensors.compressors.quantized_compressors.pack_quantized``):

  - ``weight_packed``: ``int32``, shape ``(out_features, in_features // 8)``.
    Each int32 packs 8 sequential INT4 nibbles along the input axis,
    bit-shift order ``[0, 4, 8, 12, 16, 20, 24, 28]`` (sequential —
    NOT the legacy AWQ ``[0, 4, 1, 5, 2, 6, 3, 7]`` reverse order).
  - ``weight_scale``: float (bf16/fp16/fp32), shape
    ``(out_features, in_features // group_size)``.
  - ``weight_shape``: int64 ``[out_features, in_features]`` — metadata
    for the un-packed shape (carried because the int32 pack adds
    padding that ``unpack`` must trim).
  - ``weight_zero_point`` (asymmetric only): packed int32, shape
    ``(out_features // 8, in_features // group_size)`` (packed along
    the OUTPUT axis with the same 8-per-int32 sequential layout).
    Symmetric checkpoints (most public AWQ packs we ship — Qwen 3.5/3.6,
    Llama 3 AWQ) do NOT carry this tensor; the dequant uses the
    canonical INT4 mid-point offset of 8 instead.

Numeric contract::

    int8_signed[out, in] = unpack_int4(weight_packed[out, in // 8])
                            (subtract offset=8 -> signed [-8, 7])
    W[out, in] = int8_signed[out, in] * scales[out, in // group_size]
                  (asymmetric: + zero_point per group)

Two paths share that contract: a Triton kernel for CUDA + a torch
reference for CPU smoke + parity tests. Both produce the same
``(out_features, in_features)`` dense matrix in ``scales.dtype``.

This module is the on-disk format reader; the dispatch into the right
``QuantLinearBase`` subclass and the ``apply_quant_if_present`` hook
both live in :mod:`arbi_serve.weight_quant.awq.backend` (which
auto-detects pack-quantized vs legacy AWQ from the safetensors keys
and routes binding accordingly).

Unpack int32 → int8-signed for the pack-quantized INT4 layout.

Mirrors ``compressed_tensors.compressors.quantized_compressors.
pack_quantized.unpack_from_int32`` for ``packed_dim=1``: each int32
yields 8 sequential INT4 nibbles along the input axis, then the
canonical offset (8) is subtracted to restore signed [-8, 7].

Returns ``int8`` tensor of shape ``(out_features, in_features)``;
padding is trimmed when ``in_features`` is not a multiple of 8.

Dequantize pack-quantized INT4 weights to a dense ``(out, in)`` matrix.

Picks the Triton kernel when running on CUDA + Triton is available,
otherwise falls back to the torch reference. Both paths produce the
same dense ``(out_features, in_features)`` matrix in ``weight_scale.dtype``;
integer arithmetic is bit-identical, the float multiply by
``weight_scale`` matches up to the dtype's rounding.

Triton dequant for pack-quantized INT4.

Output shape ``(out_features, in_features)`` in ``weight_scale.dtype``.
Each program handles a ``[block_m, block_n]`` tile in OUTPUT space;
the input axis (block_n) is processed 8 elements per int32 read.

RTN (round-to-nearest) INT4 quantizer emitting the pack-quantized layout.

Produces ``(weight_packed, weight_scale, weight_zero_point)`` tensors in
EXACTLY the compressed-tensors ``pack-quantized`` on-disk layout that
:meth:`AWQLinearBase.pack_quantized_load` consumes (asymmetric INT4,
per-``(row, group)`` scale + packed zero-point, group along the INPUT
axis) — the same scheme the Qwen3.6-27B-AWQ-INT4 trunk ships
(``num_bits=4, group_size=32, symmetric=false``).

Used for LOAD-TIME quantization of checkpoint-dense layers the AWQ
export deliberately skipped (``quantization_config.ignore``: ``lm_head``,
``mtp.fc``) — there is no calibration data for those layers in the
checkpoint, so data-free RTN with a min-max (zero-inclusive) grid is the
scheme. Quality gates live with the callers: the drafter-head consumer is
accept-length-gated, the target-head consumer is eval-gated.

Numerics contract (tested):
  - the stored nibble is ``u = clamp(round(w / scale) + zp_u, 0, 15)``;
  - dequant (see :func:`pack_quantized_dequantize_torch`) reads
    ``(u - zp_u) * scale`` — so ``|dequant - w| <= scale / 2`` for every
    in-range element, up to the bf16 rounding of ``scale`` itself;
  - ``scale`` is computed in fp32, rounded to the STORAGE dtype, and the
    ROUNDED value is what quantization divides by (so the dequant error
    stays centred instead of inheriting the scale's own rounding bias);
  - the min/max grid is zero-inclusive (``min(w, 0) .. max(w, 0)``) so a
    zero weight — including the loader's zero-filled pad rows of a short
    vocab shard — always round-trips to exactly 0.0;
  - the fp32 grid is materialised in ROW CHUNKS
    (:data:`_RTN_QUANT_CHUNK_BYTES`), never over the whole tensor: rows are
    independent (grouping runs along the INPUT axis), so the emitted bytes
    are identical to an unchunked upcast while peak scratch stays bounded
    by the chunk budget rather than scaling with the shard.

Rows per fp32 quantize chunk, bounded by :data:`_RTN_QUANT_CHUNK_BYTES`.

Chunking is along the OUTPUT axis, which the quantize step treats
elementwise (grouping runs along the INPUT axis), so the chunk width
cannot change the emitted nibble.

The division is :func:`~arbi_serve.runtime.row_tiling.rows_per_tile`; what
lives here is the POLICY — the fp32 upcast is ``in_features * 4`` bytes per
row, and a shard that already fits the budget is quantized whole.

Pack an unsigned-nibble ([0, 15]) int32 tensor 8-to-1 along ``axis``.

Sequential bit order (nibble ``i`` of each int32 at bits ``4*i..4*i+3``)
— mirrors llm-compressor's ``pack_to_int32`` and the unpack loops in
:mod:`arbi_serve.weight_quant.awq.pack_quantized`.

RTN-quantize a dense ``(out, in)`` weight to pack-quantized INT4.

Returns ``(weight_packed, weight_scale, weight_zero_point)``:
  - ``weight_packed``:     ``(out, in // 8)`` int32, packed along INPUT;
  - ``weight_scale``:      ``(out, in // group_size)`` ``scale_dtype``;
  - ``weight_zero_point``: ``(out // 8, in // group_size)`` int32,
    packed along OUTPUT (asymmetric, unsigned-shifted like the weights).

Rows are quantized independently (grouping is along the input axis),
so quantizing a TP output-shard locally is bit-identical to slicing a
globally quantized tensor — each rank may RTN its own vocab shard.

Deterministic: pure elementwise torch ops, no data-dependent branching.

Per-group ``(weight_scale, zp_u)`` from zero-inclusive ``(wmin, wmax)``.

Single source of truth for the asymmetric-INT4 scale/zero-point math
(shared by RTN min/max and the activation-weighted clip search):
  - ``scale = (wmax - wmin) / 15`` in fp32, rounded to ``scale_dtype``,
    and the ROUNDED value is what quantization divides by (dequant
    multiplies by the stored scale, so dividing by anything else
    de-centres the error);
  - ``zp_u = round(-wmin / scale)`` clamped to ``[0, 15]`` — lands
    ``wmin`` on 0.

Inputs already zero-inclusive (``wmin <= 0 <= wmax``). Returns
``(weight_scale[scale_dtype], zp_u[int32])`` at the ``(out, G)`` shape
of the inputs.

Emit the pack-quantized triple from a chosen per-group scale + zp.

``weight`` is ``(out, in)``; ``weight_scale`` / ``zp_u`` are
``(out, in // group_size)`` (the stored scale dtype and int32 zp). The
nibble is ``u = clamp(round(w / scale) + zp, 0, 15)`` — the exact
inverse of :func:`pack_quantized_dequantize_torch`. Shared by RTN and
the activation-weighted clip search so both round-trip identically.

Pack-quantized triple from a per-group zero-inclusive ``(wmin, wmax)``.

Convenience wrapper: :func:`group_scale_zp_from_minmax` then
:func:`pack_quantized_from_scale_zp`. The activation-weighted clip
search feeds its chosen (clipped) min/max here so the emit path is
identical to RTN's.

SmoothQuant per-input-channel activation smoothing for the W4A8 recipe.

The naive W4A8 path (per-token e4m3 activation quant, no smoothing) carries a
per-projection format floor set by whichever input channel has the largest
magnitude. Post-SwiGLU (``down_proj`` input) and attention-output (``o_proj`` /
``out_proj`` input) activations have a handful of channels 2-3 orders of
magnitude above the rest, so the per-token scale is spent on the outlier and
every other channel quantizes into a few e4m3 codes. Composed across the depth
of a 27B/31B stack this floor collapses the logits.

SmoothQuant removes the floor with an EXACT pre-quant transform. For a linear
``y = X Wᵀ`` a per-input-channel scale ``s`` migrates activation difficulty
into the weights::

    X Wᵀ == (X · diag(1/s)) · (W · diag(s))ᵀ

with ``s_j = max|X_j|^α / max|W_j|^(1-α)`` (SmoothQuant, arXiv 2211.10438).
The smoothed activation ``X·diag(1/s)`` has a far flatter per-channel profile,
so its per-token e4m3 quant floor drops by the outlier ratio. The transform is
exact BEFORE quantization — its whole value is that what it feeds the quantizer
is easier to represent.

This module is pure tensor algebra (CPU-verifiable); it computes ``s``, applies
the two half-folds, and provides the fp8 round-trip helpers the calibration /
verification path uses to measure the floor reduction. Where the smoothed
weight is served is a producer/serving concern:

  * norm-fold — ``1/s`` folds into the preceding RMSNorm weight (and any dense
    siblings sharing that norm are compensated by ``·diag(s)``); the served
    activation is already smoothed, so no runtime op and the pack is correct
    under both W4A16 and W4A8.
  * weight-fold — for projections with no clean pre-linear norm (``down_proj``,
    ``o_proj`` / ``out_proj``) the smoothed weight ``W·diag(s)`` is packed and
    ``act_pre_scale = 1/s`` is applied to the activation at serve time, ahead of
    the per-token e4m3 quant. That pack is W4A8-only.

Per-input-channel weight magnitude ``max_i |W_ij|`` over a group.

``weights`` is a list of ``(out_i, in)`` tensors that share the input
channel axis (a fused/grouped projection, e.g. Gemma q/k/v or the GDN
``in_proj_qkv`` + ``in_proj_z``). Returns ``(in,)`` — the max over every
output row of every member. This is SmoothQuant's ``max|W_j|`` term.

SmoothQuant per-input-channel scale ``s_j = max|X_j|^α / max|W_j|^(1-α)``.

``act_absmax`` and ``weight_absmax`` are ``(in,)`` non-negative per-channel
magnitudes (the activation abs-max from calibration and
:func:`channel_weight_absmax`). Returns ``s`` ``(in,)`` fp32, clamped to
``[min_scale, max_scale]`` so a dead activation channel (``max|X_j| = 0``)
or a dead weight channel cannot produce a 0/inf scale. ``alpha`` in ``[0,1]``
trades activation smoothing (→1) against weight quantizability (→0); 0.5 is
the SmoothQuant default and the a8 recipe's default.

``s`` is NOT renormalized around 1 (unlike the AWQ salience scale): the
absolute balance between activation and weight magnitude is what sets the
fp8 floor, and both half-folds are applied exactly, so any global rescale
would cancel anyway.

Fold ``1/s`` into an RMSNorm weight so the norm emits ``X·diag(1/s)``.

A ``* weight`` RMSNorm (Gemma-4) folds as ``w/s``; a ``(1 + weight)`` RMSNorm
(Qwen3.5/3.6, ``unit_offset``) folds as ``(1 + w)/s - 1`` so the effective
gain ``(1 + w_new) = (1 + w)/s``. Returns the new norm weight in the input
dtype.

Per-token (per-row) e4m3 quant→dequant of ``x`` ``(M, in)``.

Mirrors the serving ``_quant_fp8_per_token`` scale choice (row abs-max /
448, clamped) so the measured floor matches what the W4A8 kernel sees.
Returns the dequantized fp32 tensor.

Per-channel dynamic-range spread of ``x`` ``(M, in)``: the ratio of the
largest per-channel abs-max to the MEDIAN per-channel abs-max.

A few channels dwarfing the median is exactly the profile that wrecks
per-token fp8 (the row scale is spent on the outlier, so the median-scale
channels quantize into a couple of e4m3 codes); SmoothQuant collapses it
toward ~1.

Mean over input channels of the per-channel per-token e4m3 rel-error.

Unlike the energy-weighted :func:`relative_mse` (which the outlier channels
dominate and so looks fine even when the rest is destroyed), this weights
every channel equally — so it exposes the many median-magnitude channels the
outlier's per-token scale quantizes into near-zero. This is the floor that
compounds across depth.

Measure the per-token fp8 floor of ``x`` before vs after smoothing by
``s``. Returns the outlier ratio and rel-MSE for both, for the calibration
/verification report. ``x`` ``(M, in)`` real captured activations, ``s``
``(in,)`` the SmoothQuant scale.

Backend-agnostic surfaces for weight quantization.

A *quant backend* is a way to load a model whose dense ``Linear``
projection weights have been replaced on disk by a backend-specific
encoded representation (trellis grids, packed nibble groups, fp8
tiles, etc.). Every supported backend (EXL3, AWQ, NVFP4, MXFP4, FP8)
implements the :class:`QuantBackend` protocol and registers itself
via :mod:`arbi_serve.weight_quant.registry`.

Three structural pieces live here:

  - :class:`QuantLinearBase` — base class every quant Linear inherits.
    It carries the ``is_quantized`` marker (so :func:`weight_map`
    filters skip it) and locks the dense ``weight_loader`` out, so
    accidental dense binds raise loudly. TP slicing math is *not*
    here — sharding strategies differ across backends (some quant
    along output dim, some along groups) so the per-backend Linear
    classes implement their own slicing inside their per-backend
    binding API.

  - :class:`QuantBackend` — the protocol an integration implements.
    Detection from a safetensors collection, enumerating quantized
    linear paths, picking the right per-Linear class for each dense
    Linear, and binding backend-specific tensors per linear.

  - :class:`LoRAUnsupportedError` — sentinel raised by a backend that
    cannot compose with LoRA. EXL3 + LoRA works (BGMV applies on top
    of the dequant output); other backends might not.

Per-arch ``from_safetensors`` calls one entry point —
:func:`apply_quant_if_present` from :mod:`arbi_serve.weight_quant.loader`.
The arch's ``weight_map`` filters via :func:`filter_weight_map_for_quant`.
Both helpers are backend-agnostic; they consult the registry to
discover whichever backend owns the checkpoint at hand.

Common surface for every backend's quant-Linear classes.

Subclasses are expected to:

  - Carry zero ``nn.Parameter`` named ``weight``. The dense loader
    walks ``named_parameters``; a missing entry there is what makes
    :func:`filter_weight_map_for_quant` cleanly skip the swapped
    linear without per-arch wiring.
  - Register backend-specific tensors as buffers (non-persistent)
    and bind them via a backend-specific load method (e.g.
    :meth:`EXL3LinearBase.exl3_load`).
  - Implement :meth:`forward` themselves — no shared compute path
    works across trellis dequant, packed-int4 GEMM, fp8 mma.
  - Inherit the LoRA hook from :class:`LinearBase` (the BGMV apply
    accumulates on top of any base output, regardless of how it
    was computed). If a backend can't support LoRA, override
    :meth:`_maybe_apply_lora` and raise :class:`LoRAUnsupportedError`.

The ``is_quantized`` attribute is the duck-type marker the
weight-map filter checks. Set on the class so isinstance is not
required for the check (lighter coupling for downstream code).

Shared TP alignment-check helpers for the Column / Row mixins.

``tp_block_align`` is the per-rank shard granularity along the sharded
axis (AWQ INT4 pack-factor 8, EXL3 / NVFP4 16-element block, FP8 1);
``tp_align_note`` is the parenthetical reason appended to the
divisibility error when the alignment is wider than ``tp_size`` alone.

Shared ``__init__`` + ``forward`` for a backend's column-parallel
quant Linear (output dim sharded across the TP group).

The subclass sets :attr:`tp_block_align` (+ :attr:`tp_align_note`) and
implements ``_quant_forward``; per-tensor output-axis slicing stays in
the backend's own load method.

Shared ``__init__`` + ``forward`` for a backend's row-parallel quant
Linear (input dim sharded across the TP group).

The subclass sets :attr:`tp_block_align` (+ :attr:`tp_align_note`) and
implements ``_quant_forward``; per-tensor input-axis slicing stays in
the backend's own load method.

Protocol every weight-quantization backend implements.

Backends register themselves at import time via
:func:`arbi_serve.weight_quant.registry.register`. The loader iterates
the registry in *detection-priority* order — sorted by
:attr:`detection_priority` ascending (lower = runs first); the first
``detect(stc)`` to return ``True`` wins and drives the swap.

A backend may pull in heavy upstream dependencies (custom CUDA
extensions, JIT-compiled kernels). Lazy-import them — the
registration call should NOT trigger the import; only the first
Linear bind should.

What a backend can hand the stacked fused-MoE path.

``quant_kind`` is a :attr:`~arbi_serve.models.moe.FusedMoE.SUPPORTED_QUANT_KINDS`
member; ``group_size`` is the checkpoint's quantization group along
the reduction axis, and ``block_size`` its 2-D block for layouts
quantized blockwise instead.

Drop ``*.weight`` entries whose owning module is quantized.

Backend-agnostic: duck-types on the ``is_quantized`` marker that
every :class:`QuantLinearBase` carries. Any arch whose
``weight_map()`` is the standard ``{"<dotted>.weight": Spec(...)}``
shape can route through this instead of writing per-arch filter
loops.

Detect any registered quant backend on ``stc`` and untie the
``lm_head`` if that backend requires it. Returns the (possibly
updated) ``dims``.

One-call from each arch's ``from_safetensors``. For dense
checkpoints returns ``dims`` unchanged. The detection pass is
cheap (one safetensors-keys scan).

Drop ``tie_word_embeddings`` if a backend requires the untie.

Single source of truth for the untie — :func:`prepare_dims_for_quant`
calls this for every backend whose ``lm_head_untie_required()`` is
True. Returns a (possibly new) ``dims`` instance.

Register a dense ``bias`` parameter mirroring the swapped-out
dense linear's (rank-local) bias.

``like`` is the dense module's bias parameter (usually still on
meta) — only its shape is read. ``dtype`` must be the ENGINE dtype
(bf16): the dense construction dtype on ``like`` is fp32 (params
are cast at meta-materialization, which this real tensor skips),
and ``load_model_weights``'s ``param.copy_`` casts INTO the param
dtype — an fp32 bias would silently promote the layer output to
fp32 and trip the attention kernels' dtype asserts. Falls back to
``like.dtype`` only when no engine dtype was threaded (tests).

The parameter is created empty on CPU (the swap's ``.to(device)``
moves it) and filled by
:func:`arbi_serve.loader.weights.load_model_weights` through the
plain-param branch: the weight_map's bias spec carries the
``shard_dim`` so per-rank slicing happens there, exactly as for
the dense module. Application lives in the TP mixins' ``forward``
(dense-parity: column adds the local shard pre-gather; row adds
once — fused at TP=1, post-all-reduce at TP>1).

Compute the quant-GEMM for ``x`` (the per-backend dequant/matmul).

This is the single backend-specific seam the TP-scaffold mixins
(:class:`ReplicatedLinearMixin` / :class:`ColumnParallelMixin` /
:class:`RowParallelMixin`) call. Each backend's ``…LinearBase``
implements it (delegating to its own ``_awq_forward`` /
``_exl3_forward`` / ``_fp8_forward`` / ``_nvfp4_forward``); the
mixins layer the shared LoRA apply + gather / all-reduce tail on
top so those bodies are written once.

Return the active :class:`ParallelConfig`.

The Column/Row TP mixins read the parallel config through this seam
rather than calling ``get_parallel_config`` directly so that each
backend can scope the lookup to ITS OWN module namespace — the
per-backend Linear modules import ``get_parallel_config`` and
override this to call it, which keeps the unit tests' per-module
``patch.object(<backend>.linear, "get_parallel_config", ...)``
effective after the ``__init__`` scaffold moved here. This base
fallback resolves the real config for any backend that doesn't
override (none in tree today).

Return ``(start, length)`` of this rank's shard along the sharded axis.

The default is the even split every backend used inline before this
seam existed. A backend whose shard boundaries must land on a block
wider than ``tp_size`` divides evenly (EXL3) overrides this.

Run the quant GEMM, apply LoRA in the local out shard, then
all-gather the output shards across the TP group when
``gather_output`` is set.

LoRA lands BEFORE the all-gather: a column-parallel base linear's
LoRA weights are sharded along the same output dim, so each rank's
correction lives in its local out shard; the optional gather then
composes everything back to the full output.

Bias (when attached) is the rank-LOCAL output shard — the loader
sliced it along the same output dim as the weight — added before
LoRA/gather, matching the dense :class:`ColumnParallelLinear`
(which fuses it into ``F.linear``).

This rank's un-gathered output shard: the quant GEMM plus the
rank-local bias, WITHOUT the ``gather_output`` all-gather and WITHOUT
LoRA. On a LoRA-free call this is bit-identical to this rank's slice
of the gathered :meth:`forward` output (same kernel, same input).

This is the vocab-shard logits surface the shard-resident verify
(:mod:`arbi_serve.spec_decode.mtp_verify_sharded`) and the drafter's
distributed argmax consume on a weight-quantized vocab-parallel
lm_head. ``out_features`` on a quant TP mixin is the LOCAL shard
width, so the returned last dim is ``out_features``.

Run the quant GEMM, apply LoRA, then all-reduce the partial sums
across the TP group (row-parallel).

The LoRA correction is itself a partial sum over the sharded input
dim; the all-reduce below combines base + LoRA partials in one
step.

Bias (when attached) is replicated — it applies to the full output
— so it must land ONCE: fused-equivalent (pre-LoRA) at TP=1,
post-all-reduce at TP>1. Mirrors the dense
:class:`RowParallelLinear` exactly.

Return True if this backend owns the given checkpoint.

Detection is fast — one pass over the in-memory safetensors
index, no shard reads. EXL3 looks for ``*.trellis``; AWQ
looks for ``*.qweight`` + ``*.scales``; FP8 looks for tile
scales of a specific dtype. Mutually exclusive in practice.

Read this backend's tensors at ``<safetensors_path>.*`` from
``stc``, apply per-rank slicing if needed, and bind them onto
the constructed Linear instance via the backend's load method.

``compute_dtype`` is the ENGINE dtype the linear's activations
arrive in. A backend whose kernel consumes a 16-bit side tensor
(AWQ Marlin's group scales) in the activation's dtype must
materialize it at ``compute_dtype``, NOT at the checkpoint's
export dtype.

The stacked-expert quant spec this backend can serve for
``expert_paths``, or ``None`` if it cannot.

The routed experts of an MoE are the one place where a
per-linear quant backend is the WRONG shape: a per-expert module
walk host-syncs on the routed ids, which forbids cudagraph
capture of the decode step. An arch asks this before choosing a
representation — a backend that answers with a spec gets the
stacked :class:`~arbi_serve.models.moe.FusedMoE`, one whose
payload cannot be stacked (EXL3 trellis, a per-linear W4A8
pre-scale) answers ``None`` and keeps the per-expert modules.

``expert_paths`` are safetensors linear prefixes (no tensor
suffix), one per routed expert projection.

Every backend answers. A backend whose payload can never be
stacked returns ``None`` for every pack; that answer is part of
the contract, not the absence of a method.

FP8 backend for DeepSeek-V4's ``.scale`` naming.

Same payload as any other block-wise FP8 pack — E4M3 weights with a
128x128 scale grid — under a different suffix: DeepSeek-V4 writes
``<linear>.scale`` where the compressed-tensors convention writes
``<linear>.weight_scale_inv``. That is the whole difference, so this is a
subclass that renames the suffix rather than a second implementation.

Detection is deliberately narrow. A bare ``.scale`` next to a ``.weight``
is a loose signature that other exporters could match, so this backend
also requires the checkpoint's own ``config.json`` to declare
``model_type: "deepseek_v4"``. It runs BEFORE the generic FP8 backend
(lower ``detection_priority``), which would otherwise not claim these
linears at all — its suffixes do not appear in this checkpoint.

The routed experts are excluded here even though they match the
signature: they are bound as one stacked payload per layer by the arch's
weight map (MXFP4 pairs on Flash-0731, block-wise FP8 on Flash-Base), not
as per-expert linear modules.

Describe the routed-expert payload for a caller that asks.

``expert_dtype`` decides it: ``fp8`` is a block-wise FP8 stack this
backend can describe, ``fp4`` is an MXFP4 stack it cannot (a
different kernel reads it), and the arch wires that one directly.

Load-time RTN fp8 quantization of the input token embedding.

The sibling of :mod:`arbi_serve.weight_quant.head_quant` (which quantizes the
OUTPUT head) for the INPUT ``embed_tokens`` table. On a large-vocab model the
untied embedding is a multi-GiB bf16 gather table that never enters a matmul —
so, unlike a weight matrix, it needs no quant GEMM kernel: fp8-e4m3 rows are
gathered with the native ``F.embedding`` / ``index_select`` and cast back to
bf16 with a per-row (per-token) dequant scale
(:meth:`~arbi_serve.models.layers.VocabParallelEmbedding.quantize_to_fp8`).

``ServerConfig.embed_quant``:

  - ``"off"``: checkpoint-faithful bf16 embedding. Pure no-op.
  - ``"auto"`` (default): a REQUEST for fp8 that :func:`resolve_embed_mode`
    settles against the loaded model — ``fp8`` where the checkpoint allows it,
    ``off`` (warning) where it does not.
  - ``"fp8"``: the embedding table is RTN-quantized to fp8-e4m3 with per-row
    scales; the bf16 table is FREED (halving its residency), growing the KV
    budget. This CHANGES EMITTED LOGITS — the embedding is the target's layer-0
    input, so the change is not verify-corrected — and it also feeds the MTP /
    DFlash drafter (which BORROWS the target's ``embed_tokens``); the gate is a
    fixed-config eval holding (top-1 agreement / KL / perplexity) plus a flat
    accept_len, not a TPOT number.

TIE SAFETY: ``fp8`` requires an UNTIED embedding. A tied head
(:class:`~arbi_serve.models.layers.TiedLMHead`) reads the embedding tensor AS
its output projection weight, so quantizing-and-freeing the table would corrupt
the lm_head matmul. :func:`resolve_embed_mode` is the ONE place that decides
this: ``auto`` skips a tied embedding, a NAMED mode refuses it loudly. Every
consumer — the eager apply, the warm placeholder arm, the exl3-direct binder —
resolves there first, so no consumer ever sees a mode it cannot execute.

Runs at the engine-boot seam AFTER weights load and BEFORE KV sizing — so the
freed bf16 bytes are visible to the KV budget — mirroring the head-quant phase.

Locate the model's input-embedding MODULE across architecture layouts.

Walks the same paths as the drafter's ``_TARGET_EMBED_PATHS`` (the DFlash
draft head borrows the target's embedding), so on every arch that exposes a
module at the first matching path the two agree and the quantized module is
exactly the one every reader — target forward, MTP embed, borrowing drafter
— dispatches through.

The walk is MODULE-typed and the drafter's is not: an arch may expose a
plain CALLABLE at one of these names as the drafter's embedding seam (Muse
Glimmer's ``_MuseGlimmerInner.embed_tokens`` property returns the nested
``_NormedEmbedding.raw`` bound method, so the borrowing head reads the
UNNORMALIZED rows). A callable is a fine embedding seam and a useless quant
target: it owns no parameters to swap and no bytes to free. Skipping it and
continuing the walk is what lands this arch on the real nested module
instead of returning something whose ``.parameters()`` does not exist —
which is what every consumer here assumes it can call.

Whether the output head reads ``embed``'s weight as its projection.

A :class:`TiedLMHead` holds the embedding it shares; a truthful
``tie_word_embeddings`` flag is the belt-and-suspenders check for a head
that reads the tensor without wrapping the module.

The single embed-quant decision: ``(concrete_mode, embedding)``.

``embedding`` is the module the concrete ``fp8`` mode transforms, and is
``None`` whenever the mode resolves ``off``. :func:`resolve_embed_mode` is
the public face; consumers that also need the module (they would otherwise
re-walk the model to find it) call this.

Resolve to a CONCRETE mode: ``"off"`` or ``"fp8"``.

``auto`` is a REQUEST, not a mode — eligibility is a property of the loaded
model (a TIED embedding cannot be freed), so only this function can decide
it. Both paths resolve here, so no downstream consumer ever sees ``auto``.
A NAMED mode is returned unchanged and still fails loud downstream.

Apply ``cfg.embed_quant`` to a fully-loaded model IN PLACE (the EAGER
path). Returns the RESIDENT quantized-embedding bytes on this rank, or 0
when the mode resolves ``off`` / the embedding is already quantized.

:func:`resolve_embed_mode` settles the mode first: ``auto`` skips a
checkpoint it cannot transform, an unknown or ineligible NAMED mode raises
— no silent fallback to bf16 (a config asking for quant and silently not
getting it would surface as unexplained model drift, not an error).

Read the embed-quant mode a fully-built model is ALREADY in (``fp8`` if
the input embedding carries that payload, else ``off``). Inverse
of :func:`apply_embed_quant`, for the donor-share path to write the truthful
resolved mode back to ``cfg``.

Reproduce a donor's fp8 INPUT EMBEDDING on a donor-bound MEMBER, so the
member presents the quantized slots the donor can actually back. Returns 1
when the member was armed, else 0.

The embed-quant sibling of
:func:`~arbi_serve.weight_quant.head_quant.share_donor_head_quant_modules`,
and it exists for the identical reason. ``fp8`` REPLACES the input-embedding
tensor set — the bf16 ``weight`` is dropped and freed (that reclaimed half
table IS the KV win) — so a donor booted fp8 carries only ``weight_fp8`` +
``weight_scale``. A member's graph is constructed fresh under
``skip_weight_load``, which builds the DENSE ``weight``, and
:class:`~arbi_serve.loader.flat_loader.DonorWeightBinder` binds by exact
qualified name — so without this the member demands
``model.embed_tokens.weight``, the donor has no such tensor, and the bind
aborts. Every capture-affecting live config override rebuilds a member, so
that abort takes the whole rebuild tier of the admin surface down (a latched
swap fault: ``/health/ready`` 503 until ``POST /v1/admin/config_restore``).

Arms the member with the SAME empty quantized placeholders the warm
flat-dump reload uses (:meth:`VocabParallelEmbedding.arm_quant_placeholders`)
rather than sharing the module outright: the zero-element ``weight_fp8`` /
``weight_scale`` buffers are exactly the placeholder family the binder
already aliases to the donor's populated buffers, so the member keeps its own
embedding module (per-member module state stays distinct, as everywhere else
on this path) and shares only the immutable storage — zero new weight VRAM.
The head shares at MODULE granularity instead only because a head's
non-zero-element marker buffers are not a placeholder the generic binder can
alias.

Runs during the member's graph construction, BEFORE the binder walks it.
Reads the donor's REALIZED state rather than ``cfg.embed_quant``: the mode is
resolved against the loaded model, so the donor's module is the only truthful
source (an ``auto`` request that resolved ``off`` on this checkpoint leaves a
dense donor, which arms nothing here). No-op for a dense donor and for a
member already fp8.

Reshape the input embedding into empty quantized placeholders on the WARM
flat-dump reload path. Returns the CONCRETE mode reached, so the caller
accounts the armed bytes exactly when one was armed.

The warm cold-dump was captured with the embed already quantized, so the
freshly-built ``skip_weight_load`` graph must present the SAME quantized
slots (no bf16 ``weight``) for the dump to match and fill — landing the KV
win on the warm path too (the bf16 table is never allocated). Takes the same
:func:`resolve_embed_mode` decision as the eager :func:`apply_embed_quant`,
so the two paths agree by construction; a no-op for a resolved ``off`` or an
already-quantized (donor-shared) embed.

EXL3 trellis-quantization backend for arbi-serve.

Importing this subpackage registers the backend with the global
quant-backend registry. Heavy upstream deps (``exllamav3``, the JIT-
compiled ``exllamav3_ext`` CUDA kernel) only load on first detected
EXL3 checkpoint, so a vanilla bf16 run never pays the import cost.

Activation prep for the int8 EXL3 prefill GEMM: rotate, relabel, quantise.

ONE PASS, and that is the whole point. #1861's K-major problem resolves into a
16-element permutation of the CONTRACTION axis (see
:mod:`.kmajor_layout`), which is free only if it rides a pass the activation
already takes. Leg A already spends one pass turning ``x`` into an fp16
``A_had`` scratch; :func:`a_prep` replaces that pass and writes one byte per
element instead of two, so the activation traffic goes down rather than up. A
standalone permute kernel would hand back exactly what the finding saved.

THE CONTRACT WITH THE GEMM. For a linear with ``suh``/``svh`` metadata, the
int8 leg computes

    acc[m, n]  = sum_L  a_q[m, L] * w_q[SIGMA(k), n]        (int32)
    y_had[m,n] = acc[m, n] * a_scale[m] * W_SCALE

and the existing output stage (``had_r_128(y, None, svh)``) is unchanged.
``a_scale`` is per TOKEN and ``W_SCALE`` is a scalar CONSTANT of the codebook,
not a tensor: exl3's mul1 decodes to ``(bytesum - 510) * k_inv`` with
``k_inv = half(0x1eee)``, so an integer decode of ``sat((bytesum - 510 + 1) >> 2)``
carries ``W_SCALE = 4 * k_inv``. There is no per-column and no per-group weight
scale to load, which is what makes the epilogue a single fused multiply.

Per TOKEN rather than per token-and-k-group on purpose: a scale that varied
along k could not be folded after an int32 accumulation over k, so the GEMM
would have to break its accumulator. The 128-wide Hadamard is what makes one
scale per row sufficient -- it is load-bearing, not a flag.

AND IT IS ALREADY WHERE A GROUP-WISE ROTATION WOULD GO. The rotation runs
immediately before the absmax and its width IS the width a per-group scale
groups by (``A_PREP_HADAMARD_BLOCK`` == ``int8_kernel.A_GROUP``), so the trick
of rotating inside the quantisation group so that no single heavy channel can
set the group's scale is this pass. A second rotation at that width does not
add to it: the normalised Sylvester Hadamard is an involution, so it is the
first one's inverse and hands the quantiser back the raw activation. What the
rotation cannot reach is the spread ACROSS a row's blocks, because it does not
mix between them; that is what the per-token amax pays for and what a
per-group scale removes. ``tools/int8_gemm/rotation_headroom.py`` is the
measurement and ``tests/test_exl3_had_basis.py`` pins the two widths together.
See #1861.

The scale table, whose SHAPE selects the quantiser.

``[rows]`` is the shipped per-token amax.  ``[K/128][rows]`` is one scale
per (token, 128-wide Hadamard block) -- the boundary the rotation already
establishes, since it equalises magnitudes inside a block and not across
blocks.  The kernel reads the extent and refuses anything else, so there is
no flag to set on one side and forget on the other.

K-major, not row-major: the GEMM's mainloop drains one group at a time and
a warp's 32 lanes cover 8 consecutive rows, so this makes that read
contiguous where a ``[rows][K/128]`` table would make it a strided gather.

JIT-compile (or pull from cache) the a_prep extension.

Registered statically into ``torch.ops.arbi_serve_exl3_a_prep`` at dlopen,
the way the Marlin loader does it, so a captured region never re-enters the
Python build path.

Compile (or pull from cache) the a_prep extension, from OUTSIDE a step.

The public name for what every entry point here calls first. A caller that
knows it is about to enter a served step -- or a cudagraph capture -- warms
the build here instead of discovering ``ninja`` inside one; the int8 leg's
boot seam (:func:`~arbi_serve.weight_quant.exl3.int8_kernel.load_int8_kernel`)
is that caller.

``x`` (fp16, rows x in_features) -> ``(int8 activation, scale)``.

``group`` asks for one scale per (token, 128-wide Hadamard block) instead of
one per token; see :func:`_scale_buf` for why the shape is the switch.

Same pass, emitting the GEMM's ``mma.m16n8k32.s8`` A-fragment buffer.

The int8 GEMM does not read a row-major ``[M, K]`` activation: it reads
``A_frag[m/16][k'/32][lane][16]``, already in fragment order, so a
row-major drop-in would be silently misread. A lane of this kernel already
holds exactly one of those four-byte slots, so the fragment layout costs an
address and nothing else -- no shuffle, no extra pass, no extra byte.

``M`` is padded up to the GEMM's ``TILE_M`` and the padding rows are
written as zero by the same launch, so the tail tile never contracts
against stale scratch. ``out``/``scale`` may be supplied by the caller,
which is what a cudagraph-captured region needs: an allocation inside the
captured region freezes at a stale pointer.

One input, G rotated-and-quantised outputs, reading ``x`` from HBM once.

The served graph has three co-input groups -- ``q|k|v`` (G=3),
``in_proj_qkv|in_proj_z`` (G=2) and ``gate|up`` (G=2) -- and none of them is
weight-fused, because ``suh`` differs per consumer and that is precisely
why the narrow-N shapes exist. Grouping is a BLOCK-level fact: a per-linear
op sees one consumer at a time and cannot express it, so the signature is
multi-output from the start rather than retrofitted.

ONLY THE HBM READ IS SHARED, and that is a measured limit rather than a
conservative choice. ``suh``'s sign half is shared exactly across a group,
but its magnitude half is not, and a diagonal does not commute with a
Hadamard -- the one escape, a mismatch constant within a 128-block, is
closed because ~95% of the variation is inside the block. So the rotation
and the quantiser are redone per consumer; ``x`` is cached in registers and
read once.

``out``/``scale`` may be supplied, which is what a captured region needs.

``silu(gate) * up`` -> G int8 activations. The ``down_proj`` site.

The largest single item in the int8 leg's traffic budget: 4 B/elem of the
MLP intermediate, which at ``I=17408`` over 64 layers is **4.456 MB/token**.

``Qwen3_5RMSNorm(x)`` -> G int8 activations. The input-norm sites.

``folded`` is ``Qwen3_5RMSNorm.folded_weight`` -- ``1 + weight`` in fp32,
precomputed by ``prepare_for_compile``. The row reduction runs ONCE for the
whole group, which is the half of this that grouping actually buys.

Fused residual-add + ``Qwen3_5RMSNorm`` -> G int8. The post-attn site.

``residual`` IS MUTATED IN PLACE to ``residual + x``, rounded to its own
dtype -- the contract ``fused_add_rms_norm`` already declares and what the
next layer reads. The norm itself consumes the UNROUNDED fp32 sum, which is
what the compiled graph does and is not what the eager source reads.

The fp16 activation the SERVED norm hands a quantized projection.

Returns ``(activation, new_residual)``. The residual is the sum rounded to
``x``'s dtype; the activation is computed from the UNROUNDED sum. Those are
two different numbers and the compiled graph produces both from one add --
``tools/prologue_fusion/probe_norm_semantics.py`` is the measurement.

Byte offset of logical element ``(row, k_logical)`` in the A-fragment buffer.

A restatement of the GEMM's documented ``A_frag`` layout, NOT of this
kernel's store arithmetic, so a test can hold the two against each other
instead of against one another's prose.

The same result via torch, for a test to compare against.

``rotated`` lets a caller supply the rotation from exllamav3's own
``had_r_128`` so the comparison isolates the relabel and the quantiser from
the Hadamard's fp16 rounding, which the fused kernel deliberately skips.

``group`` reproduces the per-128-block quantiser and returns the K-major
``[K/128][rows]`` table.  The amax is taken AFTER the relabel, as the kernel
takes it -- which is sound only because SIGMA permutes inside a 16-block and
therefore never moves an element across a 128-block boundary.

``(x.view(-1, 128) @ H_128) / sqrt(128)``, Sylvester order.

Which IS the kernel's order: it splits the index as ``4*lane + j``, runs
``H_4`` on ``j`` and a butterfly ``H_32`` on ``lane``, and
``H_4[j][j'] * H_32[lane][lane'] = (-1)^popcount(p & p')``. The test suite
does not take that on trust -- it simulates the kernel's two stages and
compares the composite matrix element by element.

What the kernel's two shuffles and byte permute produce, per lane.

``[lane][byte] -> (source lane, source element)``. The kernel's store is the
only place SIGMA is applied, and an off-by-one there produces plausible
wrong output rather than a crash, so the map is derived here from the same
shuffle arithmetic the kernel uses and checked against
:data:`~.kmajor_layout.SIGMA` by a CPU test.

EXL3 :class:`QuantBackend` implementation.

The class implements the protocol — detection from ``*.trellis`` keys,
mapping HF-canonical safetensors paths to the LinearEXL3 mirror class,
and reading the trellis, Hadamard scales, and optional codebook seed per linear.

EXL3-specific assumptions encoded here:

  - Detection key is the ``.trellis`` suffix.
  - Tensor schema per linear: ``{trellis (i16), suh (f16), svh (f16),
    mcg|mul1 (optional seed)}``.
  - lm_head must be untied from embed_tokens (the trellis lm_head
    cannot be aliased to the dense embedding).
  - Block alignment is 16 elements along each axis (the trellis
    quantization unit).

Return the EXL3 quantized Linear subclass mirroring ``dense_cls``.

``MergedColumnParallelLinear`` (e.g. the GDN ``in_proj_qkv``, which
EXL3 ships as a single fused trellis tensor) is supported at TP=1
(single column-parallel linear) and at TP>1 (per-sub-block fused
split inside :class:`EXL3MergedColumnParallelLinear`).

``VocabParallelLMHead`` is the lm_head's dense class (EXL3 always
quantizes lm_head — see :meth:`lm_head_untie_required`). It is a
``ColumnParallelLinear`` subclass, so an exact-class table lookup
misses it; it is mapped explicitly below. Any other unregistered
dense class raises ``RuntimeError``.

Read trellis / suh / svh (+ optional codebook seed) and bind
via :meth:`exl3_load`.

``dtype=None`` is critical — the loader preserves on-disk
dtypes (int16 for trellis, float16 for suh / svh). Casting
trellis to bf16 would corrupt the codebook.

A checkpoint quantized with a non-default codebook ships a
per-linear scalar seed under ``.mcg`` (mcg codebook) or ``.mul1``
(mul1 codebook). Decoding such a trellis with the default
codebook silently produces wrong weights, so the seed is read and
threaded into the kernel when present.

Which is exactly why the seeds are also checked against what the
checkpoint DECLARES before any of them is read -- see
:meth:`_check_checkpoint_identity`. The kernel routes on seed
presence and cannot notice a disagreement; nothing downstream can
either, because a trellis decoded with the wrong codebook has the
right dtype and extent.

The pack's ``quantization_config``, cached per model directory.

``quantization_config.json`` first because exllamav3's converter writes
it last and it is the copy that carries the extra metadata; the
``config.json`` section is the fallback and the only one some packs
ship. An unreadable or absent config is an EMPTY declaration, not a
failure -- a pack that declares nothing is the pre-tag state, and
refusing it would refuse every checkpoint that exists.

Refuse a pack whose declared numerics are not the ones we serve.

Both facts checked here change what the weights DECODE TO while
leaving every tensor shape intact, which is why neither can be caught
downstream: the trellis, ``suh`` and ``svh`` of a pack in a 512-wide
basis, or quantized against another codebook, load without complaint
and contract cleanly against activations from a different vector
space. The result is a model that runs and answers wrongly.

The check is on the DECLARATION, not on the tensors, and that is the
point. The kernel already infers the codebook from seed presence, so
inferring it again here would only re-derive what it will do anyway;
comparing the inference against what the quantizer WROTE DOWN is what
turns a silent substitution into a refusal.

Refuse a codebook the kernel cannot decode, or a seed that
contradicts the declaration.

Two seeds on one linear is refused ahead of the declaration: the
kernel takes them as independent booleans and the pair is not a
codebook it has, so whichever one wins is a coin toss that the
declaration cannot adjudicate.

Return ``None``: a trellis payload has no stackable representation.

EXL3 stores each linear as its own codebook-indexed trellis with
per-linear ``suh`` / ``svh`` scale vectors, so there is no layout in
which the routed experts share one packed tensor. The routed experts
stay per-expert modules.

Record one ``binder`` — reading ONLY the safetensors HEADERS (no
weight bytes).

The binder needs each buffer's post-TP-slice numel before any data
is loaded. The trellis codebook width ``K`` and the optional codebook
seed sizes come from the on-disk header (``byte_range`` parses the
shard header without materializing the tensor); the rank-local
feature dims are already on ``lin`` (the parallel mirror set them at
construction).

exl3's codebook numbering, in one place, with no imports.

The numbers are exllamav3's, not ours: `quant/codebook.cuh` dispatches
`decode_3inst<cb>` on exactly these values. They are mirrored by the
`EXL3_I8_CB_*` defines in `csrc/exl3_i8_gemm.cu`, and
`tests/test_exl3_int8_codebook_refusal.py` reads both and asserts they agree,
so the mirror cannot drift the way the dp4a addend once did.

It ships inside `arbi_serve` because the served int8 adapter
(:mod:`~arbi_serve.weight_quant.exl3.int8_kernel`) names a codebook on every
call and `tools` is not present in the serving image.

This module imports nothing of its own so that both the checkpoint loader
(`tools/int8_gemm/ckpt`, which pulls in safetensors) and the extension loader
(which JIT-builds CUDA) can name a codebook without dragging in the other.

EXL3 trellis-GEMM as a ``torch.library`` custom op (compile / cudagraph safe).

Why this exists
---------------
arbi-serve compiles every decoder layer with ``torch.compile`` (per-layer
piecewise strategy, see :mod:`arbi_serve.compile`) and replays the result
under a CUDA graph. Both passes require that the EXL3 trellis-dequant GEMM:

  1. be **traceable** — Dynamo cannot step into the pybind C++
     ``exllamav3_ext.BC_LinearEXL3`` method, so with ``fullgraph=True`` the
     raw call site graph-breaks / errors. Wrapping the call as a
     ``torch.library.custom_op`` with a ``register_fake`` turns it into an
     opaque-but-typed node that Inductor plans allocations around.
  2. be **allocation-free in the captured region** — exllamav3's
     ``LinearEXL3.forward`` routes through ``bc.run_alloc`` which allocates
     BOTH the output ``y`` *and* (for bsz>1) an ``A_had`` scratch via
     ``at::empty_like`` on every call. Per-call device allocation is the
     classic cudagraph blocker. We instead call the no-alloc ``bc.run``
     path with a **caller-allocated** output (the custom op's fake lets
     Inductor allocate it once and reuse the buffer across graph replays).

exllamav3's own cudagraph mechanism
-----------------------------------
exllamav3 captures with the raw CUDA-graph C++ API (``BC_LinearEXL3::run_gr``
takes a ``Graph*`` and, for bsz==1, threads a **preallocated** ``xh``
scratch buffer — a member set once at ``LinearEXL3.__init__`` via
``g_tensor_cache.get(device, (1, in_features), out_dtype)`` — into
``exl3_gemm_gr`` so nothing allocates during capture). ``run_alloc`` is the
*eager convenience* wrapper that does allocate. arbi-serve doesn't use
exllamav3's C++ Graph; it uses torch.compile + torch.cuda.graph. The
equivalent for that path is: no per-call alloc inside the op + a registered
fake. The no-alloc ``bc.run`` (a thin wrapper over ``run_gr(.., nullptr)``)
gives us exactly that — the bsz==1 path reuses the preallocated ``xh``
member, and the bsz>1 path's ``A_had`` scratch is allocated by the kernel
wrapper only outside any capture (warmup) and never re-sized for a fixed
decode shape.

The ``run_alloc`` path additionally heap-corrupts (``free(): invalid
pointer``) on multi-row input for some K builds; the no-alloc ``bc.run``
path is correct there too. So this wrapper is both a cudagraph-enabler and
a correctness fix for the multi-token (prefill / MTP-verify) shape.

Register ``arbi_serve::exl3_linear_gemm`` (idempotent).

Called lazily from :meth:`EXL3LinearBase._build_inner` — i.e. the
first time an EXL3 linear is actually bound — so a vanilla bf16 run
never imports exllamav3 nor registers this op.

Rows above which upstream dispatches THIS SHAPE to reconstruct+hgemm.

``out_features`` is the KERNEL width the op was called with (the trellis's
padded one, and under TP this rank's local shard). Upstream asks the same
question of its DECLARED width; the two agree because the narrow-N bound is
a multiple of 128 and both padding and TP sharding move the width in whole
128-blocks, so neither can cross it.

``hgemm`` as the extension MODULE exposes it, never a caller's handle.

Deliberately not ``ext.hgemm``. The capability is a property of the BINARY,
and reading it off whatever object the caller passes is what made this
probe maskable: a bench that wraps ``ext.hgemm`` to count calls hands back
a plain Python closure whose ``__doc__`` is ``None``, the probe concludes
"this build has no acc_mode", and the leg serves fp32 under the armed
flag's name. An observer installed to measure the accumulator is then an
observer that disables it, and the run reports the baseline as the result.
That has already cost one investigation a wrong answer, so the probe asks
the module directly and a wrapper cannot get between them.

Whether the loaded exllamav3's ``hgemm`` takes a per-call accumulator mode.

A property of the BINARY, so it is probed once and cached. Without it the
fp16-accumulate leg is simply unavailable and the leg keeps fp32 — passing
a defaulted-in-C++ argument positionally to a pybind overload that does not
declare it is a TypeError, not a no-op.

The only alternative a stock build offers is
``torch.backends.cuda.matmul.allow_fp16_accumulation``, which is
PROCESS-wide: it would also retune the MoE routing, attention-gate and MLA
index GEMMs that share this entry point, and those are exactly the small,
ill-conditioned shapes an fp16 accumulator ruins. So the capability is a
per-call argument or it is nothing.

``ext`` is still taken, and still inspected — not for the capability, but
to say out loud when the handle the leg will actually CALL is not the
binding the probe just read. A transparent ``*args`` proxy forwards
``acc_mode`` fine; a wrapper with a fixed three-argument signature raises
``TypeError`` on the first armed call. Either way the operator should know
an instrument is in the path.

Whether this inner carries the static fp32-accumulate pin.

``None`` (a caller with no inner in hand) reads as NOT pinned, because the
pin is a positive property of a linear that was marked — a helper that
could not find the linear must not be able to arm one either, which is why
the only caller that matters (:func:`_prefill_acc_args`) is handed the
inner the op already resolved.

Whether this linear keeps the fp32 accumulator while the flag is armed.

Two clauses, and they are not the same KIND of rule.

The first is the COMPONENT exclusion: an inner carrying
:data:`_PREFILL_FP32_ATTR` keeps fp32 unconditionally. It is stamped at
bind on MLA's low-rank index GEMMs, whose output is a latent the whole
attention then reads through, and there is no exl3 DeepSeek checkpoint to
measure one on — which is precisely why they are excluded rather than
swept. Being unmeasurable is not being safe. See
``weight_quant.loader._mla_keeps_fp32_prefill`` for the policy.

The second is the per-linear-class POLICY, uniform by construction over
every linear the first clause admits, and the swept alternatives are why.

The hypothesis worth sweeping was that an fp16 accumulator's error tracks
the dot product's CONDITION NUMBER, sum|a_i b_i| / |sum a_i b_i|, so that
holding the worst-conditioned projection classes on fp32 would recover most
of the KL for a small slice of the speed. It does track the condition
number -- but no exclusion measured better than uniform END TO END, while
every one of them cost speed. RTX 4090, GMU 1.0, 16384-token prefill at chunk 2048, 10 prompts from
``build_realkl_search_basket``, decode-step GDN
The medians sit inside one band and the means
are carried by single near-tie steps, so nothing here buys back the speed
it spends.

Kept as a FUNCTION rather than deleted because it is the seam
``tools/exl3_prefill_accumulator_ladder.py`` drives to re-run that sweep on
another checkpoint: the classes are a property of the model, and a model
whose residual stream is narrower or whose MLP is deeper could land
somewhere else. It is a static property of the linear -- never the row
count -- for the reason :func:`_prefill_pins_leg_a` is: a numerics class
that moved with the row count would make one prompt's output a function of
who it batched with.

A sweep replaces this whole function, so it also replaces the first clause
— which is why :func:`_prefill_acc_args` re-checks the pin ITSELF rather
than trusting this call. That duplication is deliberate and is the
difference between an exclusion and a default.

WHICH cuBLAS accumulator the prefill leg runs a GEMM with.

An identity, not a shape. The two plans agree on every dimension and
disagree on the arithmetic, so nothing downstream can tell them apart:
changing only the compute type moves this model's output as far as
switching legs entirely. Anything that can carry one plan while another was
declared -- a cudagraph most of all, which records the kernel that was
launched while it captured and replays that one forever -- has to carry the
TAG, so a consumer checks what it got instead of checking that it got
something the right size.

Whether this BUILD's ``hgemm`` takes a per-call accumulator mode.

Asks :func:`_hgemm_acc_mode_supported` through the extension module's own
binding, so the proxy warning that helper raises cannot fire on a question
about the binary.

A build with no extension to ask answers ``False``, which is the truth
about it rather than a swallowed error: no ``hgemm`` is no ``acc_mode``.
This module imports without exllamav3 (every fork import in it is lazy)
and the flag defaults ON, so a bf16 or AWQ deployment reaches this asking
a question about a binary it does not have. On a deployment that DOES
serve EXL3 the extension is imported at bind, long before anything asks
this, so a genuinely broken build has already failed loudly elsewhere.

The plan a PREFILL forward is supposed to resolve on this build.

The flag AND the capability, because a build whose ``hgemm`` takes no
``acc_mode`` cannot honour an armed flag and keeps fp32 -- so a consumer
comparing against the flag alone would refuse a correct graph on that
build. Per-linear exclusions are deliberately NOT folded in: they are a
property of individual linears, this is the property of the forward.

A captured graph carries an accumulator the declaring forward did not.

Its own class because it is not a capture failure and must not be handled
as one: the graph recorded FINE, it recorded the wrong arithmetic. Every
replay of it serves that arithmetic, silently, for the life of the process.

Assert a cudagraph capture records the accumulator ``expected``.

A cudagraph replays the kernel that was launched while it captured. The
prefill accumulator is chosen per call from the forward's DECLARED PHASE
(:func:`prefill_phase`), so a capture whose forward declared nothing bakes
the fp32 plan into a graph that then serves every prefill -- under a flag
that reads armed, and with no shape, dtype or count anywhere that differs.
This is the check that turns that into a refusal.

``expected`` is the plan the capture INTENDS, which the caller states
rather than infers; :func:`declared_prefill_accum_plan` is what a prefill
capture passes. A capture of a class that legitimately keeps fp32 -- decode
-- passes :attr:`PrefillAccumPlan.FP32`, or does not wrap at all.

Silent when the leg never ran: a bucket whose rows all sit below the
reconstruct threshold records no leg-B call, and there is no plan to
disagree about. Silent, too, when every call it did make was an exclusion,
because that is a declared fp32 and not a substituted one.

Extra ``hgemm`` args selecting THIS call's cuBLAS compute type.

``(1,)`` asks the fork for an fp16 accumulator; ``()`` leaves the call
exactly as the pre-fork signature spells it, so a build whose ``hgemm``
takes three arguments still runs.

Scoped three ways, and the PHASE is the one that carries the safety
argument. Reaching this leg is a question of ROW COUNT, not of phase:
``rows > _auto_reconstruct_threshold(out_features)``, and that threshold is
64 on a narrow ``out_features`` — on the served 27B, the ``k_proj``/
``v_proj`` of every full-attention layer. A decode step wide enough clears
64 rows (B >= 65, or B >= 13 on an MTP K=4 verify slate) and lands those
linears here. Ungated they would take the fp16 accumulator and write the KV
the drafter reads, inside the one leg the accept-invariance proof covers.
So the accumulator asks the STEP whether it is prefill
(:func:`prefill_phase`, marked from ``ScheduledBatch.is_prefill``) and
declines otherwise, and the unmarked default is fp32 — a forward that
reaches this op without declaring its phase gets the conservative
accumulator, never the fast one.

The phase is a separate channel from :data:`_PREFILL_NUMERICS`, which is
marked only when ``exl3_prefill_row_invariant`` asks for it: gating on that
would disarm this lever on every deployment that has not opted into an
unrelated reproducibility guarantee.

Then the leg, and within it the linears the policy admits. ``hgemm`` is
shared, and its other callers are small or ill-conditioned GEMMs for which
the fp32 accumulator is not a luxury; a process-wide switch would take them
with it.

And within the leg, the COMPONENT exclusion: an inner carrying the
bind-time pin (MLA's low-rank index GEMMs — see
``weight_quant.loader._mla_keeps_fp32_prefill``) is refused here, ahead of
:func:`_prefill_keeps_fp32`, because a policy sweep replaces that function
wholesale and would take the exclusion with it. The pin is read from the
inner the OP ALREADY RESOLVED, so it cannot be re-derived — or lost — on
the way in.

``(plan, excluded)`` for this call. The ONE place the plan is decided.

``excluded`` says the fp32 was DECLARED -- a pinned linear or the class
policy -- rather than substituted. A consumer checking a captured graph
needs that split: fp32 for a stated reason is not the defect, fp32 because
nobody declared the phase is.

Rotate the activation into the basis the WEIGHTS were quantized in.

The INPUT axis rotation: sign-flip by ``suh``, then the Hadamard over the
contraction axis. Its width is a property of how the checkpoint was
quantized, and it is a separate axis from the output rotation below --
widening one says nothing about the other.

A named seam rather than a bare ``ext.had_r_128`` because ONE fork entry
point serves both rotations. The call sites are distinguished today only by
which of ``pre_scale``/``post_scale`` is non-``None``, which is an argument
position, not a statement of intent; a reader tracing a basis change has to
reconstruct the direction from the scale vector's name. These two functions
say it instead.

Rotate the GEMM result out of the weights' OUTPUT basis, in place.

The output-axis counterpart of :func:`_rotate_input_basis`; see there for
why the two are named rather than left as one entry point.

Refuse a rotation the basis does not divide, naming WHICH axis.

The fork checks the same divisibility and raises from C++ with neither the
axis nor the width in the message. Both rotations reach that one entry
point, so its refusal cannot say which of a linear's two axes is wrong --
and on a padded or TP-sharded linear those are different numbers arriving
from different code. Named here, the refusal points at the axis.

Upstream's fused gate surface, or ``None`` when this fork exports none.

``None`` is the answer for a fork that predates the surface, recorded in
:data:`_FUSED_RECONSTRUCT_FORK_GAP` and folded into
:func:`_fused_reconstruct_buildable` so no caller has to ask twice.

Rows at or above which THIS SHAPE takes the fused reconstruct.

``out_features`` is the KERNEL width the op was called with (padded, and
under TP this rank's local shard); upstream asks the same question of its
DECLARED width. The two agree because upstream's class bounds are multiples
of the Hadamard block and both padding and TP sharding move the width in
whole blocks, so neither can cross one — the same argument that makes
:func:`_auto_reconstruct_threshold` safe to ask on kernel dims.

Only meaningful on a build that CAN serve the leg — ask
:func:`_fused_reconstruct_enabled` first. On a fork that exports no fused
surface there is no threshold to report and this says so rather than
inventing one: a stand-in number here would be the arbi-side copy of a
value the import exists to keep single. The served path never asks (see
:func:`_use_fused_reconstruct`); the benches do, behind that same gate.

Whether the INSTALLED fork can serve the fused leg at all.

Three properties of the loaded build, and every one of them is asked as a
CAPABILITY question so a fork older than the pin answers "no" instead of
raising: upstream's own kill switch, the gate surface
(:func:`_fused_reconstruct_consts`), and the extension symbol. The whole
fork-facing surface of this leg is resolved HERE, which is what lets
:func:`_use_fused_reconstruct` reach it only after the flag.

Upstream's kill switch is read first: a leg turned off on purpose is not a
skew, and must not record one.

The fused surface the installed fork does not export, or ``""``.

Resolves the build probe first, so the answer is never "nobody has asked
yet". A non-empty string is a PIN/IMAGE SKEW: this process serves the
standalone reconstruct leg for a reason no operator chose, and nothing
else about the boot differs — which is why the boot ledger has to carry
it (see :data:`_FUSED_RECONSTRUCT_FORK_GAP`).

Empty for every other reason the leg can be off. ``exl3_fused_reconstruct``
and ``EXL3_NO_FUSED_RECONSTRUCT`` are requests, not degradations, and a
build that simply has no ``reconstruct_had_slice`` never claimed the leg.

Whether this large-M call takes the fused original-basis reconstruct.

``in_features`` / ``n`` are KERNEL dims (the op is called with the trellis's
padded widths), and under TP they are this rank's LOCAL widths — EXL3 shards
on whole Hadamard blocks, so a shard is 128-divisible exactly when the
unsharded axis is.

This gate is ROW-KEYED on two of the three output-width classes -- a third
algorithm choice nested inside leg B, live for every linear the prefill pin
does not cover (see :func:`_prefill_pins_leg_a`). On the WIDE single-slice
class the threshold is 0, so on those linears there is no row-keyed branch
here at all: the variant is decided by weight geometry alone and a call's
row count cannot change which algorithm runs.

Nothing below reads the row count except the threshold.
``reconstruct_had_slice`` takes no activation and no ``M``; its
preconditions are all weight geometry (both axes divisible by the Hadamard
block, ``n_offset`` divisible by it, ``suh`` and ``svh`` at least as long
as the slice) and the kernel ``TORCH_CHECK``s every one of them, so the
divisibility clauses guard a loud error rather than a silent misread. The
row count reaches only ``hgemm``, which both variants call identically.
Neither kernel checks that its inputs are CONTIGUOUS and both index from
``data_ptr``; both variants satisfy it by construction (the op flattens
``x`` with ``.contiguous()``, the weight is a view of the reconstruct
scratch from its base, and ``svh[n_start:]`` is a slice of a contiguous
1-D tensor), and none of that depends on where the threshold sits.

``pin_to_threshold`` drops the row clause: the caller declares that this
call is served at THIS SHAPE'S OWN threshold, so the variant is decided by
weight geometry alone and ``rows`` cannot move it. That is what the
numerics-class row split needs (see the call site) — the same answer the
old spelling produced by passing ``_fused_reconstruct_threshold(n)`` in as
``rows``, minus the caller's own reach into the fork.

THE ROW-INVARIANCE CLASS PINS IT TOO, and this gate reads that itself
rather than trusting a caller to pass it. ``exl3_prefill_row_invariant``
is the guarantee speculative decoding's losslessness rests on, and it
pinned the LEG (:func:`_prefill_pins_leg_a`) while leaving the leg-B
VARIANT chosen per call from the row count — so a user who asked for
reproducibility across concurrency still got two algorithms. They are not
two roundings of one computation: the fused variant evaluates
``x . (H_k . W_hat . H_n)`` with both Hadamards folded into the weight,
the standalone one ``((x . H_k) . W_hat) . H_n`` with them on the
activation and the output. That is a REASSOCIATION, equal only in exact
arithmetic, and this tree measures the gap at 10036/15360 output elements
moved, max |delta| 1.95e-3, between fused totals of 512 and 2048 rows with
the hgemm ``M`` already pinned.

Read here and not at the call site because a pin the caller has to
remember is a pin one caller will forget: the mixed-capture split already
passes ``pin_to_threshold`` and the row-invariance class did not, which is
exactly that failure. One gate, one question.

Whether ``inner`` emits FEWER columns than its trellis stores.

True only for an output window over another linear's weights
(:class:`~arbi_serve.weight_quant.exl3.linear.EXL3OutputPrefixLinear`).
The one place that distinction is decided, because three seams have to
agree on it: the leg dispatch (only ``exl3_gemm`` takes an output bound),
the reconstruct-scratch reserve (a window can never reach that leg, so it
must not size the buffer off its wide trellis) and the kernel-shape pin.

Return one inner's ``(in_features, out_features)`` in KERNEL space.

Kernel space (the trellis's own 16-element-quantized dims) is what the op
reconstructs against — a Hadamard-padded checkpoint's kernel dims exceed
the module's declared ones. Read off the trellis so the answer matches the
tensor the kernel actually decodes, EXCEPT on the output axis of a window,
where the emitted width is narrower than the weight and is what every
consumer of these dims is asking about. The descriptor's own fields are the
fallback.

Declare the widest row count each bound linear's call sites can produce.

``default_rows`` applies to every registered inner: the step token budget,
which bounds any linear reached from inside a model forward. ``head_inners``
names the logits-head linears, whose call sites gather ONE row per sequence
(or one per speculative slate entry) rather than one per token, and
``head_rows`` is that narrower bound.

The bound is what lets :func:`reconstruct_scratch_elems` skip a linear that
can never cross its own ``auto_reconstruct_threshold``: such a linear never
reaches the reconstruct leg, so it never draws a weight-shaped slab, so
sizing the shared scratch to it reserves bytes nothing can spend. Both
figures are the caller's UPPER bounds — a bound that under-states a call
site pushes the linear onto the per-call fallback in
:func:`_reconstruct_slab`, which is correct but allocates outside the boot
budget.

Whether any call site can hand ``inner`` enough rows to take leg B.

``True`` when no bound is declared: the reconstruct leg is reachable unless
something has established otherwise, which keeps a registry populated
outside the engine (tools, tests) sized exactly as it was before.

fp16 elements the widest reconstruct among ``inners`` asks for.

``in_features x min(out_features, MAX_RECONSTRUCT_SLICE_N)`` — the wide
(N-sliced) leg never materialises more than one slice at a time, so the
slice width caps each term. Linears that cannot reach the leg contribute
nothing (see :func:`_reaches_reconstruct`).

fp16 elements the widest UNSERVED EXL3 linear's reconstruct needs.

:func:`_scratch_need` over every bound linear that carries no scratch yet
and can reach the reconstruct leg. 0 when nothing is waiting, which is both
the "no EXL3 linear is bound" case (keeping this exllamav3-import-free on a
vanilla run) and the "already reserved" case — and now also the case where
every waiting linear is bounded below its own threshold.

Scoped to the unattached set rather than the whole registry because the
registry outlives a member: a second member's reservation must be sized
from ITS OWN linears, not from a parked member's.

Reserve (or grow, pre-capture) THIS member's reconstruct scratch.

``current`` is the buffer this member already reserved at an earlier boot
seam, or ``None`` on its first. Returns ``(buffer, bytes)`` — the caller
holds the buffer for the member's lifetime, which is what keeps the pool
memory the linears point at mapped. The caller also owns the pool routing:
allocate inside the named pool whose lifetime matches an address captured
graphs may bake, never in the default allocator.

Attaches the buffer to every linear bound since the last reservation. That
unattached set is one member's worth: a member's quant linears are
constructed fresh and register their own inners, so a second member never
inherits the first's buffer — which it must not, because the first's pool
is parked while the second builds.

Idempotent: a second call with nothing newly bound is a no-op. A call that
finds a WIDER requirement — a drafter attached after the base model's
reservation — allocates one grown buffer and re-points this member's
already-attached linears at it, preserving the one-buffer-per-member
property. Safe because both seams run before the capture sweep, so no
address a graph has baked can move.

Boot-time by contract. The serving path must never be the first caller: a
lazily grown buffer allocates out of headroom the KV sizer has already
given away, which is how an unbounded staging allocation has OOMed a rank
mid-collective before.

``need`` fp16 elements for one reconstruct, from ``inner``'s scratch.

Returned flat; the caller views it at the weight's ``(in_features, cols)``
shape. The view always starts at the buffer's base, so its alignment is the
allocator block's — identical to the per-call ``torch.empty`` this replaces,
which is what keeps the ``hgemm`` that reads it bit-identical.

Read off the inner rather than a module global so a linear can only ever
reach the buffer its OWN member reserved (see the section header).

Falls back to a per-call allocation (the pre-reservation behaviour) when no
buffer covers the request: no reservation ran, the request is on another
device, or a shape wider than the reservation reached the op. The fallback
is correct but unbudgeted, so it warns once naming the reservation seam.

The declared class of the running forward, or ``StepClass.DECODE``.

``None`` (nothing declared) resolves to DECODE rather than raising: the op
is reachable from paths that are not the canonical forward site (warmup
probes, the fused-reconstruct agreement test, a bench harness calling the
op directly), and every one of those wants today's legs.

Slate width of the running forward, or 0 when nothing declared one.

0 is not a measurement, it is "unknown", and every rule that reads this
treats it as below threshold — so a path that reaches the op outside the
declaration seam (a warmup probe, a bench harness calling the op directly)
keeps the shipped leg rather than being routed on a width nobody set.

Declare the enclosed forward's STEP CLASS and slate width.

Restores the previous values on exit, including on exception — a failed
prefill must not leave the process serving decode under a prefill class,
which is the direction that matters because decode writes the KV the
drafter reads.

``rows`` is the forward's flat token count. Keyword-only and defaulted so
the class alone stays a complete declaration for the callers that route on
nothing else; the production seam always passes it.

Why THIS forward declared no step class, once one previously did.

Three states reach here and they want three different actions, so they get
three names. Reading them as one — which they were, until a within-process
A/B produced the second and reported it as the first — makes an operator's
own switch look like a bypass.

``forward declared no step class``
    The declaration gate is ON and a forward still declared nothing, so a
    call site bypassed :func:`~arbi_serve.runtime.forward_declaration.
    declare_forward`. THE ONE THAT IS A BUG, and the state the branch was
    written for: whatever this forward is recording — a cudagraph capture,
    most consequentially — routes on the shipped legs while the flag reads
    armed.

``int8 leg disarmed after a class was declared``
    The operator turned ``exl3_int8_gemm`` off. ``scope="runtime"`` makes
    that a supported live action, the latch is monotone, and the
    declaration gate follows the flag — so every later forward lands here.
    Expected, and it is also what makes an OFF arm a usable CONTROL: the
    same calls that fired on the armed arm are counted, refused, on this
    one.

The gate and the leg read the SAME flag, so they cannot disagree: the
    flag is a boolean, and `int8_gemm_armed` resolves it rather than a
    separate arch default. While it was a tri-state they could, and that
    skew presented as "the flag is on and nothing fires".

The counter for THIS forward's CLASS and window — captured or eager.

Two axes, six counters, and neither axis can be collapsed. The window
decides whether a zero is evidence at all (a captured leg fires once, at
record time); the CLASS decides which flag the reading is about. A DECODE
forward has its own pair since ``exl3_int8_decode`` made it a rule: while
it was a categorical refusal it was counted with prefill as ``off-class``.

``is_current_stream_capturing`` is a cheap driver query, and it is only
reached on a forward that already declared a step class (i.e. the flag is
armed). An unarmed deployment never calls this.

The binding the int8 leg calls, or ``None`` when this build has none.

Resolved by ``getattr`` on the already-imported extension module, NOT by a
``from exllamav3... import``. That is deliberate and it is the difference
between a seam and a breakage: a named import is a promise about the
pinned revision (``tests/test_exl3_fork_surface.py``), so adding one before
the kernel exists would fail the tree that declares it. A capability probe
on the loaded binary is the same shape ``_hgemm_acc_mode_supported`` uses
for ``hgemm``'s ``acc_mode``, and it is how this leg can be wired the day
the kernel is correct without a pin bump landing first.

THE FORK FIRST, THEN OURS. No shipped exllamav3 revision exports the symbol
(an exhaustive probe of the image finds 140 exports and no
``exl3_int8_gemm``), so what actually serves today is
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.int8_gemm_adapter` over the
kernel that ships in ``weight_quant/exl3/csrc``. The fork keeps precedence
anyway, because the day it exports one it will also own the A side, and a
build that had both must not serve two different arithmetics depending on
which branch this function happened to take.

Returns ``None`` when NEITHER exists -- no fork export and no built kernel
-- which is a real state (no CUDA, a build that failed, a boot that never
reached the pin seam) and is the one the census names ``build has no
exl3_int8_gemm``. It never triggers a build: the compile belongs to
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.load_int8_kernel` at boot,
and this is read on the serving path.

Whether the int8 leg serves THIS call, counting every refusal.

The class decision has already been made — by the step, at the forward
site — so nothing here reads a row count. What is left are the per-linear
preconditions, and each one is a named refusal:

``off-class``
    The forward declared DECODE, or declared VERIFY while
    ``exl3_int8_verify`` is ``0``. Expected, and the common reading.

``verify slate below min rows``
    The forward declared VERIFY, ``exl3_int8_verify`` is ``auto``, and the
    slate is narrower than ``exl3_int8_verify_min_rows``. THE RULE
    WORKING, not a fault: a boot that never saturates spends every verify
    call here, and it has its own name so that reading is
    "refused-by-rule" rather than the "armed and dead" a shared zero would
    show. The width comes from the DECLARED forward, so this refusal is
    the same one for every linear in the step and the same one a captured
    graph recorded for its own ``(B, S)``.
``windowed output``
    An ``EXL3OutputPrefixLinear`` emits fewer columns than its trellis
    stores; only ``exl3_gemm`` takes the emitted width as a bound. Same
    exclusion the trellis leg's own dispatch makes, for the same reason.
``codebook is exl3 default`` / ``codebook seeds disagree (mcg and mul1 both
set)``
    THE ONLY PRECONDITION HERE WHOSE VIOLATION WOULD BE SILENT. The kernel
    decodes ``mcg`` and ``mul1`` — the codebook is a template parameter of
    the mainloop — and nothing else. A trellis of one codebook is
    byte-for-byte indistinguishable from another's — same shape, same
    dtype — so decoding it with the wrong arithmetic would return
    plausible, wrong weights with no shape, no dtype and no counter
    noticing. Every other refusal on this list would have surfaced as an
    error; this one would have surfaced as degraded output. See
    :func:`~arbi_serve.weight_quant.exl3.int8_policy.int8_codebook_refusal`,
    which owns the rule and is pure over the two booleans.

``build has no exl3_int8_gemm``
    The flag is armed on a build without the kernel. LOUD once.
``int8 A-fragment scratch unreserved``
    The linear is pinned but :func:`reserve_int8_scratch` never ran, so the
    adapter has no preallocated buffer. It will not allocate one on the
    serving path: a prefill graph is captured, and an allocation inside a
    captured region bakes a stale pointer.

``<class> <band>: ...``
    The pin considered this linear FOR THIS CLASS AND BAND and declined it
    by name — a recorded table decline, the logits head's role, or a widest
    call site that cannot fill one m-tile of this band's tile. Per band,
    because the tile is an input to the rule: a geometry the leg serves at
    a 16-row tile can be declined at a 32-row one.

``int8 kernel-shape pin unresolved``
    THE INVARIANCE PRECONDITION, and the reason this refusal exists before
    the kernel does. The trellis leg is row-count-invariant only because
    :func:`resolve_kernel_shape_pins` freezes one ``(TILESIZE_K,
    TILESIZE_N)`` family per geometry and lets nothing but tile DEPTH
    follow the rows. A Marlin-shaped int8 kernel has the same exposure in
    its own vocabulary — its K-tiling and any split-K factor set the
    reduction ORDER, so a picker that reads M would make one row's output
    a function of how many rows travelled with it, inside the class the
    pin was supposed to have made safe. The leg therefore refuses an
    unpinned linear rather than serving it unpinned: an unpinned linear
    can never reach the new leg, so the pin cannot be forgotten and later
    discovered.

Declare the enclosed forward's PHASE for the large-M accumulator.

Restores the previous value on exit, including on exception, so a failed
prefill cannot leave the process computing decode with a prefill
accumulator — the direction that matters, since the KV decode writes is
what the accept-invariance proof rests on.

Declare that the enclosed forward's first ``rows`` flat rows are DECODE.

``flat_tokens`` is the forward's total flat token count, and the split
applies only to a call whose row count equals it — see
:data:`_MIXED_ROW_SPLIT`. Passing ``0`` (the default) disables the split
entirely rather than applying it unguarded: an undeclared total is not a
licence to guess a layout.

Restores the previous value on exit, including on exception: a failed
mixed capture must not leave the process splitting an ordinary prefill
forward at a stale row index, which would silently change that forward's
hgemm M and therefore its output.

Run leg B's ``hgemm`` at a FIXED M per numerics class.

The single seam every leg-B ``hgemm`` goes through, so the row-class
split cannot land on some call sites and not others — a partial
application would leave a linear whose decode rows still ride the
chunk width, which is the bug this exists to prevent and the kind that
hides for a month.

With no split declared this is exactly today's single call, same
arguments, byte-for-byte. With a split at ``d`` it is two calls over
disjoint row ranges of the SAME reconstructed ``w`` — so the weight is
materialised once (the expensive half; the reconstruct is what a second
forward would have to repeat) while each block's M stays a constant of
its class.

The split is IGNORED unless this call's row count is exactly the flat
token count the caller declared it against. That is not defensive coding:
a linear running at some OTHER row count is not carrying the layout the
split point describes — the lm_head's one-row-per-sequence gather is the
ordinary case, a permuted MoE expert range the dangerous one — and
blocking at an index that means nothing there would mis-assign rows to
classes. One class is the correct answer for such a call, and the mixed
sweep refuses outright the one composition (grouped MoE) where the whole
forward lacks a usable split.

Mark the enclosed forward as the PREFILL numerics class.

The engine marks the class only when ``exl3_prefill_row_invariant`` is set (see
``arbi_serve/runtime/forward_exec.run_model_forward``); the served default
leaves it unmarked, so the dispatch keeps its row-keyed leg choice and the
op body's gate short-circuits on a module-global read.

Restores the previous value on exit, including on exception, so a failed
forward cannot leave the process computing decode on the prefill leg.

Whether the PREFILL class pins this linear to leg A.

A linear crosses the leg boundary inside the prefill row range iff its
threshold falls in that range. The NARROW class (``out_features <=
NARROW_RECONSTRUCT_MAX_N``) has a threshold of 64, which every prefill
step above a tiny tail chunk clears -- so a short prompt alone (58 rows)
and the same prompt beside three others (237 rows) land on different
legs. On the served 27B that is exactly 34 linears: the ``k_proj`` and
``v_proj`` of all 17 full-attention layers, which feed the KV cache.

Static in the linear's own geometry, and expressed ONLY through the
threshold function the dispatch itself consults -- "is this linear on the
lowered threshold?" -- so no second copy of the narrow bound exists here
to drift away from upstream's.

Geometries whose widest verify slate leaves the leg its own decode takes.

``[(out_features, threshold, widest_slate), ...]``, empty when none cross.

THE DEFECT THIS NAMES. The leg is chosen from the row count alone
(``rows > auto_reconstruct_threshold(out_features)``), a verify step runs
``B * (n_draft + 1)`` rows where the plain decode of the same tokens runs
``B``, and ``B`` is bounded by ``max_batch``. So whenever
``max_batch * (n_draft + 1)`` exceeds a geometry's threshold while
``max_batch`` does not, that geometry's verify takes reconstruct + cuBLAS
``hgemm`` and its decode takes the trellis GEMM. Those are different
ALGORITHMS, not a different reduction order, and the kernel-shape pin does
not reach the reconstruct leg at all — so speculative decoding's
losslessness fails there and no amount of pinning repairs it.

It bites the NARROW class hardest and that is the worst place for it: the
threshold is 64 for ``out_features <= 2048``, which on the served 27B is
exactly the 34 ``k_proj``/``v_proj`` that feed the KV cache.

LATENT AT EVERY DEPLOYED CONFIGURATION, which is why this reports rather
than refuses. The dispatch is ``rows > threshold``, so a widest slate of
exactly 64 does NOT cross: ``max_batch = 8`` with ``n_draft = 7`` gives 64
and stays on leg A with its decode. It arms from ``max_batch = 9`` at
``n_draft = 7``, or ``max_batch = 16`` at ``n_draft = 5``. Refusing would
break a configuration an operator may legitimately want and whose cost is a
numerics guarantee rather than a crash; the fix is a fork change plus a pin
bump (the threshold is imported from exllamav3 and deliberately never
redeclared here — see the dispatch comment), which is real work that should
happen when it is needed.

A pure function of three numbers so the CPU lane can hold it, and so the
condition lives somewhere that FIRES rather than only in a comment. A
condition recorded in an issue is a condition nobody checks.

``threshold`` defaults to :func:`_auto_reconstruct_threshold` — the same
function the dispatch consults, never a second copy of the bound.

Install the row set this boot races the kernel families over. Returns it.

``None`` clears the installation, which is the state a fresh process and a
torn-down test are both in. Called by the engine's pin seam, which is the
only caller holding both the flag and the served config that ``auto``
derives from.

BOOT SCOPE, and not by convention: the pin is resolved before the first
forward and its shapes are baked into every captured graph, so installing a
different set after that changes nothing already pinned and would only make
the next seam disagree with the graphs.

The row counts the family race is timed at in this process.

Resolution order: an installed set, then an explicit
``ARBI_EXL3_PIN_SELECT_ROWS`` list, then :data:`_PIN_SELECT_ROWS`. The flag
is read here as well as at the seam so a caller that resolves pins outside
the engine — a recorder, a test — still honours a row set the operator
named, rather than recording one key's answer under another key.

``auto`` cannot be answered here: it is a function of the served config,
which this module never sees. It reads as the shipped tuple and says so
once, because the alternative is a boot silently racing over rows nobody
chose.

The kernel-shape family every EXL3 geometry resolved to.

``{"<in>x<out>xK<bits>": (family, member_shape_indices, {family: ms})}``.
The family key is ``"K<tilesize_k>N<tilesize_n>"`` and the scores are the
selection timings, so both the choice and the margin it won by survive the
boot that made them.

The scores of a RESTORED geometry are the timings of the boot that measured
it, not of this one. :func:`kernel_shape_pin_origins` says which is which.

``{geometry: force_num_sms}`` this process pinned, ``0`` = kernel's choice.

Beside :func:`kernel_shape_pins` rather than inside it: the family and the
width are both frozen inputs to the same launch, but the family's receipt
carries the race that chose it and the width's does not — it is restored
from a recorded decision or it is absent. Merging them would put a measured
margin next to a value nothing on this boot measured.

``{geometry: "measured" | "cached"}`` for this process's resolved pins.

A geometry absent from this map was never pinned at all — the decline paths
in :func:`resolve_kernel_shape_pins` say why in the log.

Whether ONE inner carries a resolved kernel-shape pin.

The single reader of the per-inner pin attribute outside the dispatch, so
"is this linear pinned" has one answer everywhere it is asked.

``(pinned, bound)`` over ``inners``, or over every registered inner.

COVERAGE, not the delta one resolution pass added. The two differ at every
seam that runs after the first: :func:`resolve_kernel_shape_pins` is
idempotent, so a second seam with nothing left to pin returns 0 — which
reads exactly like a pin that could not be established, and was reported as
one. Whether the trellis GEMM is row-count invariant is a question about
the linears that are BOUND, and this is the function that answers it.

``{shape_idx: (TILESIZE_M, K, N)}`` from the pinned exllamav3 header.

Read through :mod:`arbi_serve.weight_quant.exl3.kernel_shape_table` rather
than restated here, so a pin move cannot leave a stale table behind. That
module SHIPS, and must: reading the table through the repo-only ``tools``
tree makes this pin inert in every built image (see its docstring). An
unreadable header returns ``{}``, which makes the pin decline rather than
guess.

``{(TILESIZE_K, TILESIZE_N): ladder}`` this geometry may be pinned to.

The candidate set the selection chooses from, and the same value
:func:`~arbi_serve.weight_quant.exl3.shape_pin_cache.restorable_ladder`
checks a restored pin against. It lives here as ONE definition because
three callers need it to agree exactly: the sweep, the restore check, and
any instrument that wants to hold a family other than the one the stopwatch
drew. A ladder derived a second time somewhere else and drifting by one
member is a pin that is silently refused — or worse, silently accepted for
a family nobody measured.

A family qualifies only if it has at least one legal member AND its
shallowest member covers a single row: the pin serves decode, and a ladder
that starts above one row cannot.

Shallowest member deep enough for ``rows``; above the deepest, the member
that costs the fewest full weight passes.

``exl3_gemm_shape_compat`` declines a tile deeper than 16 rows below 17, so
the ladder is walked shallow-first and the entry is pre-filtered to members
legal for this geometry.

Above the deepest member the shallow-first walk falls off the end, and
taking the deepest there is NOT free: the kernel tiles M serially, so the
cost is ``ceil(rows / tile_m)`` full passes over the trellis, and that count
is not monotonic in ``tile_m``. At 208 rows a 96-row tile needs 3 passes and
so does an 80-row tile -- but the 80-row tile wastes 32 fewer row-slots
doing it. Always taking the deepest costs 1.25-1.28x at 112-128 rows and
1.066x census over the served range (measured, 27B on sm_89).

So above the deepest member: minimise the pass count first, and among the
members that achieve it take the SHALLOWEST -- least partial-tile waste, and
lower register/SMEM pressure. That rule reproduces every measured band
(T64 at 112-128, T80 at 144-160 and 208-240, T96 elsewhere) from the cost
model alone, rather than pinning a fitted table that would go stale the next
time the shape table gains a member.

The trellis-leg launch a served boot makes for one geometry at ``rows``.

Every field is an input to the ``exl3_gemm`` call the dispatch above
issues: ``shape_idx`` is its ``force_shape_idx`` and ``num_sms`` its
``force_num_sms``. ``family`` and ``ladder`` are the record they were
resolved from, kept so an instrument can print WHICH pin it timed rather
than a bare index.

No shipped pin covers this geometry on this configuration.

Raised by :func:`served_trellis_pin` instead of returning the unpinned
dispatch (``force_shape_idx = -1``): that dispatch is the GEMV special case
below nine rows and the autotuner above it, free to draw a member of any
family, and no boot serves it. A measurement taken on it is a measurement
of a kernel the deployment does not run. The message names the geometry
and the recorder that produces the missing entry.

The pin a served boot passes to ``exl3_gemm`` for this geometry at ``rows``.

ONE RESOLUTION, SHARED WITH THE BOOT. The shipped table is read through
:func:`~arbi_serve.weight_quant.exl3.shape_pin_cache.shipped_pin_entries`
under the key the boot computes, the family's ladder is the one
:func:`legal_shape_ladders` derives, and the member is what
:func:`_pin_for_rows` walks to -- the three calls
:func:`resolve_kernel_shape_pins` makes, in the same order. An offline
instrument that wants "the leg the engine serves" calls this rather than
restating any of them, so it cannot time a member the boot never binds.

A GEOMETRY WITHOUT A RECORD IS REFUSED, never defaulted. The boot falls
through to its own timed race for such a geometry, and that race is a
function of the load the boot ran under, so no offline instrument can
reproduce its answer; the only leg both agree on is a recorded one.
:class:`UnservedGeometryError` names the geometry and the recorder.

``ext``, ``dims`` and ``shipped`` are injection seams so the resolution is
checkable without a card; a caller that leaves them ``None`` gets the
served extension, the pinned shape table and the shipped table for
``device``.

``{geometry: {band: shape_id}}`` the int8 VERIFY class serves this boot.

PER BAND, because that is what the deployment serves: a captured verify
graph is recorded per ``(B, S)``, so a served slate width is one of a small
known set, and each width takes the tile that wins at it. A reading that
collapsed the bands would name one tile for calls that run two.

``{geometry: {band: reason}}`` the VERIFY class will not serve.

Two-sided as above, and per band for the same reason the pins are: the row
bound a linear's call sites declare is compared against THIS band's TILE_M,
so a geometry can be served in one band and declined in the next.

``{band label: verify forwards declared}`` — the per-band census.

The counter half ``exl3_int8_verify_gemm`` cannot carry: it has one fire
count, and which TILE that count was taken on is what an operator comparing
two boots (or two threshold settings) is actually asking. Counted at the
declaration seam, once per forward, because the band is a property of the
slate and not of a linear.

``(threshold, origin)`` DECODE routes on under ``exl3_int8_decode=auto``.

``None`` means unresolved and every decode call refuses by name. The
origin is the sentence a boot log and a receipt print beside the number:
which ladders it was derived from, or which flag pinned it.

``(trellis ms, int8 ms | None, trellis launches, int8 launches)`` for one cell.

The trellis arm launches exactly what the op launches for this linear at
``rows`` -- the pinned member ``_pin_for_rows`` walks to and the pinned grid
width; the int8 arm is the served chain ``int8_gemm_entry`` resolves to,
on the band plan the declaration seam would select. ``None`` for int8
when this linear has no armed plan for that band: the leg would refuse it
by name and the trellis member would serve, so the whole-model sum charges
the trellis time to both arms for it.

The launch COUNTS are returned so the caller can prove the arms ran --
a probe whose kernels never launched reports a plausible number, and a
plausible number from a check that could not run is the defect this
codebase names most often.

``(threshold, why)`` from ``{rows: (trellis ms, int8 ms)}`` whole-model sums. Pure.

THE NARROWEST SERVED WIDTH FROM WHICH THE INT8 CHAIN WINS AT EVERY WIDER
ONE. Monotone by requirement, not by assumption: a set on which the winner
alternates has no single threshold, and a threshold that admitted a losing
rung above it would be a regression the rule itself introduced -- so that
set resolves ``None`` and the reason names the rung.

Install the DECODE threshold for this boot. Returns ``(threshold, origin)``.

THE VALUE IS THE FLAG'S OR THIS BOOT'S OWN MEASUREMENT, NEVER A LITERAL.
An operator that typed an integer into ``exl3_int8_decode_min_rows``
pinned it and nothing is probed. ``auto`` MEASURES it: for one linear per
geometry this member bound, at every row count in ``rows_set`` (the
captured decode B ladder the engine will actually present), the served
trellis member is timed against the int8 band plan the declaration seam
would select, the two are summed over the checkpoint's geometries weighted
by how many linears carry each, and
:func:`int8_decode_min_rows_from_probe` takes the narrowest width from
which the int8 chain wins monotonically. The receipt that motivated this
(``docs/receipts/int8-decode-crossover-2026-09-07.txt``) found the edge at
the served pin's deepest tile + 1 on a 4090, and that derivation
(:func:`~arbi_serve.weight_quant.exl3.int8_policy.int8_decode_min_rows_from_tiles`)
is printed in the origin as the CROSS-CHECK beside the probed value; the
router reads the probed value.

CACHED under the budget cache, keyed on everything the answer depends on:
``key_inputs`` (the trellis pin's own key: card, driver, torch, exllamav3
and arbi-serve revisions, shape table, select rows, TP size) plus the
served row set, the geometry census, every int8 plan's shape per band, the
quantiser, and the probe version. Same key, same answer, no stopwatch --
which is what keeps the leg a decode step takes from varying boot to boot
at one configuration; a changed input is a new key and a new measurement.

A boot that cannot probe -- no row set, no pinned trellis member, no armed
int8 plan anywhere -- resolves ``None`` and says why; the rule then refuses
every decode call by name rather than routing on a width nothing measured.
Called at the int8 arm seam AFTER the scratch is reserved (the int8 arm
needs armed plans) and after the trellis pin (the trellis arm needs the
member), which is the order ``engine/exl3_shape_pin.py`` runs them.

``{geometry: reason}`` for every geometry the int8 leg will NOT serve.

The two-sided half of the receipt. A boot that pinned six geometries and
declined three has said something; a boot that reports six and nothing else
has not, and the six read identically in both.

``(pinned, bound)`` int8 launch plans over the whole registry.

COVERAGE, not the delta one pass added, for the reason
:func:`kernel_shape_pin_state` is: the resolution is idempotent, so a second
seam with nothing left to pin returns 0 and that reads exactly like a pin
that could not be established.

``verify`` selects the CLASS. A keyword over a class token because there
are two and the caller is a log line; the plan attributes stay private to
this module, so nothing outside it can address a linear's plan by a name it
spelled itself.

A VERIFY linear counts as pinned when ANY band pinned it. The bands are
tiles, not capabilities: a linear the leg serves at one slate width and
declines at another is on the leg, and the per-band receipts say which
widths. Requiring every band would report a whole class dead over one
geometry's narrow-tile decline.

Install the VERIFY tile ladder and derive every table that follows it.

ONE WRITER for the band count. The plan attributes, the decline
attributes, the labels, the receipt maps and the selector the declaration
seam walks are all functions of this ladder, so deriving them here is what
stops a boot from holding a plan for a band nothing can select — or a
selector for a band nothing pinned.

A ladder that REPLACES a different one strips the old attributes off every
registered linear first. A stale plan left on a linear is armed at scratch
the new reservation does not size and no seam disarms, and the GEMM writes
through it.

Install the PACKAGE-DEFAULT ladder if the pin has not resolved one yet.

The seams that walk a linear's plans — scratch sizing, the rebind, the test
teardown — are reachable on a boot that never armed the leg, and they must
address the same attributes the pin would have used. The band EDGES do not
depend on the quantiser, so the tables this installs are the ones the pin
then confirms or replaces.

Record a geometry's decline for ONE CLASS, on the receipt AND on the linear.

Both, because two different readers ask. An operator reads the receipt to
see what this deployment serves; the census reads the linear, on the
forward, to name the refusal — and a refusal that said "pin unresolved" for
a geometry the pin had deliberately declined would report a considered
decision as a missing one.

Per class, because the classes decline for different reasons: the tile a
class is pinned to is an input to the rule. A shared reason would tell a
verify refusal the prefill story.

The recorded table for this configuration. ``{}`` when it cannot be read.

A thin wrapper over the trellis pin's own reader so both ladders take one
key and one file. An unreadable table costs the DEFAULT shape and nothing
else, which is what every geometry the table does not cover takes anyway.

The trellis's own bit width, or 0 when it cannot be read.

``B.size(2)`` is ``16 * K``; the kernel derives K the same way and refuses a
last axis that is not a multiple of 16, so reading it here asks the tensor
rather than the descriptor.

The VERIFY tile ladder this boot serves. Never raises.

THREE SOURCES, IN ORDER, and each is a decision somebody recorded rather
than a number this boot invented: the caller's own ``verify_bands`` (the
seam a test and a bench arm drive), the shipped table's entry-wide
``int8_verify_bands``, and the package default -- which is the ladder
``docs/receipts/int8-verify-small-m-2026-09-05.txt`` measured.

A recorded ladder that cannot be READ, or that does not pass
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.int8_verify_bands_refusal`,
falls back to the default and says so. That is the fail-safe direction and
the only one available: the refusal is about the ARITHMETIC being
band-dependent, so honouring the record would serve a verify slate that
does not reproduce the decode it verifies.

Freeze ONE int8 launch shape per EXL3 geometry PER STEP CLASS AND BAND.

Returns the number of launch plans newly pinned, across every class and
band -- not linears, because a linear can be pinned in one and declined in
another, and a count of linears could not say which.

A DECISION, NEVER A RACE, and the asymmetry with the trellis family beside
it is the point. ``exl3_gemm`` has no default shape -- some member of some
family must be forced, so when nothing is recorded a boot-time stopwatch is
the only instrument available. The int8 kernel has one per class
(:data:`~arbi_serve.weight_quant.exl3.int8_kernel.INT8_DEFAULT_SHAPE_ID`
and its verify sibling, the shapes ``tools/int8_gemm``'s receipts were
measured against), so racing either would INTRODUCE a boot-to-boot variable
where none exists -- #1879's defect with the sign flipped, and #1910
recorded exactly this reasoning for the grid width. A measured improvement
lands in ``shipped_pins.json`` as ``int8_shape`` / ``int8_verify_shape``;
no boot ever invents one.

ROW-COUNT INVARIANCE IS STRUCTURAL, which is why each class is one frozen
shape and not a ladder. A shape id fixes TILE_K, the stage count and the k
loop, so the reduction order over ``in_features`` is identical at every row
count; M only decides how many blocks tile the output.

THE CLASSES TAKE DIFFERENT TILES, AND VERIFY TAKES ONE PER ROW BAND, and
both are admissible for the same reason: every entry shares a
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.int8_reduction_key`, so
TILE_M moves the block count and nothing else. What the class buys is that
a verify slate -- a fraction of one 128-row m-tile at every served
concurrency -- stops launching a prefill chunk's grid. What the BAND buys
is the rest: the int8 tile pads M up to a whole number of TILE_M and the
trellis leg steps at its own boundaries, so the two legs are sawtooth in M
with their teeth out of phase and the winner alternates across the slate
span. A captured deployment presents a small known set of slate widths, so
each width can take the tile that wins at it rather than one tile carrying
the whole span.

THE LADDER IS RECORDED, never derived from this boot: ``verify_bands`` (or
the shipped table's entry-wide ``int8_verify_bands``) names the edges, the
package default is the receipt's, and a ladder whose members do not all
share the prefill tile's reduction key is REFUSED rather than served --
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.int8_verify_bands_refusal`
owns that rule.

Idempotent, and silent about geometries it cannot serve except through
:func:`int8_shape_pins_declined` and its verify sibling. A linear it
declines keeps the shipped legs, which is the pre-arm behaviour. The one
thing it refuses is an ARMED leg with no kernel on a CUDA device: that
boot would report a configuration it does not serve.

``"<band> (<tile>): <geom>=shape<id>, ...; ..."`` — the per-band pin line.

ONE CLAUSE PER BAND, because that is the decision: an operator reading this
against a boot's captured ``(B, S)`` ladder can say which tile each served
slate width takes. A line that summed the bands would name a tile for
widths that do not run it.

The VERIFY declines, per band, for the reason the pins are per band.

Two-sided like every other decline line — a class that declined nothing and
a class whose declines were dropped read identically otherwise — but ONE
sentence when no band declined anything, because a per-band repetition of
"nothing" is noise an operator has to read past to reach the bands that did.

``"<label>: <geom> (<why>), ..."``, or a sentence saying nothing was.

Two-sided by construction, for the reason the pin's own log is: a boot that
reports what it pinned and says nothing about what it did not reads
identically to one that had nothing to decline.

The widest row count any call site can hand this linear.

THE SAME BOUND leg B's reconstruct scratch is sized from
(:func:`declare_reconstruct_row_bounds`), because it is the same question --
the logits head gathers one row per sequence where a body linear sees a
whole prefill chunk. Deriving a second bound here is how the two would drift
and how one of them would be wrong.

``(int8 A-fragment elements, fp32 scale elements)`` the widest pinned linear needs.

Over ``member``'s linears that carry an int8 plan and no scratch yet — one
member's worth, for the reason :func:`reconstruct_scratch_elems` is scoped
that way, and now by construction rather than by the accident that no other
member had built yet (:func:`member_inners`).
``(0, 0)`` when nothing is waiting, which is both "the leg is not armed" and
"already reserved".

A pinned linear with NO declared row bound contributes nothing and is
reported by :func:`reserve_int8_scratch`: sizing it would mean inventing a
bound, and a scratch too small for a call site is a refusal at serve time
rather than a wrong answer, so guessing buys nothing.

Reserve (or grow, pre-capture) THIS member's int8 A-fragment scratch.

Returns ``((a_frag, a_scale), bytes)``. The caller owns the pool routing and
holds the buffers for the member's lifetime, exactly as it does for the
reconstruct scratch — same pool (``capture.io_buffers``), same reason: a
captured prefill graph bakes these ``data_ptr``s and the pool's lifetime is
the one that matches.

Runs for the ACTIVE member: at its build, and again whenever it becomes
active (:meth:`Engine.register_model_switch_hook`). The buffers live in
that member's ``capture.io_buffers`` pool because its captured graphs bake
their ``data_ptr``, so they are the member's to own — and a member swap
releases the outgoing member's pool. A binding therefore outlives the
memory it names as soon as the engine leaves that member, and the plans
are process-global objects shared by every member, so the swap BACK finds
them pointing into a pool that is gone. The GEMM writes through it: an
illegal access with the weights, the activations and the allocator all
healthy.

So the binding is re-established rather than trusted, and the plans are
DISARMED before the new buffers are bound. Unarmed is the safe state — the
linear refuses by name — and disarming first is what makes "armed at a
released buffer" unrepresentable rather than merely unlikely: a failure
part-way through leaves plans refusing, not faulting.

There is no serving-path fallback. A plan this call does not reach stays
disarmed and its linear refuses by name.

``member`` scopes the walk to the linears that member's build registered
(:func:`member_inners`). The registry is process-global and never pruned —
it has to be, because a parked member's captured graphs resolve their
linears out of it by uid on the wake — so an unscoped walk disarms and
re-binds EVERY member's plans onto the buffers of whichever member happens
to be reserving. Within a member a grow must move every plan; across
members it must move none. ``None`` is the un-namespaced member, so a boot
with no residency reaches every linear exactly as before.

Choose one kernel-shape FAMILY per EXL3 geometry and freeze it.

Returns the number of linears newly pinned. Boot-time by contract: the
selection times candidate kernels, which must not happen inside a cudagraph
capture, and a linear reaching serving unpinned runs a different kernel from
its already-pinned siblings.

MEASURED ONCE PER CONFIGURATION. The winner is a property of the card and
the kernel set, so the choice is read back from
:mod:`arbi_serve.weight_quant.exl3.shape_pin_cache` when that module holds a
reading for this exact configuration, and TIMED — and then recorded — when
it does not. What that buys is not speed: it is that two boots at one
configuration draw the SAME kernel, which is the precondition for any
cross-boot comparison to mean anything. A restored ladder is checked against
the one this build derives and REFUSED if they differ; nothing is repaired.

Idempotent — a linear already carrying a pin is skipped, so a second seam
only pins what a drafter added since. NEVER raises: a linear that cannot be
pinned keeps the row-count-dependent auto dispatch, which is the pre-pin
behaviour, and says so in the log.

Args:
    device: the card the selection is timed on.
    build_dir: override the pin-cache directory. ``None`` (default) takes
        the pin cache's own resolution — ``ARBI_SERVE_EXL3_PIN_CACHE_DIR``,
        then the user cache fallback. Deliberately independent of the
        budget cache's directory, which A/B arms move to isolate the budget
        measurement and would otherwise re-sweep the pin along with it.

``(cache_key_inputs, {geometry: record})`` for this configuration.

A cache that cannot be consulted must cost the sweep and nothing else, so
every failure here degrades to ``(None, {})``: the geometries are timed, and
``None`` inputs later suppress the write rather than persisting under a key
that was not fully built.

``{geometry: {class: shape id}}`` the int8 leg resolved this boot.

The shape the pin cache records beside the trellis family. Assembled from
the per-class receipt maps rather than kept as a third one, so a geometry
that is pinned in one class and declined in the other says exactly that
(one key present) instead of carrying a sentinel somebody has to interpret.

Write the resolved pin back, and say which half of it this boot measured.

Driven by the SWEEP COMPLETING and by nothing else — not a ``/metrics``
scrape, not an OTEL callback, not an admin route. The engine owns the
reading; every telemetry surface reads it back through
:func:`kernel_shape_pins` / :func:`kernel_shape_pin_origins`.

Writes when a geometry was actually timed, OR when the int8 leg's per-class
tiles differ from what the file already holds. A boot that restored every
pin and resolved the same tiles has nothing new to say, and rewriting the
file from it would put a fresh timestamp on a reading it did not take — but
a boot that armed the leg for the first time, or that moved a class's tile,
has changed what this deployment serves, and a record that omitted it would
describe a configuration nobody is running.

``"<geom>=<family>, ..."`` for a log line, sorted by geometry.

The geometry alone says WHICH readings a boot took; the family says WHAT
IT IS RUNNING, and that is the variable underneath every cross-boot
comparison. An arm that silently drew the other family is the failure this
line exists to make impossible to miss, so the two are never printed apart.

Pick the family, and say whether the STOPWATCH actually picked it.

``(family, how, margin_pct)``. ``how`` is ``"timed"`` when the winner beat
the runner-up by more than either reading could resolve, and
``"tie-break"`` when it did not.

WHY A TIE-BREAK AT ALL. The race is not noisy -- measured 12/12 identical
per geometry on a quiet card (#1879) -- and that is precisely the problem.
Run it again under one competing process and two of seven geometries pick a
DIFFERENT family, also 12/12. The answer is a deterministic function of
load, so it is stable within a deployment and inconsistent across
deployments, and the cache then makes the first draw permanent. A coin flip
would have been easier: it would have been visible.

So when the gap is inside what the measurement can resolve, the stopwatch
did not decide it and is not allowed to. The tie breaks on the family key
itself -- a total order on the candidate set that depends on the shape table
and nothing else -- so two deployments on the same table agree whatever
their load was. The cost of choosing wrongly inside the band is bounded BY
the band, and the band is small change: an oracle picking the best family
per geometry per row -- better than any pin can be, since one family must
serve the whole span -- beats the shipped pin by 1.011-1.023x at the rows
actually served (#1859). This trades a fraction of a percent of GEMM time
for two deployments that serve the same numerics.

Not "prefer the incumbent", which would be the other obvious rule: there is
no incumbent at the only moment this runs, which is a first sweep at a
configuration nothing has recorded.

``(median wall time, this reading's own resolution)`` for one family, ms.

Summed rather than compared per row: the pin serves the whole span with one
family, so the quantity being minimised is the span's total.

THE SECOND NUMBER IS WHY THIS RETURNS A PAIR. The winner of this race is
frequently decided by less than the race can resolve — measured across the
served 27B's geometries, six of eight by under 3.5% and the ``k_proj`` /
``v_proj`` pair by 0.14% (#1879). A margin that small is not a decision, and
treating it as one is what makes the choice a deterministic function of how
busy the card was: quiet and contended each pick stably, and they pick
DIFFERENTLY. So the caller needs to know what this reading could actually
separate, and it is measured here rather than assumed as a constant,
because it is a property of the card, the geometry and the load of this
boot -- every one of which a constant would be wrong about.

The resolution is the summed inter-quartile range of the per-row samples:
the same statistic, over the same samples, as the median it accompanies.

Drop every resolved pin (test cleanup / engine teardown).

The installed row set goes with them: it is an input to the resolution, not
a property of the process, and a set left behind would race the NEXT
engine's geometries over the previous one's rows.

The cuMem tag namespace currently open, or ``None``.

One read of the allocator, which owns the namespace — never a copy of it.
``None`` on every path with no residency build in progress, which is what
makes an un-namespaced boot indistinguishable from today.

The registered inners built under stable-VA member ``member``.

The gate every reservation walk goes through. A linear belongs to exactly
one member — the one whose build registered it — and a member that is
parked must not have its plans re-pointed at the ACTIVE member's buffers:
within a member a scratch grow has to move every plan, across members it
must never move any. ``None`` selects the un-namespaced member, so a boot
with no residency reaches every linear exactly as it did before the stamp
existed.

Name ``member`` as the owner of every inner that carries no stamp yet.

The boot model is built BEFORE any residency namespace exists, so its
linears register unstamped; every later member builds inside
``tag_namespace(key)`` and stamps itself. Called once, when the boot model
is registered as a residency record, so the boot member is addressable by
the same key its record carries and a wake can gate on one expression for
every member. Returns the number adopted.

Publish a per-linear :class:`exllamav3.LinearEXL3` against ``uid``.

Stamps the inner with the stable-VA member being built (:data:`_MEMBER_TAG_ATTR`)
so a later scratch reservation can tell this member's linears from a parked
member's — see :func:`member_inners`. An inner that refuses attributes (a
bare stand-in in a registry test) stays unstamped, which reads as the
un-namespaced member — the pre-stamp behaviour — rather than failing a bind.

Drop the inner published against ``uid`` (idempotent).

The registry holds a STRONG reference to the inner ``LinearEXL3``,
which in turn pins its ``trellis / suh / svh`` GPU storage. When the
compactor's release hook tears down ``EXL3LinearBase._inner`` ahead of
a slab allocation it must ALSO drop this registry entry, otherwise the
inner (and its original The rebind hook rebuilds
and re-registers the inner over the slab views afterwards.

Every ``_exl3_uid`` reachable from ``model`` with NO inner published.

A linear whose uid is missing from the dispatch table cannot forward: the
op's real-impl resolves by uid and raises. The gap is what a compaction
RELEASE hook leaves behind when the matching ``rebind_after_compaction``
never ran — an aborted member build is the way that happens on a model that
is otherwise fully loaded, and on a donor-share build the released modules
can be the DONOR's own (they are shared by object identity), so the member
that is supposedly still serving is the one left unable to forward.

Read as a SERVABILITY probe, so a caller reporting "this member continues
to serve" can check the claim instead of inferring it from residency
bookkeeping. Empty list = nothing missing (the ordinary state).

An UNBOUND linear is not a gap: a placeholder that never took weights owes
no descriptor (``rebind_after_compaction`` declines it on the same test),
and counting it would accuse a healthy member.

Drop every registered inner and the per-inner state hung off them
(test cleanup).

Detach BEFORE clearing: both clears walk the registry to reach the inners,
so emptying it first leaves every inner still carrying its scratch and its
shape pin.

Run the EXL3 trellis-dequant GEMM for one linear.

``x`` is fp16 with last dim == in_features. Returns fp16 with the
same leading dims and last dim == ``out_features``. ``layer_uid``
(a 0-dim int64 CPU tensor — see ``EXL3LinearBase.__init__`` for why
it's a tensor not an int) resolves the per-linear
:class:`exllamav3.LinearEXL3` (which owns the precomputed
``BC_LinearEXL3`` kernel descriptor + scratch).

Allocation discipline: the output is allocated ONCE here per call
shape, then the no-alloc ``bc.run`` writes into it. Under
torch.compile + cudagraph, Inductor plans this output buffer from
the registered fake and reuses it across replays — no per-step
device alloc inside the captured region.

The Hadamard basis an EXL3 checkpoint is quantized in, and served in.

EXL3 wraps every quantized linear in two Hadamard rotations -- one over the
INPUT (contraction) axis, one over the OUTPUT axis -- each normalised by
``1/sqrt(width)``. The width is not a tuning knob. It is part of the format:
the quantizer rotated the weights in some basis, and a kernel that rotates the
activations in a DIFFERENT one contracts two unrelated vector spaces. The
shapes still match, so nothing raises; the output is noise wearing the right
dimensions.

That is why the width lives here rather than as a literal beside each use. The
two axes are independent -- widening the input basis says nothing about the
output basis -- so a caller that means one of them says which, and reads the
scale from :func:`had_r_scale` for the width it named. A bare ``1/sqrt(128)``
at a call site records neither.

The device side of this constant is ``csrc/exl3_had128.cuh``
(``HAD_BASIS_WIDTH`` / ``had_r_scale``), and ``tests/test_exl3_had_basis.py``
reads that header and fails if the two drift. One number in two languages is
the seam, so it is checked rather than maintained by inspection.

IMPORT-LIGHT: nothing here imports torch or exllamav3, so a loader can ask
"what basis is this checkpoint in" without pulling the EXL3 forward path in.

A checkpoint's Hadamard basis is not the one this build rotates in.

Its own class because it is not a shape error and must not be caught by a
handler that recovers from one. Rotating in the wrong basis produces
correctly shaped, wholly wrong activations -- there is no fallback and no
degraded mode, only a refusal.

``1/sqrt(width)``, the normalisation of a ``width``-wide Hadamard.

Refuses a non-power-of-two width rather than returning a plausible number:
the Hadamard this normalises is a Sylvester construction, which exists only
at powers of two, so a caller holding some other width has already lost
track of what it is scaling.

``csrc/exl3_had128.cuh`` computes the same value without ``sqrt``, because
C++ has no ``constexpr`` one: it halves ``1/sqrt(2)`` down from the width,
which is exact because each step is an exponent decrement. Both spellings
land on the same fp32 value at every power of two, which is what the
kernels consume, and ``tests/test_exl3_had_basis.py`` asserts it rather
than asserting the two derivations look alike.

Refuse a Hadamard basis this build cannot serve. Returns the width.

``None`` means the source predates the field. It reads as
:data:`HAD_BASIS_WIDTH` because that is the only basis anything has ever
written -- an absent tag is not evidence of a different basis. A tag that
is PRESENT and different is refused, which is the case that matters: it is
what makes widening the basis a supported operation instead of a silent
relabelling of the weights.

``what`` names the source in the message (a checkpoint path, a linear), so
the refusal says which artefact disagrees rather than that one did.

The int8 EXL3 prefill GEMM as the served leg calls it: build it, then bridge it.

TWO HALVES OF ONE SEAM, and neither is useful without the other.

**The build.** ``csrc/exl3_i8_gemm.cu`` is the kernel #1861 measured. It was a
``tools/int8_gemm`` dev harness JIT-built by ``tools/int8_gemm/ext.py``, which
means it could not serve anything: the Dockerfile copies ``arbi_serve``
wholesale and does not carry ``tools``, so every built image had the harness's
source nowhere on disk. The source, the codebook ids and mul1's grid constant
therefore moved into the package -- one copy, imported by the harness rather
than duplicated for it.

**The bridge.** The served call site
(:func:`~arbi_serve.weight_quant.exl3.custom_op.exl3_linear_gemm`) hands nine
arguments starting with a RAW fp16 ``x``, with ``suh``/``svh`` still to be
applied. The kernel's own binding takes pre-rotated, pre-quantised int8 A
FRAGMENTS, a per-token scale, and a codebook id it refuses to assume. Bridging
those is what :func:`int8_gemm_adapter` is: one ``a_prep`` pass (suh pre-scale,
128-wide Hadamard, per-token amax, int8 quantise, SIGMA relabel, fragment-order
store) and one kernel call whose epilogue folds ``svh`` back in.

THE FRAGMENT ORDER IS THE WHOLE RISK. The mainloop reads
``A_frag[m/16][k'/32][lane][16]`` in ``mma.m16n8k32.s8`` operand order, not row
major, and the two layouts have the SAME shape and dtype -- a row-major
hand-off compiles, runs and is quietly wrong. The adapter therefore calls
:func:`~arbi_serve.weight_quant.exl3.a_prep.a_prep_frag`, whose store IS that
layout (:func:`~arbi_serve.weight_quant.exl3.a_prep.frag_offset` restates it),
and never :func:`~arbi_serve.weight_quant.exl3.a_prep.a_prep`, whose row-major
output has the identical type and size.

WHAT THE ADAPTER MAY NOT DO, AND WHY IT IS THE HARD PART. It sits on the
critical path of every prefill linear -- ~409 of them per forward on the served
27B -- and a prefill graph is CAPTURED (``prefill_capture=full``, #1885). So:

* **No host synchronisation.** ``a_scale`` is a device-side per-token amax and
  never leaves the device; nothing here reads a tensor VALUE. One ``.item()``
  would be 409 syncs per forward, which costs more than the kernel saves.
  ``tests/test_exl3_int8_adapter.py`` runs the adapter under
  ``torch.cuda.set_sync_debug_mode("error")``, so a sync introduced here fails
  a test rather than a benchmark.
* **No allocation.** The A-fragment buffer and the scale table are RESERVED AT
  BOOT (:func:`~arbi_serve.weight_quant.exl3.custom_op.reserve_int8_scratch`,
  the same treatment leg B's dequantised weight already gets) and reached
  through the launch plan. An allocation inside a captured region freezes at a
  stale pointer.
* **No lazy init.** Both extensions are compiled by :func:`load_int8_kernel` at
  the pin seam. A first call that ran ``ninja`` inside a capture is not
  survivable, so a plan with no scratch REFUSES rather than falling back --
  a fallback that works eagerly and breaks under capture is precisely the
  "inert while reporting as applied" defect #1885 found twice.
* **No data-dependent control flow.** Every branch here is over a Python int
  or a bool resolved at bind time and frozen into the plan.

WHY THE ADAPTER LIVES ON THE ARBI SIDE rather than being one more export of the
exllamav3 fork, which is where ``exl3_gemm`` and ``exl3_mgemm`` live. Its input
half is ``a_prep``, and ``a_prep`` is arbi's: it fuses the prologue (the norm,
the ``silu*mul``, the residual add) that the fork has no notion of. A fork-side
``exl3_int8_gemm`` would have to carry ``a_prep`` with it, which would put the
prologue fusion behind a fork revision bump.
:func:`~arbi_serve.weight_quant.exl3.custom_op.int8_gemm_entry` still prefers a
fork export when one exists, so this is a fallback rather than a fork.

The shape a geometry takes when the shipped table names none.

One seam, so the flag and the pin cannot disagree about which entry the
quantiser choice means -- and so a test can drive both values without a
boot.

The VERIFY tile ladder: ``((widest rows this band serves, shape), ...)``.

Narrowest band first, and the LAST band's row bound is ``None`` — it is the
band every slate above the last edge lands in, so it cannot be a number
without inventing a widest served slate.

A LADDER OVER ONE REDUCTION KEY, which is the whole licence for it:
:func:`int8_verify_bands_refusal` is the machine check that every member
contracts ``k`` the way the prefill tile does, so the band decides which
blocks do the work and never what the work computes.

Which band a slate of ``rows`` rows lands in. Pure over two ints.

Read ONCE per forward, at the declaration seam, from a width the forward
already declared — never per linear and never off an activation. A linear
that read its own ``x`` would make one row's tile follow the geometry it
landed on, and under cudagraph capture there is no host read at all.

The band's name in a receipt, a refusal and a log line.

The BOUND, not the index: an operator reading "verify rows<=48" against a
slate width knows immediately which band a call took, where "band 0" would
have to be resolved against a table that is not in front of them.

Why this ladder must not be served, or ``None``. Pure over the table.

THE LADDER IS ONLY LEGAL IF THE TILE IS NOT PART OF THE ARITHMETIC. Every
member has to share the prefill tile's :func:`int8_reduction_key`, because
a verify slate must reproduce the decode it verifies and a boot may serve
two bands to two slates of the same sequence. A ladder that mixed keys
would make one row's output a function of how many rows travelled with it —
the row-keyed defect the step-class channel exists to remove, arriving
through the band instead of through the row count.

The edges must also be strictly increasing and terminated by an open band,
or some slate width maps to no band or to two.

A recorded ladder as this module's own tuple, or ``None`` if unreadable.

Never raises: it runs on the boot path over a file an operator may have
hand-edited, and the fail-safe direction is the package default — which is
the ladder every receipt was measured against. ``null`` is how the open
last band is spelled in JSON.

How many scale groups shape ``shape_id`` reads along k. ``0`` = per token.

THE GRANULARITY SEAM. Every consumer -- the boot reservation's sizing, the
adapter's scale view, and the legality check -- asks this rather than
assuming one scale per row, so the second granularity is a change to
:data:`GROUP_SCALE_SHAPES` and ``a_prep`` and to nothing else here.

What fixes shape ``shape_id``'s REDUCTION ORDER over ``k``.

``(TILE_K, activation group width)``, ``0`` for a per-token scale. Two
entries that agree here contract ``in_features`` in the same order and fold
the same partial sums at the same boundaries, so a pin that moves between
them changes which blocks do the work and not what the work computes.

THE FUNCTION EXISTS BECAUSE THE VERIFY CLASS TAKES A DIFFERENT TILE. An
M-dependent tile pin is only sound if the tile is not part of the
arithmetic, and "TILE_M is only a block count" is precisely the kind of
claim that is true until a shape is added for which it is not. So it is a
value a test can compare rather than a property a reader has to re-derive
from the table; ``tools/int8_gemm/row_invariance.py`` is the same assertion
on the card, against the kernel rather than against the table.

Deliberately NOT the whole tuple. TILE_N, the warp split, the stage count
and the swizzle move which block computes an output element and in what
order blocks are scheduled; none of them reassociates a block's own k loop.

The tiles the VERIFY class may be pinned to, narrowest TILE_M first.

Derived from the table rather than listed beside it: an entry is a
candidate when it is FOLD-capable (the adapter folds ``svh`` in the
epilogue and has no second pass), it is not a null-decode control, it
carries the quantiser this boot selected, its reduction key matches that
class default's, and its TILE_M is below the prefill tile's — which is what
makes it a small-M candidate rather than just another entry.

The candidate SET is what a bench sweeps and what a receipt names one
member of. It is not itself a decision: :func:`int8_verify_default_shape`
is, and ``shipped_pins.json`` overrides that per geometry.

One linear's frozen int8 launch: the DECISION, plus the scratch it uses.

Everything the adapter needs that is not an argument of the call, resolved
ONCE when the linear is bound and never re-derived per call. That is not
only speed: ``cb`` derived per call would be a branch on the descriptor
inside a captured region, and ``w_scale`` read per call would be a pybind
round trip 409 times a forward.

Mutable in exactly one direction and at exactly one seam: the shape is
frozen when the pin resolves, and :func:`bind_scratch` fills the buffers
when the boot reservation runs. Between those two seams the plan is
incomplete and :func:`~arbi_serve.weight_quant.exl3.custom_op._int8_leg_serves`
refuses it by name -- ``a_frag is None`` is the sentinel, and it is a
sentinel rather than an empty tensor because an empty tensor would silently
make every ``[:0]`` view legal.

ONE PLAN PER STEP CLASS, and a linear carries one for each class the leg
serves it in. The classes ask the kernel different questions -- a prefill
chunk fills many m-tiles and a verify slate a fraction of one -- so they
are pinned to different TILE_M, which is legal only because they share a
reduction key (:func:`int8_reduction_key`). Two plans rather than one plan
with two shapes, so nothing on the serving path has to select a field: the
call site picks the OBJECT from the class it already declared, and
everything downstream reads a frozen plan exactly as it did when there was
one. ``label`` names the class for the log and the repr, and is never read
by a launch.

Name the target arch through ``TORCH_CUDA_ARCH_LIST``, never through ``-arch``.

THE POINT IS THAT THE IMAGE BAKE AND A SERVED BOOT PRODUCE THE SAME COMMAND
LINE. torch's JIT cache is ninja, and ninja rebuilds when the command
changes — so a bake compiled under ``TORCH_CUDA_ARCH_LIST=8.9`` and a boot
that passed its own ``-arch=sm_89`` would be two different commands over one
output, and every armed boot would recompile the thing the image had already
compiled. The bake would look present and buy nothing, which is the exact
shape of a lever that reports as applied and is inert.

So the arch is named in ONE vocabulary, torch's own, and this function only
supplies it when the caller has not: a builder has no device to ask and sets
the list, a serving boot has a device and no list. Both then reach
``_get_cuda_arch_flags`` with a single-arch list and get identical
``-gencode``.

Boot-time and single-threaded by contract, which is what makes mutating the
environment acceptable here; it is restored either way.

Compile (or pull from the ext cache) the int8 GEMM. ``None`` if it cannot.

Never raises. A build failure leaves the leg refusing by name rather than
failing a boot: the shipped legs serve every call correctly without it, so
an unbuildable optional kernel is a performance outcome, not an outage.

Build the kernel AND warm ``a_prep``. Boot-time; returns the module or ``None``.

Both halves, because the adapter needs both and a JIT compile is not
something a forward may discover. ``a_prep``'s extension is compiled here
for the same reason the GEMM's is: its first call would otherwise run
``torch.utils.cpp_extension.load`` inside a served step -- and inside a
cudagraph capture, which is not survivable.

The built kernel, or ``None`` — WITHOUT triggering a build.

The reader the serving path uses. It must never compile: that decision
belongs to :func:`load_int8_kernel` at boot, and a forward that discovers a
missing kernel has to refuse rather than stall a step behind ``ninja``.

Why ``shape_id`` cannot serve this geometry, or ``None`` when it can.

Every clause is a precondition the kernel itself enforces with a
``TORCH_CHECK``, asked here instead at BOOT so the answer is "this linear is
not pinned" rather than an exception inside a served step. The last one is
asked of the COMPILED kernel rather than of the table above: whether a
shape's stages fit in shared memory at a given trellis width and codebook is
a property of the build (the register allocation is part of it), and
``shape_probe`` throws for a combination that does not.

``cb`` is therefore required rather than defaulted. The codebook is a
template parameter of the mainloop, the two decodes do not cost the same
registers, and a probe run under the other codebook reports occupancy for a
kernel this linear will not run.

The ``cb`` the kernel is told, derived from the linear's OWN descriptor.

A delegate to
:func:`~arbi_serve.weight_quant.exl3.int8_policy.int8_codebook_id` rather
than a second derivation of the same two booleans: that is the one place
the numbering lives, so a linear the policy SERVES and a ``cb`` the kernel
is HANDED can never disagree.

Never a literal. The kernel decodes mcg and mul1 on a template axis, an mcg
trellis is byte-for-byte indistinguishable from a mul1 one, and the two
grids differ by 2.16x — so a hardcoded codebook is a silently wrong weight
rather than a fault.

Resolved at BIND time and frozen into the plan, so the serving path does not
branch on the descriptor. Raises for a linear with no int8 codebook, which
the pin declined before it could get here.

``(A-fragment elements, scale elements)`` one call at ``rows x size_k`` uses.

The one definition of the scratch arithmetic, so the boot reservation sizes
what the call actually slices instead of a second derivation of it. Pure
over Python ints -- it reads no tensor and cannot sync.

The nine-argument entry the served int8 leg calls. Writes ``y``, returns it.

Argument order is the CALL SITE's and is pinned by
``tests/test_exl3_int8_gemm_arity.py``; it is not the kernel binding's.
``mcg``/``mul1`` are accepted and not read: they were resolved into
``plan.cb`` when the linear was bound, and re-deriving them here would put a
descriptor branch on the critical path of 409 linears and inside a capture.

ONE PASS ON EACH SIDE. ``a_prep_frag`` folds suh's pre-scale, the 128-wide
Hadamard, the per-token amax, the int8 quantiser, the SIGMA relabel of the
contraction axis and the fragment-order store into a single read of ``x``;
the GEMM's epilogue folds ``r_scale * svh`` and the output-side Hadamard
into the store of ``y``. Neither ``had_r_128`` pass the shipped leg runs
survives.

NOTHING HERE ALLOCATES AND NOTHING HERE SYNCS. The two views are
metadata-only and always start at the buffer's base, so every captured
graph bakes one address; ``rows_pad`` is Python integer arithmetic over
``x.shape``, which is static at capture.

``narrow`` RATHER THAN ``[:n]``, and that is the difference between a loud
failure and a corrupted one. A call site wider than the row bound the
reservation was sized from would take a SHORTER slice than it asked for,
and ``a_prep_frag`` would then write ``rows_pad * size_k`` bytes into it --
past the end of the buffer, into whatever the pool holds next. Slicing
clamps and says nothing; ``narrow`` refuses, in C++, without reading a
device value. It costs nothing and it cannot be the check that could not
fail.

Drop the scratch binding, returning this plan to its unarmed state.

The safe direction whenever the buffers a plan names may no longer be
the ones it should use: an unarmed plan is REFUSED by name
(:func:`~arbi_serve.weight_quant.exl3.custom_op._int8_leg_serves`),
while one left pointing at a released buffer is written through by the
GEMM — an illegal access that poisons the context for the process.

Whether the int8 EXL3 prefill GEMM serves THIS step, resolved in one place.

Two questions, kept apart because they have different owners:

  * **Is the leg armed at all** — a tri-state operator flag over an arch
    default (:func:`resolve_int8_gemm`), the shape ``ARBI_SERVE_AWQ_NO_A8``
    established: unset means "ask the arch", ``0``/``1`` mean the operator has
    decided and the arch is not consulted. The arch default lives in
    :func:`int8_gemm_default_for_cap` and is the ONE line that moves the day
    the kernel is correct — with a reason recorded beside it, the way
    ``exl3_prefill_f16acc``'s default move records the pin bump that made the
    capability exist.
  * **Which CLASSES it serves** — :func:`int8_class_refusal`. PREFILL always,
    VERIFY and DECODE each on a tri-state (``0`` never / ``1`` always /
    ``auto``). #1861's amended scope routes verify by STEP TYPE for invariance
    while leaving open, as an empirical question, whether verify is FASTER
    there. ``docs/receipts/int8-verify-small-m-2026-09-05.txt`` answers it and
    the answer is "it depends on the slate": a verify slate is ``B x (K+1)``
    rows, and the int8 chain loses to the trellis leg on a narrow slate and
    wins on a wide one. ``auto`` is that measurement expressed as a rule — the
    leg serves a verify forward whose OWN rows reach
    ``exl3_int8_verify_min_rows`` and refuses it by name below that.

    DECODE IS THE SAME RULE OVER THE SAME PHYSICS, and it was a categorical
    refusal until ``docs/receipts/int8-decode-crossover-2026-09-07.txt``. A
    decode step is ``B`` rows, and ``B`` at ``--max-batch 256`` is 256 rows —
    every row count a verify slate presents and more. What decides the leg is
    not the class but the row count against the served trellis pin: the
    trellis GEMM walks M in ``TILESIZE_M`` steps with a ``grid.sync`` between
    them and re-reads the whole packed weight per step, so from one row past
    the deepest tile the pin serves it pays a SECOND full weight pass, and the
    int8 chain — one pass at any M — wins from there on every 27B geometry.
    That is why the DECODE threshold is DERIVED from the served ladder
    (:func:`int8_decode_min_rows_from_tiles`) rather than typed: it is
    parametrized on the one input it depends on, and it moves the day a deeper
    family is pinned. The step class survives as the ALIGNMENT constraint and
    nothing else: every int8 tile shares one reduction key, so a decode row
    and a verify row the leg serves take one arithmetic.

    WHAT ``auto`` GIVES UP, stated rather than left to be discovered. Under
    ``1`` (and under ``0``) a sequence's verify arithmetic is a property of
    its class alone. Under ``auto`` it also depends on the slate width, and a
    slate is ``B x (K+1)`` — so a request's verify numerics move with the
    concurrency it happened to share a step with. That is bounded, not
    unbounded: the verify graphs are captured per ``(B, S)`` and an
    intermediate ``B`` pads UP to the next captured rung, so the decision is
    taken once per graph at record time and every replay serves what was
    recorded. A deployment that needs the strict property sets ``0`` (verify
    takes decode's arithmetic) or ``1`` (verify takes the int8 leg at every
    width); ``auto`` buys throughput with it.

IMPORT-LIGHT by construction. Nothing here imports torch, exllamav3 or the
engine at module scope, so the forward path can ask "is this armed" without
dragging the EXL3 module into a bf16 deployment's every step. Every RULE is
pure over its arguments and unit-testable without a GPU; only
:func:`int8_gemm_armed`, the convenience that binds those rules to this
process, touches the flags registry and the device — and it re-reads the flag
every call, because the param is ``scope="runtime"`` and a latched read is how
``exl3_fused_reconstruct`` came to report an admin override as applied while
the served leg never moved.

The VERIFY half is read the same way here and is nonetheless CAPTURE-AFFECTING
at the config layer: verify runs inside captured graphs, so a live flip would
be inert while reading as applied. Re-reading it costs nothing and keeps this
function pure over the flags; what stops the inert flip is the param's
``scope="backend"``, which rebuilds and recaptures. ``auto``'s row threshold
is capture-affecting for the same reason and by the same mechanism.

Arch default for the int8 prefill GEMM. Pure over the capability.

``True``: the kernel is this package's own (``csrc/exl3_i8_gemm.cu``, built
as one torch extension), so it ships wherever arbi-serve does, and on the
Ada parts served today it is the faster prefill leg.

The parameters are the seam for an arch whose vendor path does the exl3
format natively — there the right answer is that path, not this kernel.

Normalise whatever the flag carries into one of :data:`VERIFY_MODES`.

The field is a string tri-state, but the value reaching the forward path
can also be a bool: a live override overlay takes the operator's parsed
value, and the flag was a plain bool before ``auto`` existed. Folding both
spellings here rather than at each reader is what keeps "the flag says on"
from meaning two things in two places.

An UNRECOGNISED value resolves to ``0`` — verify then takes exactly
decode's arithmetic, which is the fail-safe direction and the one the leg
is compared against. A typo fails loud at boot instead, in the env parser;
this is the serving-path read and it must not raise inside a forward.

Whether a VERIFY slate of ``rows`` rows takes the int8 leg. Pure.

THE RULE ``auto`` IS. The int8 chain and the trellis leg cross over as the
slate widens — narrow slates go to the trellis leg, wide ones to int8 —
and ``docs/receipts/int8-verify-small-m-2026-09-05.txt`` is the
measurement, per linear and per tile, that locates the crossover on the
config this defaults for. ``min_rows`` is that crossing point as a knob, so
a different card, tile or draft depth is a re-measurement rather than a
code change.

``>=`` and not ``>``: ``min_rows`` names the narrowest slate the leg is
asked to serve, so the rung it is set to is INSIDE the served set.

The narrowest decode batch the int8 leg is asked to serve, from the served pin.

``deepest_tiles`` is, per pinned geometry, the deepest ``TILESIZE_M`` its
trellis-leg ladder reaches. The trellis GEMM reads its whole packed weight
once per ``TILESIZE_M`` rows (``exl3_gemm_kernel.cuh`` tiles M serially
with a ``grid.sync`` between tiles), so one row past the deepest served
tile the trellis leg pays a SECOND full weight pass while the int8 chain
still pays one. ``docs/receipts/int8-decode-crossover-2026-09-07.txt`` and
``docs/receipts/int8-verify-slate-geometry-2026-09-07.txt`` measure that
the int8 chain wins from that row on every 27B geometry against the served
member, and loses or ties below it. The threshold is therefore the
mechanism's own edge — ``max(deepest tile) + 1`` — and not a stopwatch's:
it is parametrized on the one thing it depends on, so a boot that pins a
deeper family (K16N128's 96-row members) moves it to 97 without a code
change, and it can never be a literal that expires when the ladder moves.

``max`` over the geometries and not ``min``, deliberately: the rule is one
threshold for the whole forward (every linear in a step sees the same
rows), and below the deepest geometry's second pass the int8 chain loses
on the wide-N geometries that dominate the weight bytes — the receipt's
whole-model rows at 24-32 read as a wash. ``None`` when nothing is pinned:
there is then no served ladder to derive from, and
:func:`int8_class_refusal` refuses by name rather than guessing.

Why the int8 leg does not serve THIS step, or ``None`` when it does.

A reason rather than a bool because the refusals are different facts an
operator acts on differently: :data:`OFF_CLASS` is "this class is not
routed here", :data:`VERIFY_BELOW_MIN_ROWS` / :data:`DECODE_BELOW_MIN_ROWS`
are "the class IS routed here and the rule declined this width", and
:data:`DECODE_MIN_ROWS_UNRESOLVED` is "the rule could not run". Pooled into
one name, a boot that never saturates would read identically to a boot
with the flag off.

DECODE takes the SAME rule as VERIFY on its own rows, under its own
tri-state (``decode``) and its own threshold (``decode_min_rows``, derived
at boot by :func:`int8_decode_min_rows_from_tiles` or pinned by the
operator). Routing on the declared step's rows is not the row-keyed defect
the class channel removed: that defect was a per-LINEAR read of the
activation that made one row's arithmetic follow the geometry it landed
on; this is one host int per FORWARD, declared once, identical for every
linear in the step and identical at record and replay.

``rows`` is the slate's own row count, taken from the DECLARED forward
(:func:`arbi_serve.runtime.forward_declaration.declare_forward`) and never
from the linear's activation: a linear that read its own ``x`` would make
one row's arithmetic follow the geometry it landed on, and a per-linear
read cannot be answered under capture at all. It is ignored outside
``auto``.

``rows`` and ``min_rows`` carry NO DEFAULTS, deliberately. A defaulted
threshold would have to restate the field literal here — a second source
for the shipped crossing — and a caller that forgot it would silently get
``min_rows=0``, i.e. mode ``1`` under ``auto``'s name. Naming both at every
call site is what makes the rule's inputs impossible to omit.

:func:`int8_class_refusal` as a boolean, for callers that want no name.

Kept as the one-line derivation of the reason rather than a second rule:
two functions each spelling the class table would be two places for the
tri-state to drift.

The named refusal for a linear's CODEBOOK, or ``None`` when served.

The int8 kernel decodes exl3's ``mcg`` and ``mul1`` codebooks, and the
codebook is a TEMPLATE PARAMETER of the mainloop rather than a constant
baked into it — ``exl3_i8_gemm.cu``'s ``i8_decode<CB, ELIDE_CLAMP>``, on
its own axis independent of the bit width, the way exl3's own
``dq_dispatch<bits, cb>`` is.

``mcg`` is exl3's DEFAULT codebook (``exl3_convert_mamba_hybrid.py``:
``--codebook choices=("mcg","mul1") default="mcg"``) and the one every
checkpoint this repo owns uses; the leg shipped decoding ``mul1``, the
legacy one, and therefore refused every model we serve. It now serves both.

WHAT DOES NOT CHANGE IS WHY THIS FUNCTION EXISTS. The codebook is still the
one precondition whose violation would be SILENT: a wrong bit width changes
the trellis's last dimension and a wrong dtype changes its type, so both are
caught by shape, while an ``mcg`` trellis is byte-for-byte
indistinguishable from a ``mul1`` one — same shape, same dtype, same
contiguity. Adding mcg does not retire the check, it moves the boundary:
exl3's ``default`` codebook (no seed at all) has no int8 decode, and a
descriptor carrying BOTH seeds is corrupt. Those still refuse, by name.

The seam is ``mcg`` XOR ``mul1``, never "not mcg" — the latter reads as a
two-codebook world and would serve every default-codebook checkpoint with
mul1's arithmetic. Two distinct names, not one, because *which* wrong
codebook is what an operator reading the counter needs: a default-codebook
checkpoint is an older export, and both-seeds-set is a corrupt descriptor
rather than a policy outcome.

Pure over its arguments — ``mcg`` and ``mul1`` are the booleans exllamav3's
``LinearEXL3`` carries (``self.mcg = self.mcg_tensor is not None``), which
are the same two the shipped ``reconstruct`` call is given.

The kernel's ``cb`` for a linear, or a raise. One place, not two.

The numbering is exllamav3's own (``decode_3inst<cb>``): 1 = mcg, 2 = mul1.
It is derived from the SAME two booleans :func:`int8_codebook_refusal`
reads, so a linear the policy serves and a ``cb`` the kernel is handed can
never disagree — which is the shape of the defect that would be silent, and
the reason this is not left to each call site to spell out.

Raises rather than returning a default: the caller is expected to have
consulted :func:`int8_codebook_refusal` first, and a fallback codebook here
would be a check that cannot fail.

Is the int8 leg armed for THIS process, right now.

Reads the tri-state fresh every call. Cheap enough for the forward path:
one flags-registry attribute read plus a memoised capability, which is what
the two EXL3 channels beside it already cost.

Pool routing for the EXL3 kernel descriptor's bsz-1 input row (``xh``).

WHAT THE BUFFER IS. ``exllamav3.LinearEXL3.__init__`` hands its C++
``BC_LinearEXL3`` descriptor a preallocated ``(1, kernel_in_features)`` fp16
row — the staging buffer the trellis GEMM Hadamard-transforms a SINGLE input
row into. It comes from ``exllamav3.util.tensor.g_tensor_cache``, a
process-global ``{f"{device}/{shape}/{dtype}/": (refcount, tensor)}`` dict, so
every linear sharing a kernel input width shares ONE row: a model is charged
one buffer per DISTINCT width, not one per linear. On a 27B EXL3 checkpoint
that is five buffers totalling ~114 KiB.

WHY IT NEEDS A HOME. The descriptor is constructed inside
:meth:`~arbi_serve.weight_quant.exl3.linear.EXL3LinearBase._build_inner`, which
runs during the weight load and at every compaction rebind — both deliberately
OUTSIDE ``model.weights`` (a cuMem pool never returns a freed block, so the
graph-build transients must not land in it). An allocation with no pool active
goes to the torch DEFAULT allocator, and a live tensor there is what the
``unpooled.torch_default_pool`` ledger row reports: real, in use, and
unbudgeted by construction. What routing buys is not the bytes — it is an
``unpooled_owners`` list that reads EMPTY, so any entry on it is a signal
rather than residue an operator learns to scroll past.

WHICH POOL. ``capture.io_buffers``, for the reason that pool exists: the
bsz==1 decode leg calls ``bc.run`` -> ``run_gr(.., nullptr)``, which uses this
buffer WITHOUT allocating, so every captured single-row decode graph bakes its
``data_ptr`` into its kernel arguments. That is the pool's definition
("persistent kernel I/O buffers captured graphs reference by ``data_ptr``",
kept out of the capture mempool so no address reuse can corrupt a replay), it
is where the sibling EXL3 tenant already lives
(:mod:`arbi_serve.engine.exl3_reconstruct_scratch`), and it is cuMem-backed
with a stable VA, so a park/wake round-trip remaps the bytes at the identical
address the graph baked. A per-step arena is the wrong home for the mirror
reason: the arena re-homes what it holds onto fresh addresses.

WHOSE. That pool belongs to ONE member, and the cache it is reached through is
process-global and keyed by (device, shape, dtype) with no member
discriminator — so a second member with the same width would otherwise inherit
the FIRST member's row and write into a parked member's unmapped VA. The cache
is therefore registered in
:data:`~arbi_serve.engine.member_scratch_retire.MEMBER_SCOPED_DEVICE_CACHES`,
which parks it with its member and hands it back on that member's wake, and
:func:`set_kernel_scratch_pool` is repointed at each seam that changes which
member is active (build, wake, member reset) — the same contract
``cublas_warmup.set_cublas_workspace_pool`` carries.

Point the kernel-row allocation at ``pool`` (``None`` disarms).

Called at every seam that changes which member owns the pools: the build
(before the weights load builds the first descriptor), a residency wake,
and the member reset. Leaving a previous member's pool armed is the fault
this exists to prevent — its physical is unmapped behind its VA while it is
parked, so a descriptor built against it would write unmapped memory.

Route the descriptor's row allocation into the armed pool.

Wraps the ``LinearEXL3(...)`` construction and NOTHING else: that call's
only device allocation is the ``g_tensor_cache`` row (its scale buffers are
passed in already materialised), so the context cannot sweep an unrelated
tensor into a pool that is capped at the freeze.

A miss — no pool armed, or a pool that refuses to be entered — falls back
to the default allocator. Correct, and not silent: the row then shows up by
name in the ``unpooled.torch_default_pool`` owner walk, which is the
surface that reports exactly this. Only the ENTRY is guarded: an exception
out of the construction itself is the caller's, and swallowing it would
hide a failed weight bind behind an accounting helper.

Distinct EXL3 kernel input widths bound under ``roots``, sorted.

The KERNEL width (``_kernel_in_features``), not the declared one: a linear
whose ``in_features`` is not a Hadamard-block multiple carries a padded
trellis, and the descriptor's row is that padded width.

Only linears whose descriptor EXISTS are counted: an unbound placeholder
(a meta-graph build before the fill, a checkpoint with no EXL3 tensors)
has constructed no descriptor and therefore drawn no row.

Duck-typed on the two attributes rather than an ``isinstance``, so pricing
the pool never imports exllamav3 on a run that has no EXL3 weights.

The pinned exllamav3 kernel-shape table, read from its header.

WHY THIS SHIPS. BOOT, inside
the serving image, to freeze one ``(TILESIZE_K, TILESIZE_N)`` family per EXL3
geometry. So the reader must live in the shipped package, never in ``tools/``:
``pyproject.toml`` ships ``packages = ["arbi_serve"]`` and the Dockerfile copies
only ``arbi_serve`` and ``client``, so a ``tools`` import resolves under a dev
bind-mount that puts the repo root on ``sys.path`` and raises
``ModuleNotFoundError`` in every shipped image — where the pin then declines and
leaves the row-count-dependent auto dispatch it exists to remove.

The table is still defined in exactly one place; ``tools.exl3_shape_receipts``
re-exports these names rather than restating them, so a pin move still cannot
leave a stale copy behind. The dependency simply points the other way, from the
repo-only tool into the shipped package, so the boot path never needs ``tools``.

Locate the pinned header WITHOUT importing exllamav3.

Importing the package pulls in ``exllamav3.ext``, which re-enters torch's
JIT and rebuilds the CUDA extension from source when no prebuilt one is
cached — minutes of compile inside what is meant to be a table lookup.
``find_spec`` resolves the package directory without executing it.

``{shape_idx: "M,K,N"}`` parsed from the pinned exllamav3 header.

Empty when the header cannot be read. Every refusal built on this keys on
the raw tag, so an empty table degrades a message and never a check.

``EXL3_GEMM_NUM_SHAPES`` as the pinned header declares it.

Read from the same file as :func:`shape_table` so a drift test can compare
the two without a built extension, and therefore without a GPU.

``(usable, reason)`` — why :func:`shape_table` is empty, when it is.

``TILESIZE_M`` row-block table). Three causes want three different fixes,
so they must not collapse into one message — nor into "not configured",
which is a fourth. The caller logs this; it never raises.

``{shape_idx: (TILESIZE_M, TILESIZE_K, TILESIZE_N)}``.

The boot pin's view of :func:`shape_table`: same source, parsed into the
tuple the family grouping needs. Empty exactly when the table is, which
makes the pin decline rather than guess.

The exl3 trellis layout, as the int8 (``mma.m16n8k32.s8``) prefill GEMM reads it.

WHY THIS EXISTS, AND WHY IT DOES NOT REPACK ANYTHING.

exl3's ``dq_dispatch`` extracts eight weights per lane and hands back two
``FragB``. Read as one object those eight weights are ``2 n-columns x k=16``,
while an int8 MMA fragment wants ``1 n-column x k=32`` — the same eight weights
on a different axis, which reads like a demand for a K-major repack of the
bitstream.

It is not one. The two fragments were never one k16 x n16 object: the mainloop
already consumes them as two INDEPENDENT n8 fragments, and their n axis is
already exactly the s8 B fragment's n axis. The missing sixteen k rows are the
NEXT k-tile. What is genuinely wrong is the ORDER of the rows inside each half,
and order along a contraction axis is free -- ``sum_k A[m,k] B[k,n]`` is
invariant under any permutation of k applied to BOTH operands. So the fix is a
relabel of k, and it lands on the activation, where the int8 quantiser owns its
store and the permutation is a register byte-shuffle before a store that has to
happen anyway.

The consequence is that the weight format does not change: not one byte of a
shipped checkpoint moves, no tensor is added, the bits per weight are identical,
and every existing consumer (leg A, leg B, the fused reconstruct, decode,
verify, TP sharding) keeps working because it is reading the same file.

This module is the machine-readable form of that contract: the position map, the
relabel, the per-fragment gather the kernel performs, and a reader for the packed
bitstream. It ships inside ``arbi_serve`` rather than ``tools`` for the same
reason :mod:`.kernel_shape_table` does — a boot-time guard may need it and a
``tools`` import does not resolve in the serving image.

``position -> r*16 + c`` inside one trellis block, r = k, c = n.

A restatement of exllamav3's ``quantize.tensor_core_perm``, which is the
``mma.m16n8k16`` f16 B-fragment layout applied twice: lane ``t`` owns the
eight positions ``8t..8t+7``, the first four being n-column ``t//4`` and the
last four n-column ``t//4 + 8``, each over k rows
``[2a, 2a+1, 2a+8, 2a+9]`` for ``a = t % 4``.

Restated rather than imported because importing ``exllamav3`` pulls in the
compiled extension, and because this map is a claim about the PINNED fork
that a test should be able to falsify against it.

Physical k row inside a 16-group -> the logical k row the s8 kernel calls it.

A rotation of bits 1..3, which is exactly what turns exl3's f16 fragment
rows ``[2a, 2a+1, 2a+8, 2a+9]`` into the s8 fragment's ``[4a, 4a+1, 4a+2,
4a+3]``, in element order, for every ``a``.

Apply the relabel to an activation's k axis: ``out[..., SIGMA[r]] = x[..., r]``.

The reference implementation of what the int8 activation quantiser must do
at its store (or the 128-wide Hadamard at its own). Present so a null control
can run the permuted and unpermuted contractions against each other; the
serving path never calls it.

``(word index, right shift)`` for each of the 256 trellis windows.

A block is 256*K bits with NO header and NO padding, read as a big-endian bit
stream over its ``uint16``s with each adjacent pair swapped in storage (the
``SWAP16`` on the ``uint32`` view in ``pack_trellis_kernel``). Weight ``p``
occupies stream bits ``[pK, pK+K)`` MSB first, so ``uint32`` word ``i`` is
stream bits ``[32i, 32i+32)`` with the earliest bit at bit 31.

The trellis word decoded at position ``p`` is the 16-bit window ENDING at the
end of weight ``p`` -- stream bits ``[(p+1)K-16, (p+1)K)``, circular within
the block. That circularity is why the block carries no window overhead, and
why cutting an eight-position run out of it costs ``7K+16`` bits.

Decode the 256 16-bit trellis words of each block, exactly as ``dq`` reads them.

``packed_u16`` is ``(..., 256*K/16)`` as stored on disk. Returns
``(..., 256)`` ``uint16``, indexed by trellis POSITION -- map to ``(k, n)``
with :func:`tensor_core_perm`.

The s8 B-operand gather, as ``(lane, n_half, reg, elem) -> (k_tile_parity, position)``.

One ``mma.m16n8k32.s8`` B operand for logical k super-tile ``u`` and n-half
``h`` is filled by TWO ``dq_dispatch`` calls, on blocks ``(2u, ntile)`` and
``(2u+1, ntile)``, taking fragment ``h`` from each::

    b0 = pack_s8(dq_dispatch(block[2u  ][ntile], lane<<3).frag[h])
    b1 = pack_s8(dq_dispatch(block[2u+1][ntile], lane<<3).frag[h])

Element to byte is the identity -- decoded ``d0..d3`` go to bytes 0..3 -- so
there is no cross-lane movement and no shuffle. Returns the position each
``(reg, elem)`` slot reads, so a test can check it against
:func:`tensor_core_perm` and :data:`SIGMA` rather than against prose.

EXL3 weight-quantization Linear classes — concrete impl of the
:mod:`arbi_serve.weight_quant.base.QuantLinearBase` contract.

Bridge between arbi-serve's :class:`LinearBase` family and
:class:`exllamav3.modules.quant.exl3.LinearEXL3`. EXL3 stores each
projection as ``{trellis (i16), suh (f16), svh (f16), mcg|mul1
(optional codebook seed)}``
instead of a dense ``(out, in)`` weight; the inner ``LinearEXL3``
holds those tensors and a precomputed ``BC_LinearEXL3`` kernel
descriptor.

The binding API is :meth:`exl3_load`, called by
:meth:`arbi_serve.weight_quant.exl3.backend.EXL3Backend.bind` with the
whole per-linear tensor set at once — ``trellis``, ``suh``, ``svh``, and
the codebook seed when the checkpoint carries one. The dense
:meth:`weight_loader` is locked out by the parent
:class:`QuantLinearBase`.

dtype handling. The EXL3 kernel ships fp16-only on Ada/Hopper;
arbi-serve's hidden state is bf16 by default. We cast at the layer
boundary: ``x.to(half) → kernel → out.to(orig_dtype)``. Two extra
elementwise ops per Linear, but it lets the model dtype stay bf16
everywhere else (norms, activations, residuals).

TP. The trellis quantization unit is 16 elements along each axis.
Column-parallel slices trellis dim 1 + svh; row-parallel slices
trellis dim 0 + suh. Both slicings happen inside :meth:`exl3_load`
based on the active :class:`ParallelConfig`. The slicing math is
unit-tested at simulated TP>1 ranks; multi-GPU end-to-end is not
verified on a single-GPU build host.

LoRA. EXL3 + LoRA composes correctly — :class:`QuantLinearBase`
inherits :meth:`LinearBase._maybe_apply_lora` which dispatches to the
BGMV kernel; the kernel accumulates the correction on top of the
trellis-dequant output regardless of how that output was computed.

Tell Dynamo not to specialize on these tensors' shapes.

Each populated buffer's every dim is marked dynamic so no
``size == const`` guard is emitted. Best-effort: a tensor already
captured by a live compiled graph (or a torch build without the API)
is skipped silently. Empty placeholders are skipped (nothing bound
yet). See the call site in ``_build_inner`` for why this is
serving-perf-neutral.

Common surface for the three EXL3 parallel variants.

Carries the trellis / suh / svh buffers (persistent — captured in
the warm-reload flat dump) and the lazily-built
:class:`exllamav3.modules.quant.exl3.LinearEXL3` kernel descriptor
(derived, rebuilt by :meth:`rebind_after_compaction`).

EXL3 equivalent of :class:`ReplicatedLinear`.

Used for the lm_head when EXL3 export quantizes it (most EXL3
Qwen3 dumps do — they untie embeddings from lm_head because the
embedding stays dense).

``__init__`` + ``forward`` come from :class:`ReplicatedLinearMixin`.

A leading output window of an already-bound EXL3 linear, over ITS tensors.

``out_features`` is the window width; ``exl3_load`` takes the SOURCE
linear's ``trellis`` / ``suh`` / ``svh`` unchanged and binds them by
reference, so this module adds no device bytes at all.

Output columns live on the trellis's middle dimension in output order —
the same fact :meth:`EXL3ColumnParallelLinear.exl3_load` slices a TP shard
on — so a window at offset 0 needs the leading ``out_features //
_TRELLIS_BLOCK`` blocks and ``svh``'s leading ``out_features`` entries.
Both are leading, so no slice has to be materialized: ``exl3_gemm`` takes
the window as its ``size_n_out`` bound, keeps B's k-row pitch at the
stored width, and reads exactly those blocks.

The window is a whole number of :data:`_HAD_BLOCK` blocks because the
runtime Hadamard along the output axis is block-diagonal in 128 — the
same alignment :class:`_EXL3HadAlignedShardMixin` enforces on a TP shard.

Several 128-aligned output windows of one EXL3 head, as one row.

:class:`EXL3OutputPrefixLinear` covers the window at offset 0, which needs
no slice because ``exl3_gemm`` bounds the output at ``size_n_out`` and
reads the LEADING trellis blocks. A window at a NON-zero offset needs the
same tensors cut at that offset — the cut
:meth:`EXL3ColumnParallelLinear.exl3_load` already performs for a TP output
shard. The kernel needs no output-row offset argument and gets none; the
offset is expressed on the caller's side, which is what makes this a
binding change and not a kernel change (measured: a trailing 3-block window
of the 27B lm_head reproduces the full head to within half the leading
window's own error — ``experiments/drafter_union_reach/``).

The rows are the windows' outputs concatenated IN WINDOW ORDER, so column
``j`` of this module's row is NOT token ``j``. The map lives in
:class:`~arbi_serve.draft_vocab_prefix.DraftReach`, carried here as
``_draft_reach`` so every consumer reads one owner rather than
reconstructing the offsets.

A non-leading window is MATERIALIZED: a slice of the trellis's middle
dimension is not contiguous, and the kernel reads a packed block range. The
leading window still costs nothing, so the resident cost is exactly the
non-leading windows — 1.41 MiB for the 3-block tail this exists for.

The GLOBAL output window ``[0, N)`` over an ALREADY vocab-sharded head.

:class:`EXL3OutputPrefixLinear` is replicated: it windows one whole head.
When the head it would window is column-parallel, no rank holds ``[0, N)``
— rank ``r`` holds global columns ``[start_r, start_r + width_r)`` — so
the window is taken PER RANK, as the intersection of the global window
with the rank's own shard, and the ranks compose it back.

The intersection is always a LEADING window of the rank's shard, because
the global window is a prefix and the shards are contiguous in rank order:
``window_r = clamp(N - start_r, 0, width_r)``. So each rank's half is
exactly what the replicated window already does — an ``exl3_gemm`` output
bound — over the tensors that rank already holds. Nothing is copied and
nothing moves between ranks.

Composition is the column-parallel all-gather the un-windowed head already
fires, at the per-rank window width instead of the shard width. The
windows are not equal across ranks, so the gather runs shape-uniform at
``max(window_r)`` — never wider than a shard, hence never wider than the
un-windowed head's — and each rank's contribution is trimmed back out on
the way to the composed row. The trim reads the STACKED gather rather than
``all_gather``'s concatenated one, so the whole composition is a single
``N``-wide copy instead of a full-width one followed by a narrow.

Two ranges of ``N`` behave differently and both are correct:

* ``N <= width_0`` — only the low ranks carry a window, the high ones
  carry none, and both the GEMM and the gather shrink with ``N``.
* ``N > width_0`` — the low ranks' windows ARE their whole shards, so
  their GEMM and the gather are exactly what they were without the lever;
  what the lever still buys is the straddling rank's narrowed GEMM, and
  what it still MEANS is that the drafter cannot propose an id ``>= N`` —
  the same drafter it would be at TP=1.

A rank whose window is empty binds no kernel descriptor and runs no GEMM;
it joins the gather with an unread buffer so the collective stays
shape-uniform. Its columns are trimmed away on every rank, so their
contents are never read.

Shard the parallel axis on :data:`_HAD_BLOCK` boundaries, allowing
uneven shards.

EXL3 applies a 128-wide Hadamard along BOTH axes at runtime
(``had_r_128``: ``y = (x.view(-1, 128) @ had_128) * scale``). The
transform is block-diagonal in 128, so a shard reproduces the unsharded
result exactly when — and only when — its offset and width are whole
128-blocks; ``exllamav3`` refuses anything else at the first reconstruct.
The 16-element trellis block is the PACKING unit and is too weak a
constraint to use here.

Whole 128-blocks need not divide evenly across ranks, so this splits the
block count and gives the remainder to the lowest ranks: a 3712-wide
projection (29 blocks) shards at TP2 as 1920 / 1792.

The split runs over the axis's HADAMARD-PADDED width, not its declared
width, because that is the width the checkpoint's trellis actually
carries (see :data:`_HAD_BLOCK`). A ``moe_intermediate_size`` of 1856 is
stored as 1920 = 15 blocks and shards at TP2 as 1024 / 896 in kernel
space, which the declared dim sees as 1024 / 832. The padded width is
derived arithmetically, so it is known before any weight is bound.

EXL3 equivalent of :class:`ColumnParallelLinear`.

Output dim is sharded across the TP group. ``out_features`` is the
GLOBAL output size; this class stores the LOCAL slice and slices
the trellis tensor along dim 1 (output blocks) inside :meth:`exl3_load`.

TP slicing math (single source of truth for this codebase):

  - ``trellis: (in_features // 16, out_features // 16, K*16)``
    → slice dim 1: ``[:, rank * local_out_blocks : (rank+1) * local_out_blocks, :]``
  - ``suh: (in_features,)``  → unchanged (input dim is shared)
  - ``svh: (out_features,)`` → slice ``[rank * local_out : (rank+1) * local_out]``

Constraint: ``out_features % (tp_size * 16) == 0`` so each rank's
local out aligns to a full 16-element trellis block
(:attr:`tp_block_align` = 16). ``__init__`` + ``forward`` come from
:class:`ColumnParallelMixin`.

LoRA + TP: applies BGMV correction in the LOCAL out shard before
the optional all-gather (the LoRA B matrix is sharded along the
same output dim by the LoRA loader — same convention as the dense
:class:`ColumnParallelLinear`).

Multi-GPU end-to-end is unverified on this build (single-GPU box);
the slicing math is unit-tested.

EXL3 equivalent of :class:`RowParallelLinear`.

Input dim is sharded across the TP group. ``in_features`` is the
GLOBAL input size; this class stores the LOCAL slice and slices
the trellis tensor along dim 0 (input blocks) inside :meth:`exl3_load`.

TP slicing math:

  - ``trellis: (in_features // 16, out_features // 16, K*16)``
    → slice dim 0: ``[rank * local_in_blocks : (rank+1) * local_in_blocks, :, :]``
  - ``suh: (in_features,)``  → slice ``[rank * local_in : (rank+1) * local_in]``
  - ``svh: (out_features,)`` → unchanged (output dim is shared)

Constraint: ``in_features % (tp_size * 16) == 0`` so each rank's
local in aligns to a full 16-element trellis block
(:attr:`tp_block_align` = 16). ``__init__`` + ``forward`` come from
:class:`RowParallelMixin`.

LoRA + TP: applies BGMV correction before the all-reduce (the
LoRA correction itself is a partial sum over the sharded input
dim; it composes with the base partial via the same all-reduce).

Multi-GPU end-to-end is unverified on this build; the slicing
math is unit-tested.

EXL3 mirror of :class:`MergedColumnParallelLinear`.

Some EXL3 exports ship a fused projection (e.g. Qwen 3.5/3.6 GDN
``in_proj_qkv``, whose HF tensor stacks ``[Q, K, V]`` along dim 0)
as ONE trellis tensor. At TP=1 the swap binds it as a single
column-parallel linear with ``out_features = sum-of-shards``;
consumers ``.split()`` the output downstream and see correct
per-shard slices.

At TP>1 a naive contiguous output-dim cut lands mid-sub-block when
the sub-block widths differ (``in_proj_qkv`` widths
``[key_dim, key_dim, value_dim]``: rank 0 at TP=2 would get
``[full Q + half K]``), corrupting the per-shard semantics the
GDNBlock relies on (it ``.split()``s ``in_proj_qkv``'s output by
``[key_dim_local, key_dim_local, value_dim_local]``). The fix mirrors
:class:`arbi_serve.weight_quant.awq.linear.AWQMergedColumnParallelLinear`:
slice EACH sub-block independently along the trellis output dim
(dim 1, in 16-element block units) and svh (the per-output scale),
then concatenate the rank-local slices in declared order. ``suh``
(the per-input scale) is shared across all shards — passed through
unchanged.

Per-sub-block constraint: every sub-block's GLOBAL offset and its
per-rank width must be whole :data:`_HAD_BLOCK` blocks, because each
sub-block window is reconstructed and Hadamard-transformed on its own
(see :class:`_EXL3HadAlignedShardMixin`).

A sub-block narrower than one Hadamard block cannot be split at all, so
it is REPLICATED: every rank loads its whole 128-wide window and the
consumer slices its own portion out of the activation. Mamba's ``dt``
(one element per head, 64 wide) is the case this exists for; replicating
64 of 10304 columns costs 0.6% duplicate compute on the projection.

This rank's ``(in, out)`` in KERNEL space, for slab pre-planning.

``header_in`` / ``header_out`` are the checkpoint trellis's GLOBAL
dims, which already include any Hadamard-alignment padding. An
unsharded axis keeps its header dim; a sharded axis reports this
rank's kernel-space shard width, which is what ``exl3_load`` binds.

Validate ``kernel`` (the bound trellis's block-rounded size on
one axis) against ``declared`` (the module's real in/out_features).

Equal is the common case (an unpadded checkpoint) — returned as-is,
every existing EXL3 checkpoint in the repo hits only this branch.
A ``kernel`` that lands exactly on ``declared`` rounded up to the
next :data:`_HAD_BLOCK` multiple is the Hadamard-alignment-padding
pattern (see module docstring) — allowed only at ``tp_size == 1``:
the TP-slicing subclasses (:class:`EXL3ColumnParallelLinear` /
:class:`EXL3RowParallelLinear`) compute per-rank block ranges from
the DECLARED global dim (``self._global_out_features`` /
``self._global_in_features``), which would mis-slice the padding
tail at ``tp_size > 1`` — unverified/unimplemented, so refused
loudly rather than silently mis-sharded. Anything else (kernel <
declared, or a gap inconsistent with rounding declared up to the
next block) is a real mismatch (mis-sharded TP slice, wrong
checkpoint) — raised loudly, same as the pre-padding-support code.

Derive ``_kernel_in_features`` / ``_kernel_out_features`` from
``trellis``'s own block counts and cross-check ``suh`` / ``svh``.

The trellis IS the authority on kernel-space geometry: a checkpoint
carrying Hadamard-alignment padding (see :data:`_HAD_BLOCK`) has
block counts wider than the module's declared dims on one axis, and
the forward path's zero-pad / trim (:meth:`_exl3_forward`) plus the
:class:`exllamav3.LinearEXL3` descriptor (:meth:`_build_inner`) both
key off these two numbers. Deriving them here — rather than at one
binding site — keeps them correct on EVERY path that ends in a
descriptor build: the cold :meth:`exl3_load`, the warm-reload /
post-compaction :meth:`rebind_after_compaction`, and
:meth:`_exl3_forward`'s lazy rebuild.

Shape expectations: ``trellis[in//16, out//16, K*16]``, ``suh[in]``,
``svh[out]`` — in KERNEL space. For column-parallel after slicing,
out is the local out shard; for row-parallel, in is the local in
shard.

Bind the EXL3 tensors and lock-in the kernel descriptor.

``mcg`` / ``mul1`` are optional per-linear codebook seeds — present
only when the checkpoint was quantized with the corresponding
non-default codebook; passed straight through to the kernel.

Subclasses override to apply per-rank slicing before delegating
back here. At TP=1 every subclass calls straight through.

Return a codebook-seed buffer, or ``None`` if it is the empty
"absent" placeholder. The kernel descriptor takes ``None`` for the
default codebook; an empty buffer (the unbound / no-seed state, or
a warm-reloaded "absent" slot) maps to that.

Drop the inner kernel host BEFORE the compactor allocates slabs.

:meth:`_build_inner` constructs an ``exllamav3.LinearEXL3`` — a
PLAIN object, NOT an nn.Module, so the compactor's persistent-
tensor walk never visits it — holding its OWN refs to
``trellis / suh / svh`` plus a ``BC_LinearEXL3`` C++ descriptor
bound to them. That secondary reference pins the ORIGINAL GPU
storage; left in place it co-resides with the slab while the slab
is allocated and filled, spiking the LOAD GPU storage owned solely by the model buffers).

The compactor calls this hook ONCE, BEFORE any slab is allocated
(:func:`arbi_serve.loader.flat_loader._run_compaction_release_hooks`).
We tear the inner down so the buffers are owned only by the module;
the host-staging guard can then free the The inner is rebuilt over the
final slab views by :meth:`rebind_after_compaction`. Idempotent;
``_loaded`` records that a rebuild is still owed so the rebind hook
fires even though ``_inner`` is now ``None``.

Dropping ``self._inner`` alone is NOT enough: the process-global
``_INNER_REGISTRY`` (custom-op dispatch table) also holds a STRONG
reference to the inner keyed by ``_exl3_uid``, and that reference
keeps the inner — and its ``trellis / suh / svh`` Deregister the uid here too; :meth:`rebind_after_compaction`
re-registers the rebuilt inner over the slab views.

Re-derive the inner kernel host against the (now slab-backed)
buffers, AFTER the compactor has rebound every buffer to its slab
view.

The compactor rebinds ``self.trellis`` / ``self.suh`` / ``self.svh``
to views into a packed weights-pool slab. The inner (torn down by
:meth:`release_secondary_refs_for_compaction`, or stale from load if
no release ran) must be rebuilt so the forward path reads off the
slab and the original storage is fully unreferenced.

:meth:`_build_inner` reads ``self.trellis`` / ``self.suh`` /
``self.svh`` and constructs a fresh descriptor over the slab views,
re-publishing it against the same ``_exl3_uid`` (the registry
overwrites).

Keys on whether the weight buffers are populated rather than on
``_inner is not None`` so it rebuilds after the release hook set
``_inner = None`` AND after a WARM RELOAD, where the meta graph
never called ``exl3_load`` (so ``_loaded`` is still ``False``) but
the flat fill DMA'd the persistent ``trellis`` buffer into place.
A no-op if no buffers were ever bound (still the empty placeholder).

Dequantize to the dense ``(out_features, in_features)`` weight in
``F.linear`` layout (row-major over output).

Mirrors :meth:`AWQLinearBase.dequantized_weight` — the contract the
GDN block's fused dense ``[B|A]`` prebuild relies on
(:meth:`GDNBlock.prepare_fused_ba_dequant`) when a checkpoint
quantizes the sub-tile ``in_proj_b`` / ``in_proj_a`` projections.
Known EXL3 GDN exports ship b/a dense, so this is a latent-path
guarantee, exercised once at build time — never a hot path.

Routes through exllamav3's full trellis reconstruct
(``LinearEXL3.get_weight_tensor`` → ``ext.reconstruct`` + the
pre-applied Hadamard/sign transforms), i.e. the same dense weight
the trellis GEMM computes against. Native dtype is fp16 (the EXL3
kernel dtype); ``dtype`` casts on top.

Bind the source shard's tensors, or nothing on an empty window.

An empty window has no columns to compute, so there is no kernel
descriptor to build over and no GEMM to run — :meth:`local_logits`
never reaches one. Binding a zero-width descriptor would be a kernel
whose only correct behaviour is to not be called.

This rank's window of the row, un-gathered — ``(..., window_r)``.

The vocab-shard logits surface the drafter's distributed argmax reads
(:meth:`~arbi_serve.weight_quant.base.ColumnParallelMixin.local_logits`).
Global column ``i`` of this rank's output is token id
``self._tp_out_start + i``.

The composed global window ``(..., N)``, identical on every rank.

A LoRA state is REFUSED rather than dropped. This module exists only
as a drafter head over ``model.lm_head``, which carries no LoRA target
in any build that reaches here — but a correction silently discarded
would make the drafter read a different projection from the verifier
with nothing to say so.

Per sub-block ``(global_start, kernel_width, declared_width)``.

``kernel_width`` exceeds ``declared_width`` only for a replicated
sub-block padded up to a whole Hadamard block, and only the LAST
sub-block may do so — the forward's trim to ``out_features`` is a
leading slice, which is correct only when the slack is trailing.

Per-sub-block + per-rank slice of the fused trellis / svh along
the OUTPUT axis, then bind via the base (TP=1) loader.

Trellis layout ``(in//16, out//16, K*16)`` — output blocks live on
dim 1, so a sub-block at output offset ``sub_offset`` of width
``sub_dim`` occupies trellis blocks
``[sub_offset//16, (sub_offset+sub_dim)//16)``; this rank takes its
``tp_shard_range`` chunk within that. ``svh`` (the per-output
scale, ``(out,)``) slices on the same output offsets in element
units. ``suh`` (per-input, shared) passes through unchanged. mcg /
mul1 are scalar codebook seeds — shared, passed through.

The one definition of mcg's int8 weight grid. Everything else imports it.

It ships inside ``arbi_serve`` because the served extension loader compiles the
kernel with the four ``-DEXL3_I8_MCG_*`` constants read from here, and ``tools``
is not present in the serving image.

``mcg`` is exl3's DEFAULT codebook -- ``scripts/exl3_convert_mamba_hybrid.py``
spells it ``--codebook choices=("mcg","mul1") default="mcg"`` -- and it is the
codebook of every checkpoint this repo owns. ``mul1`` is the legacy one. The
int8 leg (#1861) shipped decoding ``mul1`` alone, so it supported the codebook
we do not use.

WHY AN EXACT INTEGER DECODE EXISTS
----------------------------------
exl3's ``decode_3inst<1>`` is three instructions::

    x = idx16 * 0xCBAC1FED
    x = (x & 0x8FFF8FFF) ^ 0x3B603B60  # one lop3, LUT 0x6A
    v = __hadd(low_half(x), high_half(x))

The lop3 leaves each 16-bit half with its sign bit, ``0b011`` in the top three
exponent bits, and the low twelve bits taken from ``x``. So the exponent field
takes exactly FOUR values, 12..15, and each half is

    (-1)^s * (1024 + M) * 2^(E - 25),   E in {12,13,14,15}, M in [0, 1024)

-- an 11-bit integer times a power of two. Scaled by 2^13 both halves, and
therefore their exact sum, are INTEGERS. That is what makes an exact integer
decode possible; the kernel's comment used to claim mcg "has no analogue" of
mul1's integer decode, and that false sentence closed off the drafter for the
whole effort (#1861).

Dearer than mul1's single ``__dp4a``, though, and this module prices both forms
over the full 65536-index domain rather than picking on preference:

===========================  ==========  ========================================
form                          rel-MSE     what it costs
===========================  ==========  ========================================
fp16 decode + magic FMA       5.2668e-05  exl3's own 3 instructions + 1 ``hfma2``
integer, ``(V + 128) >> 8``   5.2668e-05  per-half sign/exponent/mantissa split
                                          and a VARIABLE shift, both halves
===========================  ==========  ========================================

**They tie on accuracy.** They are not the same map -- they disagree on 1028 of
65536 indices -- but every one of those is an exact half-way case, where
round-half-even and round-half-up cost the same squared error. So the choice is
made on instruction count, not on numerics, and the fp16 form wins there: it
reuses the decode exl3 already emits and adds one ``hfma2`` per two weights,
while the integer form has to take each half apart.

For scale, mul1's shipped grid scores 5.9745e-05 on the same measure. Read that
correctly: it says int8 quantisation costs mcg **no more than** it costs mul1,
so no accuracy law blocks an mcg leg. It does NOT say mcg is a better codebook,
and it is a decode-grid rel-MSE over the codebook domain, not an end-to-end
distribution result.

THE MAGIC FMA, AND WHY THE CLAMP IS DEAD HERE TOO
-------------------------------------------------
``__hfma(v, 32.0, 1152.0)`` lands in [1025.6, 1278.4] over the WHOLE domain,
which is inside the fp16 binade [1024, 2048) where the ulp is exactly 1. So the
fma's own round-to-nearest IS the quantisation, the result is an exact integer,
and its low byte is ``q + 128``: the int8 weight is that byte XOR 0x80, with no
conversion instruction and no saturation. ``q`` spans [-126, 126] and never
reaches an int8 bound, so the clamp is dead for mcg exactly as it is for mul1 --
proved over the full domain in ``tests/test_exl3_mcg_int8_grid.py``, not
asserted here.

``W_SCALE`` is DERIVED below rather than typed: it is the finest power-of-two
step whose grid still holds the codebook inside int8.

``2^13 * (low + high)`` as EXACT int64, before the fp16 add rounds.

Each half is ``(-1)^s * (1024 + M) * 2^(E-25)`` with ``E`` in 12..15, so
``2^13`` times it is ``(-1)^s * (1024 + M) << (E - 12)``: an integer. This
is the quantity the pure-integer decode form works on, and it is NOT the
same number as ``codebook_fp16`` -- the fp16 add rounds, this does not.

The finest power-of-two step that keeps the whole codebook inside int8.

Written as a search over the domain rather than as the literal 0.03125, so
a codebook constant that moved would move this with it instead of silently
saturating the served weight.

Form 1: ``__hfma(v, 32.0, 1152.0)``, then take the low byte.

The fma's round-to-nearest-even is the quantisation, so this is exactly
``rint(v / W_SCALE)`` -- and therefore the NEAREST int8 grid point to the
fp16 codebook value, which is the least squared error any decode can reach
on this grid.

EXL3 grouped multi-expert GEMM as a cudagraph-safe custom op.

``exllamav3_ext.exl3_mgemm`` runs one trellis-dequant GEMM across a set of
experts, selecting each row's expert from a device-side index tensor. It is
the sparse MoE analogue of the single-linear GEMM wrapped in
:mod:`arbi_serve.weight_quant.exl3.custom_op`.

``min_index``/``max_index`` give it an EXPERT RANGE. Selections outside
``[min_index, max_index)`` are skipped and the rest are rebased by
``min_index``, so ``B``/``suh``/``svh`` may be a pointer table local to one
expert shard. Slots keep their position: a skipped slot's output row is
simply never written, which is why ``out`` is allocated zeroed here.

Upstream reaches this op only through ``BC_BlockSparseMLP.run_bszN``,
whose captured wrapper is bounded by ``MAX_BSZN = 8``. The op itself takes
``num_tokens`` and per-token index/weight tensors, so the bound belongs to
that wrapper, not to the kernel. Driving the op directly lifts it.

Whether this process serves the routed combine reproducibly.

Imported lazily: this module is imported during weight load, before the
runtime-flag registry is necessarily populated, and a hard import at module
scope would freeze the answer at import time rather than read it per call.

Device-pointer view over one projection across all local experts.

Mirrors ``exllamav3.modules.MultiLinear``: the kernel takes the trellis
and scale buffers as pointer tensors, so the group holds those plus the
shared codebook constants every expert must agree on.

Register ``arbi_serve::exl3_moe_fused`` (idempotent).

Wraps ``exllamav3_ext.exl3_moe``, the EXPERT-major fused MoE kernel: it
dequantises each expert once and streams that expert's tokens through it,
where :func:`exl3_moe_gemm` re-reads an expert's weight per (token, expert)
slot. The whole block — gate/up, activation, down, and the routed combine —
is one launch.

``exl3_moe`` takes ``num_active`` but NO expert range, so an
expert-sharded caller must hand it routing already rebased to the local
pointer table with off-shard slots neutralised — see
:meth:`~arbi_serve.models.moe_grouped_exl3.GroupedRoutedExperts._masked_local`.

Grouped trellis GEMM: row ``i`` runs through expert ``indices[i]``.

``x`` is ``(S, in_features)`` fp16 and ``indices`` is ``(S,)`` — one
SLOT per row, already expanded by the caller. The up stage passes
token rows replicated ``top_k`` times; the down stage passes each
slot's own activation. ``weights`` is an optional ``(S,)`` scale
applied in-kernel.

``min_index``/``max_index`` are an expert range in the frame
``indices`` is expressed in; ``(-1, -1)`` disables filtering and
``indices`` are then local to the group's pointer table. Rows whose
expert falls outside the range come back zero.

See ``docs/exl3-mgemm-binding-contract.md``: indices go in flat and
int64, ``A`` carries one row per slot, ``A_had`` is sized for every
slot, and ``svh`` gives a PADDED width that is trimmed here.

Routed expert mixture for ``x``, expert-major.

``indices`` are LOCAL expert ids ``(rows, top_k)``; ``weights`` their
routing weights. Returns ``(rows, hidden_size)`` fp32.

The expert-sorted layout is built here with device-side ops only, and
the per-expert scratch is bounded by ``rows``: ``_route`` picks
``top_k`` DISTINCT experts per row, so no expert can be handed more
than one slot per row. Both properties are what make this capture-safe
— no host sync, no shape that depends on the routing.

The one definition of mul1's int8 weight grid. Everything else imports it.

The constant below was, until this module existed, written out four times --
in the kernel's ``dp4a`` addend, in ``gate1_decode``'s Python mirror, in
``qref``'s reference table, and in prose in three docstrings. That is the shape
of defect this codebase hits most often: a correction lands on one copy, the
copies that validate it move with it, and the gate still passes.

WHAT THE CONSTANT IS, AND WHY IT IS 1 AND NOT 2
-----------------------------------------------
``mul1``'s decode is ``__hfma(h, k_inv, k_bias)`` with ``h = 1024 + bytesum``,
``k_inv = half(0x1eee)`` and ``k_bias = half(0xc931) = -10.3828125``. It is an
affine function of a byte sum, which is what makes ONE ``__dp4a`` enough for the
integer int8 decode. That cheapness is mul1's and is why this module exists; it
is NOT a claim about the other codebooks. This paragraph used to say ``mcg``
"admits none", and that is FALSE -- ``mcg``'s LOP3 output is an fp16 bit pattern
whose exponent field takes exactly four values, so each half is
``(-1)^s * (1024 + M) * 2^(E-25)``, an 11-bit integer times a power of two, and
an exact integer decode follows. It is dearer than one ``__dp4a``, not
impossible, and the false version of this sentence closed off the drafter for
the whole effort (#1861). Possible is not exact:

* ``k_bias`` is NOT ``(-1024 - 510) * k_inv = -10.380996704``. The fp16 binade
  [8, 16) has ulp 0.0078125 and the neighbours are -10.375 and -10.3828125, so
  ``k_bias`` is the CORRECTLY ROUNDED value -- there is no better fp16. The
  codebook therefore sits **0.268320 k_inv = 0.067080 int8 steps** below the
  ideal lattice ``(bytesum - 510) * k_inv``, on every one of its 913 values.
* The integer decode ``sat((bytesum - 510 + c) >> 2)`` targets that IDEAL
  lattice. Against the ideal, ``c = 1`` and ``c = 2`` are exactly TIED
  (mean -0.125 and +0.125 of a step, identical MSE) -- the choice is free, and
  ``c = 2`` is the natural round-half-up. Against the REAL codebook the tie
  breaks, because the 0.067-step offset pushes ``c = 2`` away from centre and
  ``c = 1`` towards it.

Over the full 65536-index domain, against the real fp16 codebook:

===  ============  ================  ==============
 c     rel-MSE      bias (int8 step)  vs shipped c=2
===  ============  ================  ==============
 0    1.2668e-04       -0.307918          +50.4%
 1    5.9745e-05       -0.058071          -29.1%
 2    8.4244e-05       +0.191929            0.0%
 3    2.0024e-04       +0.441914         +137.7%
===  ============  ================  ==============

``c = 2 -> c = 1`` changes 16384 of 65536 indices, costs nothing (same
instruction, same shift, same addend width) and is worth -29.1% rel-MSE and a
3.3x smaller bias.

THE REMAINING +-0.058 IS A FLOOR, NOT A TARGET. ``bytesum`` is an integer, so
every finer-resolution rewrite (``(4bs - 2033) >> 4``, ``(8bs - 4066) >> 5``,
``(16bs - 8133) >> 6``) collapses onto the same partition and gives the
IDENTICAL rel-MSE 5.9745e-05. The tie set is a quarter of the domain and moves
as a block, so the residual cannot be centred away. Do not re-attempt it.

``test_exl3_mul1_int8_grid.py`` DERIVES this constant from the codebook rather
than restating it, so the value cannot drift back.

It ships inside ``arbi_serve`` because the served extension loader compiles the
kernel with ``-DEXL3_I8_DP4A_ADDEND=`` read from :data:`DP4A_ADDEND`, and
``tools`` is not present in the serving image.

``__hfma(h, k_inv, k_bias)``: one rounding, exactly as the device does it.

``h`` and ``k_inv`` are exactly representable and the product needs 22
significant bits, so float64 carries product and sum exactly and the single
cast to fp16 IS the fma's rounding.

Measure-once persistence for the EXL3 trellis-GEMM kernel-shape pin.

WHY THIS EXISTS. :func:`~arbi_serve.weight_quant.exl3.custom_op.
resolve_kernel_shape_pins` picks one ``(TILESIZE_K, TILESIZE_N)`` family per
EXL3 geometry by TIMING every legal candidate on the live card. The winner is a
property of the card and the kernel set, not of the boot — but a timing is a
measurement, and where two families sit inside each other's noise the winner
FLIPS from one boot to the next at an unchanged configuration. That has two
costs, and the second is the one that hurts:

  * the sweep is re-paid on every boot for an answer that did not change; and
  * WHICH FAMILY THE BOOT DREW is an uncontrolled variable underneath every
    cross-boot comparison, because the family sets the decode kernel. A paired
    A/B that boots one arm and then the other cannot tell a real effect from a
    pin flip, and MTP acceptance grouped by the pin rather than by the arm is
    what that looks like from the outside.

So the choice is MEASURED ONCE per configuration and read back afterwards. The
write mechanism is the budget cache's, not a second one: same
tempfile→fsync→rename write, same version field, same refuse-on-mismatch rule
— see :mod:`arbi_serve.engine.memory_budget.graph_pool`, whose
``_atomic_write_json`` this module calls rather than restates.

WHERE IT LIVES, AND WHY NOT WITH THE BUDGET CACHE. Its OWN directory
(``ARBI_SERVE_EXL3_PIN_CACHE_DIR``, else ``~/.cache/arbi-serve/exl3-kernel-pin``
— the home cache root every shipped compose file persists on a named volume,
so the pin survives the container without a deployment change). The pin is a
property of the CARD, the driver and the kernel library; a memory budget is not
one of its inputs and does not appear in its key, so it must not be one of its
LOCATIONS either. A cache directory is an ISOLATION LEVER: giving each arm of an
A/B its own budget-cache dir is the ordinary way to isolate the budget
measurement, and a pin that lives in that dir is split by it — each arm then
times and locks a different kernel family, and the comparison reads the family
rather than the change. That artefact is stable and reproducible per arm, so it
does not read as noise. A deliberate re-sweep is a deletion, and the log names
the file.

WHAT IS PERSISTED, AND AT WHICH GRANULARITY. One file per KEY; inside it one
record per GEOMETRY (``"<in>x<out>xK<bits>"``). The geometry is the model's
quantisation shape, and it belongs INSIDE the file rather than in the key
because it is exactly what makes two records different readings rather than two
readings of the same thing: a boot that binds a geometry the file does not carry
sweeps that geometry and leaves the rest alone, which is what lets the second
pin seam (a drafter checkpoint attaching its own linears) add to an entry
instead of invalidating it.

WHAT THE KEY CARRIES. Everything that changes what the timing is a reading OF:

  * ``device_uuid`` — the physical card. The reading is a stopwatch on silicon,
    and per-card is also what keeps two ranks of a TP boot, sweeping at the same
    moment, from writing one file over each other: distinct cards, distinct
    keys, distinct files.
  * ``device_name`` / ``device_capability`` / ``device_sms`` — the card's model,
    compute capability and SM count. Provenance an investigator diffs, and the
    properties that decide whether a reading could ever have transferred.
  * ``driver_version`` — the driver JITs and launches; a driver upgrade is a new
    reading. Read through NVML, and an unreadable driver keys itself (see
    :func:`_driver_version`) rather than collapsing onto a real version string.
  * ``torch_version`` / ``cuda_runtime`` — the kernel-library identity, the same
    pair ``_graph_pool_cache_key`` keys on; ``torch_version`` carries the CUDA
    build tag and ``cuda_runtime`` names the toolkit outright.
  * ``exllamav3_version`` / ``exllamav3_rev`` — the fork holds the EXL3 GEMM
    kernel map, and two different pins both report one dist version, so the rev
    stamp is the identity and the version is provenance beside it.
  * ``arbi_serve_rev`` — our own build. Which candidates are legal, how the
    ladder is walked and what the timing loop does are all our code.
  * ``shape_table`` — a digest of ``{shape_idx: (TILESIZE_M, K, N)}``. The
    candidate SET itself, and the one input that is exact in EVERY context.
    Both revision stamps are ``/etc`` files the build writes, so they name the
    IMAGE: absent on a host checkout, and under a bind-mounted tree they name
    the image rather than the code actually running. That is the same trade
    ``_read_identity_file`` and ``_driver_modules_cache_key`` already make — the
    guard is for images, where staleness ships — and the table digest is what
    keeps a shape-table move from ever being reused across regardless.
  * ``select_rows`` — the row counts the families are timed over. Change which
    rows the pin has to serve and a different family wins them.
  * ``tp_size`` — under TP the peer rank sweeps its own card at the same moment
    in the same chassis, sharing power and thermal headroom; a tp=1 reading was
    taken without that neighbour.

WHAT IS NOT IN IT, AND WHY NOT. The model id: the pin is a property of the
GEOMETRY, and two checkpoints that carry the same geometry on the same card want
the same answer — keying on the model would re-sweep for no reading. Thermal
state: it moves the timing but it is not a configuration, so keying on it would
re-sweep every boot and give back exactly the boot-to-boot flip this module
exists to remove.

THERE IS NO RE-SWEEP CONDITION AND NO TIMER. Every input that changes the
answer is in the key, and every input that merely changes the READING — how hot
the card was, what else was resident — is deliberately not, because re-sweeping
on those is the flip. What this does mean is that the first sweep's draw is the
one the configuration keeps, worse draw included: persistence makes the choice
STABLE, it does not make it BEST. Choosing the best would take an instrument the
selection does not have (it minimises GEMM time, and the families differ in
their arithmetic as well as their speed), and that instrument needs stable boots
to run at all — which is this. Deleting the entry is the re-sweep; the log names
the file.

The NVIDIA driver version string, or :data:`_UNREADABLE`.

NVML rather than a CUDA-runtime integer because the driver is what the
string names and what an operator upgrades. A box where NVML will not load
still gets a stable key — every boot on it reads the same sentinel — it
just cannot say which driver the reading was taken under, and it re-sweeps
once if NVML later starts answering.

``nvmlInit`` / ``nvmlShutdown`` are reference-counted per process, so this
balanced pair cannot pull the library out from under
:mod:`arbi_serve.server.nvml_metrics`, which holds device handles for the
life of a served process.

The dict :func:`shape_pin_cache_key` hashes. See this module's docstring.

Every value is JSON-serialisable so the digest reproduces byte-for-byte
across boots at one configuration, which is the whole property the cache
rests on.

Pick the pin-cache directory. Its own ladder, keyed on its own inputs.

Resolution order:
  1. Explicit ``build_dir`` argument (tests, and a deliberate re-sweep).
  2. ``ARBI_SERVE_EXL3_PIN_CACHE_DIR``.
  3. :data:`_PIN_CACHE_DIR_FALLBACK`.

``ARBI_SERVE_BUDGET_CACHE_DIR`` is deliberately NOT in that ladder — see
this module's docstring for what reading it costs.

``{geometry: record}`` established for this configuration, or ``None``.

REFUSES rather than adapts. A file whose ``version`` is not this module's,
or whose recorded ``cache_key_components`` are not byte-identical to
``cache_key_inputs``, is IGNORED and the caller re-sweeps. The components
check is not redundant with the filename: the name is a TRUNCATED digest,
and a cache directory is a thing people copy between hosts and volumes — a
pin timed on other hardware, used because two short hashes agreed, is worse
than paying for the sweep.

Record the resolved pin for this configuration. Returns whether it wrote.

Called by the sweep, from the sweep's own completion — never by a scrape, a
metrics callback or an admin route. The engine owns the reading; telemetry
reads it back through
:func:`~arbi_serve.weight_quant.exl3.custom_op.kernel_shape_pins`.

Writes the WHOLE resolved map, restored geometries included, so the second
pin seam's superset replaces the first seam's subset in one file rather than
accumulating partial entries. Not monotonic and not a high-water: unlike a
budget bound there is no direction in which erring is safe, and the value is
a CHOICE — the point is that it stops moving.

``int8_shapes`` is ``{geometry: {class: shape id}}`` for the int8 leg, and
it is written as a RECEIPT rather than restored as an input. Both of that
leg's sources — the shipped table and the package default — are already
deterministic, so reading a shape back from a per-host file could only
introduce the host-local variation the pin exists to remove; what the
record buys is that a boot's per-class tile is on disk beside the family it
ran with, so two deployments can be COMPARED without re-deriving what each
would have chosen.

Say so when the DEFAULT location holds a different family for this key.

A boot only reaches here by having TIMED something, and a boot writing
somewhere other than :func:`_default_pin_cache_dir` was pointed there — by
``ARBI_SERVE_EXL3_PIN_CACHE_DIR`` or by a caller's ``build_dir``. The same
key means the same card, driver, kernel library and shape table, so two
locations naming different families are not two configurations: they are
one configuration whose kernel choice depends on which directory the boot
was aimed at. That is exactly the variable a cross-boot comparison has to
hold still, so it is NAMED — both families, both paths — rather than left
to be inferred from throughput.

Reads only, and never raises: it is a receipt on a record already on disk.
Silent when the two agree, so relocating the cache on purpose does not warn
every boot for doing nothing wrong.

One geometry's recorded launch decisions. All frozen inputs to a launch.

``int8_shape`` joins ``family`` and ``num_sms`` rather than living in a
parallel table, for the reason the width joined the family: they are read
together, gated together, and a decision restored without the configuration
it was recorded under is a number about a kernel this boot may not run.

They are NOT the same KIND of decision, and that is recorded here rather
than left to be inferred: ``family`` is RACED at boot when nothing covers
the geometry, while ``num_sms`` and the two int8 shapes are never raced.
Each of those has a deterministic default that a deployment already serves,
so racing them would introduce a boot-to-boot numerics variable where none
existed — #1879's defect with the sign flipped.

``int8_verify_shape`` is the int8 leg's tile for the VERIFY class, and it
is a separate entry because it answers a separate question: a verify slate
fills a fraction of one m-tile where a prefill chunk fills many, so the two
classes want different TILE_M off one reduction order. One field per class,
so a geometry can move one without restating the other.

``{geometry: ShippedPin}`` the shipped table records for THIS config.

FOUR VALUES, ONE ENTRY, and they are not separable. ``num_sms`` is the
``force_num_sms`` the geometry's linears pass to ``exl3_gemm``; ``0`` means
the kernel's own grid choice, which is what every deployment has run since
the pin existed. The width's WINNING VALUE is family-dependent -- measured
on ``5120x1024``, the same win survives all four legal families while the
best width moves between them -- so a width read without the family it was
raced under is a number about a kernel the boot may not be running. They
travel together for that reason and are never merged from two sources.

Empty when nothing matches, which is the fail-safe direction: an unmatched
configuration is timed exactly as before rather than served an answer that
was chosen for a different card or a different candidate set.

MATCHING IS EXACT ON EVERY KEY FIELD the entry declares. ``device_name``
because the reading is a stopwatch on one card model and compute capability
alone spans parts with different SM counts and clocks; ``shape_table``
because it pins the candidate menu, so a fork bump that adds a tile stops
matching instead of applying an answer chosen without it; ``select_rows``
because changing which rows the pin must serve changes which family wins
them.

Deliberately NOT keyed on driver, torch or the model id. This table is a
DECISION about which numerics a deployment serves, not a prediction about
what this machine would time — and a decision does not go stale when the
driver moves. Keying it like a measurement would make it un-hittable, which
is how a recorded default quietly becomes dead code.

An unreadable or malformed TABLE is reported and treated as empty, because
failing a boot over a performance-neutral default would be a worse outcome
than timing the pin. A malformed value for the escape-hatch FLAG is a
different thing and fails fast, exactly as every other boolean flag does —
the operator asked for something the process cannot honour.

An entry's ``pins`` value is either a bare family string — every entry
written before widths existed, and every geometry whose width was measured
and REFUSED — or a dict carrying ``family``, ``num_sms`` and the per-class
``int8_shape`` / ``int8_verify_shape``. Any field the dict omits reads as
unrecorded, which is exactly what those deployments ran, so an older table
is restored unchanged rather than being treated as unmatched.

The shipped table entry for THIS configuration, or ``None``.

One matcher for every field an entry carries, so a per-geometry decision
and an entry-wide one can never be read out of two different entries.

The recorded VERIFY tile ladder for this configuration, as recorded.

ENTRY-WIDE, not per geometry, and that is the shape of the decision rather
than a convenience. The band EDGES are where the int8 leg and the trellis
leg cross as the slate widens, so they are a property of the row rungs a
deployment captures; the TILE that wins inside a band can still be a
geometry's own, and ``int8_verify_shape`` is where a geometry says so. One
edge vector per boot is also what lets the band be resolved once per
forward, from the width the forward declared, instead of per linear.

Returned as recorded rather than parsed here: the vocabulary belongs to
:func:`~arbi_serve.weight_quant.exl3.int8_kernel.normalise_int8_verify_bands`,
which is also what validates it, so this module does not acquire a second
opinion about what a ladder is.

A :class:`ShippedPin` from either recorded form; a bad field reads as its default.

Never raises: this runs on the boot path over a file an operator may have
hand-edited, and the fail-safe direction is what the deployment already
ran. A width or an int8 shape that cannot be parsed is dropped, not guessed
at, and the family beside it is still honoured — they are recorded
together but only the family is load-bearing for correctness.

An unparseable int8 shape reads as UNRECORDED, never as
:data:`INT8_SHAPE_DECLINED`: the leg then takes its own default, which is
the shape every receipt was measured against. A typo in a performance table
must not silently turn a leg off, and it must not silently turn one on
either — which is why the default is a shape and the decline is explicit.

The ``(family, ladder, scores, num_sms)`` a record names — iff it applies.

``legal`` is the ladder THIS build computes for each family at this
geometry, so the check is an equality against a freshly derived value rather
than a plausibility test: the family must still exist for this geometry AND
its ladder must be exactly the one recorded. Anything else — a family the
shape table no longer offers, a ladder whose membership moved, a malformed
record — returns ``None`` and the caller sweeps.

A ladder is never REPAIRED to the nearest legal one. A repaired pin is a
choice nobody measured, wearing the receipt of one that was.

Import-time shim for ``exllamav3`` weight-quantization integration.

We pull in :class:`exllamav3.modules.quant.exl3.LinearEXL3` and the
underlying ``exllamav3_ext`` CUDA extension, but we do **not** want the
rest of exllamav3 (its generator, attention, tokenizer, grammar
filters) — every one of those is fronted by another inference engine
that arbi-serve already owns.

The trouble: importing any submodule of ``exllamav3`` triggers
``exllamav3/__init__.py``, which eagerly imports the generator chain.
That chain imports ``exllamav3.generator.filter.formatron``, which
type-annotates a class signature with ``FormatterBuilder`` from
``formatron.formatter``. Annotation evaluation is **not** guarded by
the file's local ``try/except``, so ``FormatterBuilder`` (and friends)
must be name-resolvable at class-definition time.

Two upstream issues compounded this on exllamav3 0.0.x:
  1. ``formatron`` 0.5.0 (the only version satisfying exllamav3's
     ``>=0.5.0`` floor) is incompatible with ``pydantic>=2.11`` — its
     ``schemas/dict_inference.py`` references the removed
     ``pydantic.typing.Type``.
  2. ``flash_attn`` is imported at module scope in
     ``exllamav3/modules/attn.py``; we don't use the attention surface
     at all (arbi-serve has its own attention backends), but the
     namespace package on this machine ships without symbols.

We sidestep both by installing minimal ``sys.modules`` stubs **before**
the exllamav3 package init runs. The stubbed names only need to satisfy
class-body annotation lookup; nothing actually calls into them because
``formatron_available = False`` short-circuits the relevant code paths.

NEITHER issue applies at the ``EXLLAMAV3_MIN_VERSION`` floor: 1.0.0 dropped
flash-attn/xformers outright, and 1.4.1 replaced formatron+kbnf with
llguidance and defers every one of those imports into
``generator/filter/formatron.py::_load_formatron()`` (module globals are
``None`` at class-definition time, and the file carries
``from __future__ import annotations``). The stubs are therefore
belt-and-braces on 1.x, not load-bearing — they are kept because they
cost one dict write each and are restored immediately, and because they
keep this surface importable if a transitive re-introduces either name.

Because of these stubs, exllamav3 is the ONLY third-party runtime
dependency of the EXL3 weights surface: with ONLY exllamav3 installed --
none of formatron / kbnf / general-sam / frozendict / marisa-trie /
xformers present -- ``LinearEXL3`` imports, loads, runs forward and
dequantizes correctly. Hence the ``[exl3]`` extra in pyproject is empty
and exllamav3 is installed ``--no-deps`` (see :data:`_INSTALL_HINT`).

Deferred import. Importing this module is cheap and side-effect-free:
the stub install and the ``exllamav3`` import (which JIT-compiles the
``exllamav3_ext`` CUDA extension on first touch) are deferred until a
caller actually resolves :data:`LinearEXL3`. Resolution happens lazily
through module ``__getattr__`` and is cached, so the first attribute
access does the (one-time) compile and every subsequent access is a
dict lookup. This keeps mere ``import``/collection of the EXL3 surface
free of any CUDA build — only an actual checkpoint load
(``EXL3LinearBase._build_inner`` → ``shim.LinearEXL3``) pays for the
kernel. Production already defers the shim import to ``_build_inner``,
so this is consistent with the existing design.

Pop our stubs out of sys.modules so other code paths see the
real packages (or the bare namespaces). For ``flash_attn`` this is
critical: leaving the stub in place would feed a fake
``flash_attn_varlen_func`` class into arbi-serve's flash-attn varlen
resolver, which imports the symbol lazily at engine build time.

Slots one ``exl3_mgemm`` call may filter by expert range, else ``None``.

``None`` means the installed extension carries no expert-range slot list
at all: its filtering path renumbers slots into a fixed 128-entry device
array with no bound check, which mis-maps every slot above the first
off-range one and writes out of bounds past 128 slots.

Read this, never ``__version__``: the pinned revision and the PyPI wheel
both report ``1.4.2``.

Pull ``libcudart`` into the process-global symbol namespace.

Thin wrapper over :func:`arbi_serve._libcudart.ensure_global_libcudart`
(single source of truth) so the EXL3 kernel and an in-process TileLang
resolve against the same runtime regardless of load order.

Clear a stale torch FileBaton lock for ``exllamav3_ext``.

Delegates to the shared :func:`clear_stale_extension_lock` — the same
guard now also runs globally for every JIT extension via the
``cpp_extension.load`` wrapper in :mod:`arbi_serve._prebake_loader`.
Kept as belt-and-suspenders for the explicit pre-import call here.

Install the stubs, import ``exllamav3``, and return ``LinearEXL3``.

Deferred so that merely importing this module (or collecting the EXL3
test surface) does not JIT-compile the ``exllamav3_ext`` CUDA
extension. First call does the (one-time) compile and caches the
class; subsequent calls return the cached class.

Lazily resolve :data:`LinearEXL3` on first access.

PEP 562 module ``__getattr__`` — only invoked for names not found as
real module globals, so the cheap constants (``_INSTALL_HINT``, the
stub helpers) resolve normally without triggering the compile.

FP8 weight-quantization backend for arbi-serve.

Importing this subpackage registers the backend with the global
quant-backend registry. Pure PyTorch — no third-party dep, no JIT
kernel; the dequant is a single broadcast multiply and the matmul
runs through the standard ``F.linear`` in the activation dtype.

The backend handles three compressed-tensors / native FP8 layouts:

  - **per-tensor**:  ``weight_scale`` is a scalar (shape ``()`` or ``(1,)``).
  - **per-channel**: ``weight_scale`` has shape ``(out_features, 1)`` —
    one scale per output row.
  - **block-wise**:  ``weight_scale_inv`` has shape
    ``(ceil(out/128), ceil(in/128))`` — one fp32 scale per 128x128 tile
    (DeepSeek-V3 / Qwen3.6 native FP8). Dequant broadcasts each tile's
    scalar over its tile; the block size is read from
    ``quantization_config.weight_block_size`` (default 128x128).

On-disk shapes::

    weight            float8_e4m3fn   [out, in]
    weight_scale      bf16/fp16/fp32  [out, 1]  or  [1] / scalar
    weight_scale_inv  fp32            [ceil(out/128), ceil(in/128)]

Detection key is a ``*.weight`` tensor with dtype ``float8_e4m3fn``
(or ``float8_e5m2``) paired with a matching ``*.weight_scale`` (per-
tensor / per-channel) or ``*.weight_scale_inv`` (block-wise). This
covers the canonical compressed-tensors ``float-quantized`` layout
(llm-compressor) and the DeepSeek-V3 / Qwen3.6 native block-scaled FP8.

Activation-side fp8 ("dynamic" variant) is **transparent** to this
backend: weights are dequantized to bf16/fp16 at the layer boundary
and the matmul runs in the activation dtype.

FP8 :class:`QuantBackend` implementation.

The class implements the protocol — detection from fp8-typed weight
tensors paired with matching ``weight_scale``, mapping HF-canonical
safetensors paths to the FP8Linear mirror class, and reading the
(weight, weight_scale) pair per linear.

FP8-specific assumptions encoded here:

  - Detection: any ``*.weight`` tensor whose dtype is ``float8_e4m3fn``
    or ``float8_e5m2`` AND has a matching ``*.weight_scale`` companion.
    This identifies the canonical compressed-tensors fp8 layout used
    by llm-compressor and the bulk of public fp8 packs.
  - Tensor schema per linear: ``{weight (fp8), weight_scale (bf16/fp16)}``.
  - lm_head DOES NOT need to be untied — most fp8 packs leave lm_head
    dense (it sits inside ``modules_to_not_convert``). When a pack does
    quantize lm_head, the swap picks it up via the same fp8 detection.
  - Block-wise FP8 (``weight_scale_inv`` + ``weight_block_size`` of
    e.g. [128, 128] — DeepSeek-V3 / Qwen3.6 native FP8) is supported via
    the bf16-dequant path: each 128x128 tile's fp32 scale broadcasts over
    its tile. The block size is read from
    ``quantization_config.weight_block_size`` (default 128x128).
  - "Dynamic" activation fp8 is transparent at this layer: weights
    dequantize to bf16 / fp16 and the matmul runs in the activation
    dtype.

Read ``quantization_config.weight_block_size`` (``[bo, bi]``) for a
DeepSeek-V3 / Qwen3.6 native-FP8 pack, or ``None`` for per-tensor /
per-channel packs (which ignore it anyway).

Looked up in ``config.json`` (HF canonical) with an ``hf_quant_config
.json`` fallback (modelopt-style). Cached per model directory (parsed
once). Missing / malformed ⇒ ``None`` ⇒ the consumer defaults to 128x128
for any block-wise scale it sees — the universal native-FP8 tile.

Scale suffixes this backend recognises, in lookup order.

Per-tensor / per-channel first, then block-wise — the order
:meth:`bind` prefers them in. A checkpoint family that names its
scale grid differently overrides this instead of re-implementing
detection, listing and binding.

Return the FP8 quantized Linear subclass mirroring ``dense_cls``.

``MergedColumnParallelLinear`` is supported at any TP degree — the
fused tensor (e.g. the GDN ``in_proj_qkv`` in a MIXED_PRECISION pack)
binds as one column-parallel FP8 Linear whose ``fp8_load`` shards
each sub-block independently across ranks (see
:class:`FP8MergedColumnParallelLinear`), so each rank keeps whole
heads even when the ``[Q, K, V]`` widths differ. Any other
unregistered dense class raises ``RuntimeError``.

Read weight + weight_scale and bind via :meth:`fp8_load`.

``dtype=None`` is critical for both tensors — weight is fp8
(a cast would corrupt) and weight_scale's dtype carries the
export's choice (bf16 / fp16 / fp32). Block-wise packs ship the
scale under ``weight_scale_inv``; the tile shape comes from
:meth:`_weight_block_size`.

Return ``{tensor_name: (shape, dtype_tag)}`` for every buffer an
FP8 linear holds AFTER :meth:`bind`, derived analytically from the
safetensors headers — no weight materialization, no GPU.

FP8 performs NO kernel-side repack: the bound ``weight_fp8`` keeps
the source ``(out, in)`` fp8 layout and ``weight_scale`` keeps its
source per-tensor/per-channel shape. So the post-bind shapes equal
the on-disk shapes one-to-one. (Marlin-permuted backends — AWQ /
EXL3 / NVFP4 — cannot use this analytic path; their repack kernels
define the output layout, so they keep the load->stream_compact
route. FP8 is the deterministic subset.)

Used by the cold-boot allocate-then-fill planner to pre-size flat
slabs for FP8 checkpoints; a parity test asserts these predictions
equal the real bound buffer shapes on a live fixture.

``fp8_w8a8`` + the pack's scale-tile shape when EVERY routed-expert
projection is the block-wise native-fp8 layout, else ``None``.

The fused kernel consumes one fp8 expert stack per stage against one
``weight_scale_inv`` grid laid out on the same axes as the weight, so
the whole expert set must ship that grid and nothing else. A
per-tensor / per-channel ``weight_scale`` projection (whose single
scalar has no per-tile axis for the stacked kernel to index), a
projection carrying no scale, and any mix of the two answer ``None``
and keep the per-expert modules.

The tile shape is the pack's declared
``quantization_config.weight_block_size``, falling back to the
universal native-fp8 128x128 tile when the config omits it.

FP8 weight-quantization Linear classes — concrete impl of the
:mod:`arbi_serve.weight_quant.base.QuantLinearBase` contract.

Bridge between arbi-serve's :class:`LinearBase` family and the FP8
on-disk format. Each FP8 projection is stored as a pair
``{weight (fp8), weight_scale (bf16/fp16)}`` and dequantized at the
layer boundary into the activation dtype, then run through the
standard ``F.linear``.

The dequant cost amortizes over a few tokens of matmul; on Ada with
the bf16 GEMM path this is competitive at the small batches typical
of online serving.

The binding API is :meth:`fp8_load`, called by
:meth:`FP8Backend.bind` with the two tensors. The dense
:meth:`weight_loader` is locked out by the parent :class:`QuantLinearBase`.

dtype handling. Activations stay at their incoming dtype (bf16/fp16);
the dequantized weight is materialized in that dtype and cached. Re-
binding via :meth:`fp8_load` clears the cache.

TP. FP8 packs the output dim along the leading axis of ``weight``
(out, in). Column-parallel slices the OUTPUT axis (weight dim 0,
weight_scale dim 0); row-parallel slices the INPUT axis (weight dim 1
only — per-channel scales are output-axis only, so they're unchanged
on row-parallel; per-tensor scales are scalar). No special block
alignment is required — the fp8 unit is a single byte per element.

LoRA. FP8 + LoRA composes correctly — :class:`QuantLinearBase`
inherits :meth:`LinearBase._maybe_apply_lora`, which dispatches to
the BGMV kernel; the kernel accumulates the correction on top of the
dequantize-then-matmul output regardless of how it was computed.

Import the ``triton_gemm`` package once so BOTH
``torch.ops.arbi_fp8.block_fp8_mm`` and
``torch.ops.arbi_fp8.triton_fused_mm`` are registered before first use.

The package is the op-registration site (a module-import side effect) and
is otherwise pulled in only by the boot-time kernel warmup — which does not
run when the fp8 fast paths engage on a config the warmup didn't
pre-register (e.g. a per-channel weight-only pack whose prefill takes the
SM_89 triton-fused path). Importing it here — idempotently, memoized via
``sys.modules`` and the module flag — makes both fast paths self-contained
for the engine and any non-engine caller. Runs on the first eager forward
(before cudagraph capture), so it never imports inside a captured region.

Common surface for the three FP8 parallel variants.

Carries the weight (fp8) + weight_scale (bf16/fp16) buffers
(non-persistent — they come fresh from safetensors at load) and a
cached dequantized weight tensor that is built lazily on first
forward. Re-binding via :meth:`fp8_load` invalidates the cache.

FP8 equivalent of :class:`ReplicatedLinear`.

Used for the lm_head when an FP8 export quantizes it. Most public
FP8 packs leave lm_head dense (it lives in
``modules_to_not_convert``); when one does quantize lm_head it
ships its own ``lm_head.weight`` + ``lm_head.weight_scale`` and
the swap picks it up via the FP8 detection.

``__init__`` + ``forward`` come from :class:`ReplicatedLinearMixin`.

FP8 equivalent of :class:`ColumnParallelLinear`.

Output dim is sharded across the TP group. ``out_features`` is the
GLOBAL output size; this class stores the LOCAL slice and slices
the FP8 tensors along the output axis inside :meth:`fp8_load`.

TP slicing math (per-rank range from :func:`tp_shard_range`):

  - ``weight: (out_features, in_features)`` →
    slice dim 0: ``[rank * local_out : (rank+1) * local_out, :]``
  - per-channel ``weight_scale: (out_features, 1)`` →
    slice dim 0: same as weight
  - per-tensor ``weight_scale: (1,)`` → unchanged

No special pack-boundary alignment — fp8 is one byte per element
(:attr:`tp_block_align` = 1). ``__init__`` + ``forward`` come from
:class:`ColumnParallelMixin`.

Multi-GPU end-to-end is unverified on this build (single-GPU box);
the slicing math is unit-tested.

FP8 equivalent of :class:`~arbi_serve.models.linear.VocabParallelLMHead`.

Vocab-aware column-parallel head: ``out_features`` is the PADDED vocab (a
tp_size multiple — same convention as the dense head it mirrors),
``org_vocab_size`` is the real vocab, and :meth:`forward` trims the padded
logit columns after the gather so they never leak into sampling. For a
tp-divisible vocab the trim is a no-op.

Built by load-time FP8 RTN head quantization
(:mod:`arbi_serve.weight_quant.head_quant`): each rank quantizes its
ALREADY-LOCAL dense vocab shard and binds it via :meth:`fp8_load_local`.
Row-wise RTN scales are computed per output row (or per ``bo x bi`` tile),
so a rank's independently-quantized shard is bit-identical to its slice of
a globally quantized tensor.

The drafter's distributed greedy argmax
(:meth:`Qwen3_5MtpHead._distributed_greedy_argmax`) consumes
:meth:`~arbi_serve.weight_quant.base.ColumnParallelMixin.local_logits` —
the rank-local shard row WITHOUT the gather — and reconstructs the global
winner from ``(max, idx)`` pairs; that path duck-types on this class's
``local_logits`` + ``org_vocab_size`` surface.

FP8 equivalent of :class:`RowParallelLinear`.

Input dim is sharded across the TP group. ``in_features`` is the
GLOBAL input size; this class stores the LOCAL slice and slices
the FP8 weight along the input axis inside :meth:`fp8_load`. The
per-channel scale lives on the OUTPUT axis and is therefore
unchanged on row-parallel; per-tensor scale is scalar.

No special pack-boundary alignment — fp8 is one byte per element
(:attr:`tp_block_align` = 1). ``__init__`` + ``forward`` come from
:class:`RowParallelMixin`.

TP slicing math:

  - ``weight: (out_features, in_features)`` →
    slice dim 1: ``[:, rank * local_in : (rank+1) * local_in]``
  - ``weight_scale: (out_features, 1)`` or ``(1,)`` → unchanged

FP8 mirror of :class:`MergedColumnParallelLinear` (TP-aware).

Used for the GDN ``in_proj_qkv`` projection in Qwen 3.5 / 3.6 when the
checkpoint quantizes it to FP8 (e.g. ``Qwen3.6-27B-NVFP4``, a
MIXED_PRECISION pack whose attention/GDN projections are FP8). HF stores
it as a single fused tensor stacking several column-parallel
sub-projections along the OUTPUT axis (dim 0) with potentially-distinct
widths (``[Q(key_dim) | K(key_dim) | V(value_dim)]``; gate_up_proj:
``[gate | up]``) and one ``weight_scale`` (per-tensor scalar or
per-channel ``(out,)``/``(out,1)``).

At TP=1 the fused tensor maps to a single column-parallel FP8 Linear
with ``out_features = sum(_shard_out)``; consumers ``.split()`` the
output downstream and see correct per-shard slices.

At TP>1 a NAIVE contiguous output-dim cut splits mid-sub-block (rank 0
receives ``[full Q + half K]`` at TP=2 with QKV widths
``[key_dim, key_dim, value_dim]``), corrupting the per-shard semantics
the GDN block depends on. The fix — mirrors
:class:`AWQMergedColumnParallelLinear` /
:class:`NVFP4MergedColumnParallelLinear` — is to shard EACH sub-block
independently across TP ranks, then concatenate the per-rank slices in
declared order.

FP8 is one byte per element (no pack factor) and the scale indexes only
the OUTPUT axis (per-channel) or is scalar (per-tensor), so the output
shard never straddles a quant boundary; the only per-sub-block
constraint is ``sub_dim % tp_size == 0`` (:attr:`tp_block_align` = 1).
A per-tensor scalar scale is SHARED verbatim across ranks (its dequant
is identical for every row), so ``dequant(shard) ==
dequant(full)[this rank's rows]`` exactly; a per-channel scale is sliced
per-sub-block alongside the weight rows. The constructor asserts the
divisibility so a bad shape fails loudly rather than truncating a
sub-block.

Bind the (weight, weight_scale) pair and validate shapes.

``weight_block_size`` is the ``(bo, bi)`` tile for a block-wise
(``weight_scale_inv``) pack — threaded from the backend's read of
``quantization_config.weight_block_size`` (defaults to 128x128 when
omitted, the DeepSeek-V3 / Qwen3.6 convention). Ignored for
per-tensor / per-channel scales.

An exponent-container scale (E8M0) is widened to fp32 here. Every
E8M0 code is a power of two fp32 represents exactly, so the widening
changes no value, and it is what lets the dequant and ``scaled_mm``
paths hold one floating scale dtype.

Subclasses override to apply per-rank slicing before delegating
back here. At TP=1 every subclass calls straight through.

True when the bound scale is a block-wise ``weight_scale_inv``.

Derived from the (weight, scale) shapes + ``self._block_size`` — so it
is correct after a warm flat-dump boot (which restores the persistent
``weight_fp8`` / ``weight_scale`` buffers without re-running
:meth:`fp8_load`). Block-wise routes to the native block-scaled fp8
GEMM (:meth:`_fp8_forward_block`); the ``torch._scaled_mm`` /
``triton_fused`` fast paths assume a per-tensor / per-channel
``scale_b`` and are gated off for it below.

(Re)build the prebuilt RowWise ``scale_b`` (1, N) fp32 buffer read
by :meth:`_fp8_forward_scaled_mm`.

Idempotent — a pure function of ``weight_scale`` (bound in
:meth:`fp8_load`). Called ONCE more post-load, post-device-placement
by :func:`arbi_serve.engine.build.prepare_derived_weights` so the
buffer's ``data_ptr`` is stable for the compiled + captured decode
graph (and so a warm flat-dump boot — which restores ``weight_scale``
without re-binding — repopulates it). No-op before a bind (empty
``weight_scale``): the scaled_mm path is unreachable until a bind
populates ``weight_fp8``.

The value is byte-identical to the inline build it replaces (the
RowWise ``scale_b`` in :meth:`_fp8_forward_scaled_mm`):

  * **per-tensor** scalar → constant ``(1, N)`` column vector
    (``expand`` then ``contiguous`` — ``torch._scaled_mm`` requires a
    contiguous RowWise ``scale_b``);
  * **per-channel** ``(out,)`` / ``(out, 1)`` → reshape to ``(1, N)``.

Dequantize to the dense ``(out_features, in_features)`` weight in
``F.linear`` layout (row-major over output).

Mirrors :meth:`NVFP4LinearBase.dequantized_weight` /
:meth:`AWQLinearBase.dequantized_weight` — the contract the GDN
block's fused dense ``[B|A]`` prebuild relies on
(:meth:`GDNBlock.prepare_fused_ba_dequant`) when a checkpoint
quantizes the sub-tile ``in_proj_b`` / ``in_proj_a`` projections to
FP8. Public FP8 packs keep b/a dense (sub-tile ``out`` width), so this
is a latent-path guarantee, exercised once at build time — never a hot
path.

Byte-identical to the dequant-fallback arm of the forward
(:meth:`_ensure_dequantized` → :func:`dequantize_fp8`): both upcast
fp8 → ``dtype`` then broadcast-multiply by ``weight_scale``. Default
dtype is bf16 (the serving dtype). Does NOT touch the forward's
``_ensure_dequantized`` cache — a fresh tensor, so it never perturbs
the ``_dequant_dtype`` the hot path may have warmed.

Raises (never silently degrades) when called before :meth:`fp8_load`.
FP8 has no Marlin repack (buffers are never released), so — unlike the
NVFP4/AWQ overrides — there is no Marlin guard to trip.

Mirror :func:`triton_fused.is_eligible` plus the env-var gate.

Inlined here (rather than calling ``is_eligible`` directly)
because the latter wants the already-quantized fp8 inputs +
scale tensors, which we'd otherwise materialize before
knowing whether to bother. Shape-gate semantics are
byte-equivalent to ``triton_fused.is_eligible`` for per-channel
scales; per-tensor (scalar) scales are additionally restricted to
wide-N shapes (see the numel + skinny-N comments below).

Engage ``torch._scaled_mm`` (FP8 tensor cores) instead of the
bf16-dequant GEMM.

Two regimes, split by compute capability:

  * **Blackwell (sm_100/sm_120, cc[0] >= 10)** — engage
    UNCONDITIONALLY for a bf16 input. Blackwell has first-class
    FP8 tensor cores, and the scaled_mm path never materializes the
    dense bf16 weight, so it is the only FP8 route that does NOT
    hold a bf16 mirror resident. Both per-tensor and per-channel
    scales are handled below.

  * **Ada (cc==(8,9))** — the middle path only when it pays: the
    per-token dynamic-quant overhead amortizes against the
    bf16-GEMM it replaces only on wide shapes. Gate: per-channel
    scale, NOT triton_fused-eligible (that wins above), and fp8

FP8 mma via PyTorch's built-in ``torch._scaled_mm``, bf16 out.
Skips the bf16 dequant cache entirely — no dense weight is ever
materialized on this path.

Activation quantization is ALWAYS **per-row (per-token)** — the
activation scale ``sa`` is ``(M, 1)`` (one max-abs per row), NEVER a
batch-global ``x_2d.abs().amax()`` scalar. A batch-global activation
scale couples every row of the batch through one shared divisor: a
single large-magnitude sibling row inflates the shared amax and drives
every OTHER row's fp8 activation toward zero, collapsing that row's
logits → argmax token 0. That sibling row is reachable on any
``max_batch > 1`` decode — most acutely a :mod:`decode_pad` pad-up
scratch row that borrows a FREE recurrent slab row and reads stale GDN
state (large garbage), but also an ordinary batched-decode outlier.
Per-row scaling is exactly the per-row independence the decode-cudagraph
pad-up contract (:mod:`arbi_serve.runtime.decode_pad`) depends on, and
it matches every other activation-quant path in the tree (the RowWise
branch below, ``triton_fused``, and the NVFP4 dynamic quant). It is also
strictly finer than a per-tensor activation scale (less quant error).

The scale RECIPE must match on both operands (``torch._scaled_mm``
rejects a rowwise × tensorwise mix), so both weight granularities run
the **RowWise** recipe — they differ only in how ``scale_b`` is built:

  * **per-tensor** (scalar weight scale — the modelopt static-FP8
    layout, e.g. the 27B mixed pack's ``linear_attn`` / attn
    projections): the scalar broadcasts to a CONSTANT ``(1, N)`` column
    vector. RowWise computes ``out_ij = sa_i * sb_j * (x_fp8 @ b)_ij``,
    so a constant ``sb`` recovers exactly ``per-row-act × per-tensor-
    weight`` — mathematically identical to the intended per-tensor
    weight scaling, with a per-row (not per-batch) activation scale.

  * **per-channel** (scale shape ``(out,)`` / ``(out, 1)``): reshapes
    to ``(1, N)`` directly — the Ada wide-N decode path, unchanged.

Engage the native block-wise fp8 GEMM for a ``weight_scale_inv`` pack.

Preferred UNCONDITIONALLY over the bf16-dequant fallback on an
fp8-tensor-core card (Ada sm_89 or newer) whenever the flag is on and
the bound scale is block-wise: the kernel runs the fp8 cores directly
off the on-disk bytes, so it never materializes the dense bf16 weight
mirror. There is no per-shape gate — the memory saving is decisive
even at small M where the raw GEMM latency would be a wash (mirrors
the Blackwell ``_scaled_mm`` rationale). Non-fp8 cards / flag-off
fall back to bf16 dequant.

Block-scaled fp8 GEMM path: per-``(token, K-block)`` dynamic-quantize
the input, apply the per-``(N,K)``-block ``weight_scale_inv`` inside the
accumulation, return bf16. No bf16 dequant of weights — the dense bf16
mirror is never built on this path.

Bind ALREADY rank-local FP8 tensors.

The inherited :meth:`FP8ColumnParallelLinear.fp8_load` slices GLOBAL
checkpoint tensors by rank; the load-time RTN path quantizes the rank's
local dense shard directly, so its tensors must bind WITHOUT
re-slicing. Delegates straight to the base binder (shape-validated
against the LOCAL ``out_features``).

Slice this rank's output rows for one sub-block out of ``tensor``.

Applies to ``weight`` (``(out, in)`` fp8, dim 0) and a per-channel
``weight_scale`` (``(out,)`` or ``(out, 1)``, dim 0) alike — both
are output-major on dim 0. The rank's chunk within the sub-block is
``[rank*local_sub, (rank+1)*local_sub)``, shifted by the sub-block's
offset in the fused output dim.

Per-sub-block + per-rank slice of the fused FP8 tensors along the
OUTPUT axis (dim 0), concatenated in declared shard order, then
bound via the base load.

A per-tensor (scalar) ``weight_scale`` is passed through unchanged —
it applies to every row identically. A per-channel scale is sliced
per sub-block alongside the weight rows. A BLOCK-WISE
``weight_scale_inv`` (``weight_block_size`` tiles — Qwen3.6 native
FP8's fused GDN ``in_proj_qkv``) is sliced in BLOCK units on dim 0,
which requires every sub-block boundary AND every per-rank shard cut
to land on a ``bo``-row block edge (asserted loudly). At TP=1 all
slicing is the identity and the fused tensors bind unchanged.

Delegates to :meth:`FP8LinearBase.fp8_load` directly (NOT ``super()``,
which would re-run :class:`FP8ColumnParallelLinear`'s contiguous
output-axis slice and double-shard); the per-sub-block slicing here
already produced this rank's local weights.

FP8 dequantize helpers + scale-shape classification.

Three FP8 layouts are accepted by the loader:

  - **per-tensor**:  ``weight_scale`` is a scalar — one float per linear.
  - **per-channel**: ``weight_scale`` is shape ``(out_features, 1)``
    or ``(out_features,)`` — one scale per output row.
  - **block-wise**:  ``weight_scale_inv`` is shape
    ``(ceil(out/bo), ceil(in/bi))`` — one fp32 scale per ``bo x bi``
    tile of the weight (DeepSeek-V3 / Qwen3.6 native FP8, ``bo == bi ==
    128``). Despite the ``_inv`` suffix the stored value is applied by
    **multiplication**:
    ``dequant[i, j] = weight[i, j].to(dtype) * scale[i // bo, j // bi]``.
    The last block along each axis is **partial** when the weight dim is
    not a multiple of the block size — the scale tensor rounds UP
    (``ceil``) and the trailing block simply covers fewer rows / cols.

For per-tensor / per-channel the dequant is a single broadcast multiply;
block-wise broadcasts each tile's scalar over its tile. Both fp8 dtypes
(``float8_e4m3fn`` and ``float8_e5m2``) are handled identically — the
upcast to bf16/fp16 preserves the float value bit-for-bit (within the
target dtype's range), and the multiply by scale produces the
dequantized bf16/fp16 weight.

Decide which dequant broadcasting shape to use.

Inputs:
  weight: fp8 (out_features, in_features)
  scale:  fp16/bf16/fp32 — shape encodes the granularity
  block_size: (bo, bi) block shape used to recognise a block-wise
    ``weight_scale_inv`` (default 128x128 — the DeepSeek-V3 / Qwen3.6
    native-FP8 convention).

Returns the :class:`FP8ScaleKind` for this linear. Unrecognised
shapes raise ``ValueError`` so a checkpoint with an unexpected
layout fails loud rather than silently mis-dequantizing.

Ambiguity note: per-channel ``(out, 1)`` is matched **before**
block-wise, so it always wins its exact shape. A block-wise scale is
``(ceil(out/bo), ceil(in/bi))`` — for a real weight (out >> bo) that
is distinct from ``(out, 1)``; the two only collide on a degenerate
``out <= bo`` single-block-row weight, where per-channel is the
intended reading anyway.

Broadcast a ``(ceil(out/bo), ceil(in/bi))`` block scale up to the
dense ``(out_features, in_features)`` weight grid.

``repeat_interleave`` on each axis tiles a block's scalar over its
``bo`` rows / ``bi`` cols; slicing back to the exact feature counts
trims the (possibly partial) trailing block — so a non-block-multiple
dim is handled with no special-casing.

Dequantize an fp8 weight to ``target_dtype`` (bf16 or fp16).

Returns the dense ``(out_features, in_features)`` matrix. Handles
per-tensor, per-channel and block-wise (``weight_scale_inv``) scales;
``block_size`` selects the tile shape for the block-wise case.

Weight-side RTN quantization of a dense bf16/fp16 matrix to FP8 e4m3.

The inverse of :mod:`arbi_serve.weight_quant.fp8.loader`: that module reads a
checkpoint's fp8 payload + scale, this one PRODUCES the identical pair from a
dense weight at load time (data-free round-to-nearest — no calibration set).
The output is indistinguishable from a checkpoint-supplied pair: same dtype
(``float8_e4m3fn``), same scale shapes :func:`classify_scale` recognises, and
— critically — the same **multiplier** convention. Despite the checkpoint's
``weight_scale_inv`` name the stored value is applied by MULTIPLICATION::

    weight_dense ≈ weight_fp8.to(dtype) * scale[broadcast]

so the emitted scale is ``amax / 448`` (NOT ``448 / amax``). Reading it the
other way is ``scale**2`` off yet still produces finite, plausible activations,
which is why the convention is pinned by measurement in the tests rather than
by a comment.

Three granularities, one per :class:`FP8ScaleKind`:

  - ``PER_TENSOR``   — scale ``(1,)``: one fp32 scalar for the whole matrix.
  - ``PER_CHANNEL``  — scale ``(out, 1)``: one fp32 scalar per output row.
  - ``BLOCK_WISE``   — scale ``(ceil(out/bo), ceil(in/bi))``: one fp32 scalar
    per ``bo x bi`` tile, the DeepSeek-V3 / Qwen3.6 native-FP8 layout the
    ``arbi_fp8::block_fp8_mm`` kernel consumes directly. A trailing partial
    tile is covered by a full scale entry over the rows/cols that exist.

Peak memory is bounded by ROW CHUNKING (:data:`FP8_RTN_ROW_BLOCK`): output rows
quantize independently under every granularity except ``PER_TENSOR`` (whose
global amax is gathered in a separate, allocation-free first pass), so a chunk
of rows can be quantized straight into the final tensors. The live transient is
ONE chunk of fp32 upcast plus its fp8 result — never a second copy of the whole
matrix. The amax reductions use ``amax``/``amin`` rather than ``abs().amax()``
for the same reason: ``abs()`` would materialise a second full chunk.

Shape of the scale tensor :func:`rtn_quantize_fp8` emits.

Pure arithmetic on the weight shape — no allocation — so a boot budget can
price the transform before it runs, and so a test can assert the emitted
layout against :func:`classify_scale` without materialising a weight.

``(need, resident, chunk_scratch)`` for the fp8 RTN transform.

The resident result is the fp8 payload plus its fp32 scale; the only
transient is ONE row chunk of fp32 upcast plus that chunk's fp8 staging
copy. There is no repack step (fp8 keeps the source ``(out, in)`` layout),
so — unlike the INT4 head path — the peak carries no second copy of the
packed matrix. Pure arithmetic, testable without a device.

Max-abs of ``x`` over ``dims``, WITHOUT materialising ``x.abs()``.

``x.abs().amax(dims)`` allocates a full-size copy of ``x``; the whole point
of the row chunking is that the transient stays one chunk, so the amax is
taken as ``max(amax, -amin)`` — two reductions, each allocating only the
reduced result.

Quantize one row chunk into ``out_fp8``, returning its scale slice.

``scale`` is the pre-computed scale for ``PER_TENSOR`` (a global reduction
the caller already did); ``None`` for the row-local granularities, whose
scale is derived here. The fp32 upcast is divided and clamped IN PLACE so
only one chunk-sized float buffer is ever live.

RTN-quantize a dense ``(out, in)`` weight to ``(fp8_e4m3, fp32 scale)``.

The returned pair is byte-for-byte the shape/dtype contract a checkpoint
ships, and satisfies the SAME multiplier convention
(``weight ≈ fp8.to(dtype) * scale``), so it binds through the unmodified
:meth:`FP8LinearBase.fp8_load` and dequantizes through the unmodified
:func:`~arbi_serve.weight_quant.fp8.loader.dequantize_fp8`.

Rows are processed in ``row_block`` chunks straight into the final tensors:
the result is identical to quantizing the whole matrix at once (rows are
independent under per-channel and — because ``row_block`` is a multiple of
every supported block height — under block-wise too), while the transient
stays one chunk. ``PER_TENSOR`` takes a first pass to reduce the global
amax; that pass allocates only per-chunk scalars.

Raises on a non-2-D weight or an already-fp8 input (double-quantizing is
always a wiring bug, never a request).

RTN-quantize ``weight`` into caller-provided ``out_fp8`` / ``out_scale``.

Byte-for-byte identical math to :func:`rtn_quantize_fp8` (same row chunking,
same multiplier convention) but writes into PRE-ALLOCATED outputs instead of
allocating them — so the peak carries no second copy of the fp8 result, only
one ``row_block`` chunk of fp32 scratch. The cold-boot input-embedding fp8
path uses this to quantize the checkpoint table STRAIGHT into its slab views
(``weight_fp8`` / per-row ``weight_scale``), so the bf16 table is never
resident in the weights slab.

``out_fp8`` must be ``float8_e4m3fn`` with ``weight``'s shape; ``out_scale``
must have the :func:`fp8_rtn_scale_shape` layout and may be ANY float dtype
(the per-tile fp32 scale is cast on store — e.g. bf16 for the embedding's
gather-dequant scale).

Custom GEMM kernels — fp8 Triton mm (rowwise fused + block-wise).

Importing this package registers ``torch.ops.arbi_fp8.triton_fused_mm``
and ``torch.ops.arbi_fp8.block_fp8_mm`` as a side effect. The arbi-serve
FP8 linear path dispatches through these ops on fp8-capable cards.

Public API:
  - :func:`triton_fused.fp8_triton_fused_mm` -- rowwise (per-token act ×
    per-channel weight) fp8 GEMM (registered as
    ``torch.ops.arbi_fp8.triton_fused_mm``), SM_89.
  - :func:`triton_fused.is_eligible(a, b, sa, sb, out_dtype)` --
    predicate for the rowwise fast path.
  - :func:`block_fp8.block_fp8_mm` -- block-wise (per-``(token,K-block)``
    act × per-``(N,K)``-block weight) fp8 GEMM for DeepSeek-V3 / Qwen3.6
    ``weight_scale_inv`` packs (registered as
    ``torch.ops.arbi_fp8.block_fp8_mm``). Never materializes the bf16
    dequant mirror.
  - :func:`block_fp8.is_eligible(x, weight, weight_scale)` -- predicate for
    the block-wise fast path.
  - :func:`warmup` -- run tiny GEMMs to JIT-compile the kernels before the
    first real request.

Rowwise default ON; set ``ARBI_SERVE_FP8_TRITON_GEMM=0`` to disable.
Block-wise default ON; set ``ARBI_SERVE_FP8_BLOCK_GEMM=0`` to disable
(both fall through to the bf16 dequant path).

Block-wise fp8 GEMM (W8A8, per-``(N,K)``-block weight scale × per-``(token,
K-block)`` activation scale) — the native path for DeepSeek-V3 / Qwen3.6-style
``weight_scale_inv`` packs.

Semantics (matching the arbi bf16-dequant reference in
:func:`...fp8.loader.dequantize_fp8`):

    out[m, n] = Σ_kb  As[m, kb] · Ws[n // bo, kb] · Σ_{k∈kb} A_fp8[m, k]·B_fp8[n, k]

where the K axis is split into blocks ``kb`` of width ``block_k`` (== the
activation group size), the weight scale ``Ws`` is one fp32 scalar per
``(bo × bi)`` weight tile, and the activation scale ``As`` is one fp32 scalar
per ``(token, block_k)`` group. The scale of each K-block is applied to that
block's partial ``tl.dot`` BEFORE it is summed into the fp32 accumulator — this
is the defining difference from the row/col-wise fused kernel
(:mod:`.triton_fused`), where a single epilogue scale suffices.

Unlike the per-tensor / per-channel fp8 fast paths, this kernel NEVER
materializes the dense bf16 weight: it runs the fp8 tensor cores directly off
the on-disk fp8 bytes instead of holding a bf16 dequant mirror resident on top
of its fp8 storage.

The activation is dynamically quantized per ``(token, block_k)`` group
(:func:`_per_token_group_quant`) — never a batch-global amax, which would
couple rows and let one large-magnitude sibling row zero every other row's fp8
activation (see the discussion in ``fp8.linear._fp8_forward_scaled_mm``).

Pick ``(BLOCK_M, BLOCK_N, BLOCK_K, num_stages, num_warps)``.

``BLOCK_K`` is pinned to ``block_k`` (the weight/activation K-block, 128 on
every public native-fp8 pack) so each K-tile lies inside exactly one scale
block — the ``offs_ks = k // (block_k // BLOCK_K)`` indexing then reduces to
``offs_ks = k``. ``BLOCK_N`` need NOT divide the weight's ``bo``: the kernel
loads the weight scale with a per-column ``offs_bn // bo`` gather, so a wide
N-tile straddling a block boundary still gets the right per-column scale.

Tuning follows the rowwise kernel's shape logic (N ≫ M for Linear layers):
keep BLOCK_M small to hold SM occupancy, BLOCK_N=128, 3 pipeline stages.

Per-``(token, K-group)`` dynamic fp8 quant.

One program per ``(row, group)``. ``scale = max(|x_group|) / 448`` (clamped
to ``EPS``); ``x_q = clamp(x / scale, ±448) → fp8_e4m3``. A partial trailing
group (``K`` not a multiple of ``GROUP``) is masked — the out-of-range lanes
read ``0`` and do not perturb the amax.

Block-wise fp8 GEMM: ``out = quant(x) @ dequant(weight).T`` on fp8 cores.

Args:
    out: ``(M, N)`` bf16, pre-allocated.
    x:   ``(M, K)`` bf16 activation (dynamically quantized here, per
         ``(token, group_k)`` group).
    weight: ``(N, K)`` fp8_e4m3 row-major (the on-disk ``weight`` tensor).
    weight_scale: ``(ceil(N/group_n), ceil(K/group_k))`` fp32/bf16
         ``weight_scale_inv`` — one scalar per ``(group_n × group_k)`` tile.
    group_n, group_k: the ``(bo, bi)`` weight block shape (128, 128).

Return True iff the block-wise fp8 kernel can run for these operands.

Requires an fp8-tensor-core card (Ada sm_89 or newer), a bf16 activation,
an fp8 weight and a 2-D ``weight_scale_inv``. There is no shape gate: the
kernel is preferred UNCONDITIONALLY over bf16-dequant on eligible cards
because it also drops the multi-GiB bf16 dequant mirror — the memory win is
decisive even where a tiny-M GEMM would be a wash on raw latency (the same
reasoning the Blackwell ``_scaled_mm`` path uses in ``fp8.linear``).

Triton fused fp8 GEMM with rowwise (per-token) act scale × per-channel weight
scale + optional bias, for SM_89.

Single kernel: tl.dot(fp8, fp8, fp32) → epilogue applies sa[m] * sb[n] + bias[n]
inside the kernel before the bf16 store. No HBM round-trip, no dual launch.

Autotune over a few BLOCK_M / BLOCK_N / BLOCK_K / num_stages tuples covering
the decode (M=8..128) + prefill (M=512..2048) regimes.

Pick (BLOCK_M, BLOCK_N, BLOCK_K, num_stages, num_warps) for given M, N.

M-axis: BLOCK_M=16 only at M<16 (mma.m16n8k32 minimum); BLOCK_M=32
everywhere else. Linear shapes are N>>M, so M-tiling barely registers
next to the N-tile count; a smaller BLOCK_M keeps SM occupancy high.

N-axis: BLOCK_N=64 for N<=4096, BLOCK_N=128 for wider N.

3 pipeline stages.

One-shot fused fp8 GEMM with per-row act × per-col weight scaling.

Args:
    out: (M, N) bf16, pre-allocated.
    a:   (M, K) fp8_e4m3 row-major.
    b:   (K, N) fp8_e4m3 col-major (== (N, K) row-major .t()).
    scale_a: (M,) or (M, 1) fp32, per-row.
    scale_b: (N,) or (1, N) or (N, 1) fp32, per-col.
    bias: optional (N,) bf16.

Look up an AOT artifact for ``(BLOCK_M, BLOCK_N, HAS_BIAS)``.

Memoized per shape key in :mod:`arbi_serve.kernels._aot` — the
ASTSource is built lazily only on the first launch of a given
``(BLOCK_M, BLOCK_N, HAS_BIAS)`` bucket; subsequent launches are a
dict lookup.

Return True iff the kernel is expected to beat stock CUTLASS at this
(M, K, N).

Triton has a fixed per-launch floor (kernel launch + scale-load +
``tl.dot`` setup) that CUTLASS undercuts on small GEMMs, so eligibility
is gated by shape rather than by a blanket ``M >= 8`` predicate. The
gating below routes only the shapes where Triton wins:

| bucket        | rule                        |
|---------------|-----------------------------|
| very-wide-N   | ``N >= 16384``              |
| medium-wide-N | ``N >= 8192 and M >= 128``  |
| long-K        | ``K >= 4096 and M >= 64``   |
| narrow        | else → stock CUTLASS        |

Load-time RTN quantization of the checkpoint-dense head weights.

The Qwen3.6-27B AWQ export deliberately skips ``lm_head`` and ``mtp.fc``
(``quantization_config.ignore``) — they run bf16. The bf16 lm_head is a
large decode kernel and vocab shard, and the DRAFTER re-reads that
shard 5×/step through the greedy pair-argmax. There is no calibration
data for these layers in the checkpoint, so the scheme is data-free RTN
(:func:`rtn_pack_quantized_int4` — same asymmetric INT4 / group-32 grid
as the trunk's AWQ layers), and the quality gates live with the SERVING
modes rather than a calibration contract:

``ServerConfig.lm_head_quant``:

  - ``"off"`` (default): checkpoint-faithful bf16 heads. No change.
  - ``"drafter-int4"``: the DRAFTER projects through an RTN-INT4 copy of
    the lm_head vocab shard (installed via
    :meth:`Qwen3_5MtpHead.set_draft_lm_head_override`) and ``mtp.fc`` is
    swapped to INT4; verify keeps the bf16 lm_head. Draft-QUALITY risk
    only — verify corrects every emitted token — so the gate is
    accept_len staying flat. Costs +qweight bytes (the bf16 shard stays
    for verify); cuts the drafter's 5×/step shard traffic ~4×. The copy
    is READ-ONLY (a private packed module), so this mode is safe on a
    TIED-embedding head too (small models — Qwen3.5-0.8B): the copy reads
    the shared embedding weight, the tie stays fully intact (input-embed
    read + bf16 verify keep using the original tensor), and only the
    drafter projects through the private INT4 copy.
  - ``"full-int4"``: ``model.lm_head`` itself is swapped to RTN-INT4 —
    verify and drafter share it, the bf16 shard is freed (growing the KV
    budget). This CHANGES EMITTED LOGITS —
    the gate is a fixed-config eval holding, not a TPOT number.
  - ``"drafter-fp8"`` / ``"full-fp8"``: the same two placements with an
    RTN **FP8-e4m3** head (:func:`~arbi_serve.weight_quant.fp8.rtn
    .rtn_quantize_fp8`, per-output-row scale) instead of INT4. Half the
    shrink of INT4 (2× vs 4×), a far smaller WEIGHT error (e4m3 keeps 3
    mantissa bits per row-scaled value — a half-ULP floor of 2**-4
    relative, where INT4-g32 quantizes to 16 levels per group), and —
    for a memory-tight boot — a much smaller load-time TRANSIENT: fp8 does
    no Marlin repack, so its spike is ONE row chunk, where INT4's is a
    second full copy of the packed head plus a chunk
    (:func:`_rtn_int4_head_peak_bytes` vs
    :func:`~arbi_serve.weight_quant.fp8.rtn.fp8_rtn_peak_bytes`).

    Scope that memory claim honestly: it is a TRANSIENT win, not a
    total-peak win. The fp8 head is RESIDENTLY larger than the So fp8 installs where INT4 OOMs when the boot is short on the
    SPIKE; it does not help a boot that is short overall, and there INT4
    (or ``off``) is the lever.

    ACTIVATION ASYMMETRY — the reason "fp8 is gentler than INT4" is NOT a
    blanket claim. The INT4 vocab head pins ``_a8_eligible = False``, so it
    serves W4A16: quantized weights, **bf16 activations**. The FP8 head has
    no such lane — the only fp8 routes that avoid materializing a dense
    bf16 mirror (``torch._scaled_mm`` at M=1, ``triton_fused_mm`` at the
    verify/draft rows) are **W8A8**, so they also quantize the activation
    per token. So an fp8 head trades a much smaller weight error for an
    activation-quantization term the INT4 head never pays. On the weight
    side fp8 is clearly closer to bf16; end-to-end on emitted logits the
    two are not ordered a priori, which is exactly why:

    ``full-fp8`` CHANGES EMITTED LOGITS and carries the SAME fixed-config
    eval gate as ``full-int4`` — no weaker, for the activation reason
    above. ``drafter-fp8`` keeps verify in bf16, so emitted logits are
    unchanged under strict rejection sampling exactly as ``drafter-int4``
    is, and its W8A8 error is confined to draft QUALITY (the accept_len
    gate). ``drafter-fp8`` is therefore the mode to reach for first.

The FP8 modes are EXPLICIT-ONLY: ``auto`` never resolves to one. The
head lives in a checkpoint's ``modules_to_not_convert`` because the
export deliberately kept it dense; converting it is an operator
decision, never a default.

``"auto"`` (the default) resolves at the boot seam via
:func:`resolve_head_quant_mode`. The FIRST gate is
whether MTP / spec-decode is ENABLED for this server (``cfg.mtp.enabled``
— ``--enable-mtp`` / ``--mtp-n-draft>=1`` / ``ARBI_MTP``), NOT whether the
model merely carries a bundled ``mtp.*`` head. A bundled head is a model
CAPABILITY (built whenever the checkpoint ships ``mtp.*`` keys), not a
request: with MTP off the drafter never runs (``eng.spec_decode`` stays
``NoSpecDecode``), so a drafter INT4 head copy would be DEAD VRAM. MTP off
⇒ ``off`` (``mtp_disabled_no_drafter``), for TIED and UNTIED alike.

With MTP on, ``auto`` resolves ``drafter-int4`` where the read-only copy
can be built — a dense ``mtp.fc`` plus a vocab-shard lm_head, either a
DENSE untied ``VocabParallelLMHead`` (the AWQ lane — checkpoint
``quantization_config.ignore`` leaves it dense — or any bf16-trunk MTP
rig) OR a TIED head reading a ``VocabParallelEmbedding`` (small models;
the copy rides with the MTP the operator already opted into). It
resolves ``off`` (with a recorded, specific reason) where the lever is
absent:

  - MTP disabled → no drafter to redirect (``mtp_disabled_no_drafter``);
  - the checkpoint already ships a QUANTIZED head (EXL3 trellis / NVFP4 /
    AWQ-lm_head export) — the drafter already projects through a quantized
    head, so there is nothing to add (``head_already_quantized``);
  - a non-vocab-parallel head (a plain unwrapped ``nn.Linear``) with no
    shard surface to pack a copy from (``head_not_vocab_parallel``);
  - the MTP head's ``fc`` is not the dense class the RTN swap targets
    (``mtp_fc_not_dense``).

The scoping keeps the default honest: ON exactly where a live drafter can
use the copy, absent (not silently ignored) elsewhere. Explicit modes keep
their strict fail-loud semantics — a drafter/full mode with MTP DISABLED
raises (dead VRAM, no silent install); ``--lm-head-quant drafter-int4`` on
a TIED head with MTP on is SUPPORTED (builds the read-only copy); and
``--lm-head-quant full-int4`` on a TIED head raises loudly (the REPLACE
path frees the bf16 shard, which a tied head shares with the embedding and
cannot free).

Runs at the engine-boot seam AFTER weights load and BEFORE KV sizing
(so the freed / added bytes are visible to the KV budget), and BEFORE
``build_mtp_driver`` (whose ``head._lm_head is model.lm_head`` identity
assert the ``full-int4`` swap preserves via the tied-module rebind).

Resolve ``cfg.lm_head_quant`` to a concrete mode for ``model``.

Returns ``(concrete_mode, reason)``. ``reason`` is ``None`` for every
non-``auto`` input (explicit modes pass through untouched and keep
their fail-loud semantics in :func:`apply_head_quant`) and for an
``auto`` that resolved ON; an ``auto`` that resolved ``off`` carries
the machine-readable reason (also recorded as a counter refusal by the
caller seam in ``engine/build.py``).

``mtp_enabled`` is the authoritative gate: whether MTP / spec-decode is
ENABLED for THIS server (``cfg.mtp.enabled`` — ``--enable-mtp`` /
``--mtp-n-draft>=1`` / ``ARBI_MTP``), NOT whether the model merely
carries a bundled ``mtp.*`` head. A bundled head is a model CAPABILITY,
built whenever the checkpoint ships ``mtp.*`` keys (``qwen3_5.py``),
not a request: with MTP off the drafter never runs (``eng.spec_decode``
stays ``NoSpecDecode``), so the INT4 drafter copy would be DEAD VRAM.
MTP off ⇒ every head-quant mode resolves ``off`` (``mtp_disabled_no_drafter``).

With MTP on, ``auto`` eligibility is MECHANICAL: the drafter-int4 lever
exists iff ``mtp.fc`` is dense and the lm_head exposes a vocab shard to
build the read-only INT4 drafter copy from — a dense untied
``VocabParallelLMHead`` (the AWQ lane / any bf16-trunk MTP rig) OR a
TIED head reading a ``VocabParallelEmbedding`` (small models; the copy
rides with the MTP the operator already opted into). It is absent only
on EXL3 (trellis lm_head — already a quantized drafter head) and on
non-vocab-parallel heads.

``auto`` never resolves to a bare silent ``off``: every ``off`` it
returns carries a specific, honest reason naming the missing mechanism.
It also never resolves to an FP8 mode — converting a head the checkpoint
deliberately left dense is an operator decision, so ``drafter-fp8`` /
``full-fp8`` are reachable only by naming them.

``(packed_bytes, meta_bytes)`` for the INT4 form of an ``(out, in)`` shard.

``packed_bytes`` is the int32 nibble matrix (``out × in/8 × 4``);
``meta_bytes`` is the bf16 per-group scale plus the packed int32
zero-point. Pure arithmetic on the shapes — no allocation — so the boot
budget can price the transform BEFORE it runs.

``(need, resident, repack_copy, transient)`` for the INT4 transform.

The peak is the resident INT4 head PLUS one more full copy of the packed
matrix: :meth:`pack_quantized_load_local` builds the Marlin-laid-out
``_marlin_qweight`` (same byte count as the packed head) while the raw
packed source is still live, so two full packed copies coexist at the
repack. A bounded ``transient`` rides on top — whichever stage's working
set is larger:

  * the RTN pack's fp32 upcast block
    (:data:`_LM_HEAD_RTN_ROW_BLOCK` rows × ``in`` × 4 B), and
  * the Marlin repack's per-output-block scratch
    (:func:`~arbi_serve.weight_quant.awq.marlin.pack_quantized_to_marlin`
    holds one tile-N-aligned standard-layout slice plus its Marlin-layout
    result, together bounded by ``_MARLIN_REPACK_CHUNK_BYTES``; a head that
    fits in one block instead holds the whole ``2 × packed`` transiently,
    which for a single-block head is itself ≤ that budget).

The two never coexist (Stage A's scratch is freed before the repack), so
the peak takes the MAX, not the sum. The W4A16 repack does not unpack the
weight to a dense ``(in, out)`` nibble grid (8× the packed bytes), so the
priced peak matches what the load actually allocates. Pure arithmetic so
the budget is testable without a device.

Refuse the head-quant transform LOUDLY when this rank cannot fund it.

The INT4 pack is a mid-boot allocation on a device already holding the
dense weights, so at a tight ``--gpu-memory-utilization`` it can be the
allocation that OOMs — an opaque ``CUDA out of memory`` inside an RTN
kernel, tens of seconds into a load. Price it from the shard shape first
and raise with the NUMBERS instead: what the transform needs, what is
free, and the two knobs that fix it. See
:func:`_rtn_int4_head_peak_bytes` for what the peak is made of.

Row-chunked :func:`rtn_pack_quantized_int4` — same result, bounded peak.

Output rows are quantized independently (grouping is along the INPUT axis),
so packing 8-row-aligned output blocks into the corresponding output-axis
slices of the final tensors is bit-identical to packing the whole
``(out, in)`` matrix — while the fp32 working set is one block, not the
full shard, and nothing ever holds two copies of the packed head.
``out_features`` is a multiple of the pack-factor 8 (the INT4 zero-point
packs 8 output rows into one int32), so every block except possibly the
last is ``_LM_HEAD_RTN_ROW_BLOCK`` rows and the last is the 8-aligned
remainder; both keep the zp packing local to aligned 8-row groups. Falls
back to a single whole-matrix pack when the shard fits in one block.

Refuse the FP8 head transform LOUDLY when this rank cannot fund it.

Same contract as :func:`_assert_vocab_head_quant_fits` (price the peak from
the shard shape, raise with the numbers instead of OOMing tens of seconds
into a load) over the FP8 peak: the resident payload + scale, plus ONE row
chunk of fp32 upcast and its fp8 staging copy. No repack term — fp8 keeps
the source ``(out, in)`` layout — which is the whole reason this transform
lands where the INT4 one does not.

RTN-FP8 an ALREADY rank-local dense vocab shard into an
:class:`FP8VocabParallelLMHead`.

Argument contract is :func:`_build_rtn_int4_vocab_head`'s. The emitted
``(fp8, scale)`` pair is indistinguishable from a checkpoint-supplied one —
same dtype, same :func:`~arbi_serve.weight_quant.fp8.loader.classify_scale`
granularity, same multiplier convention — so it binds through the
unmodified FP8 load path and every fp8 forward kernel serves it with no
special-casing. Row scales make the rank-local quantization bit-identical
to slicing a globally quantized tensor.

RTN-INT4 an ALREADY rank-local dense vocab shard into an
:class:`AWQVocabParallelLMHead`.

``weight`` is this rank's ``(per_shard, hidden)`` shard (a
``VocabParallelLMHead.weight`` when untied, or the tied
``VocabParallelEmbedding.weight`` when the head is tied — the two expose
the identical shard surface). ``padded_vocab`` is the GLOBAL padded vocab
the dense head stores as ``out_features``; ``org_vocab_size`` is the real
(tokenizer) vocab. Bit-identical to slicing a globally quantized tensor —
RTN groups along the INPUT axis, so output rows quantize independently.

RTN packs the shard in bit-identical vocab-row chunks (rows quantize
independently — grouping is along the INPUT axis). This bounds the fp32
working set to one block, which fits the GPU default caching allocator's
slack even on a WARM boot — where the dense weights DMA straight into the
cuMem weights pool, so the default allocator never forms the large
reservation the cold scattered-load leaves behind, and a whole-matrix
(cuMem's VA reservation is counted against the default allocator;
expandable_segments is force-disabled with cuMem). The chunk transients + the small INT4 result
are freed to the default allocator afterward (``empty_cache`` at the boot
seam), so the warm boot reaches KV profiling with the SAME free VRAM as the
cold one.

Build a quantized vocab head from a dense UNTIED vocab-parallel
lm_head for the ``full-*`` REPLACE path.

This is the head that will REPLACE ``model.lm_head`` (verify + drafter share
it, bf16 shard freed), so it hard-requires an untied ``VocabParallelLMHead``:
a tied head shares its tensor with the input embedding and cannot be freed —
quantizing-and-replacing it would corrupt the embedding read. That guard is a
real safety invariant and stays. (The drafter COPY path,
:func:`_rtn_quantize_drafter_head`, is tied-safe because it only READS.)

Build the read-only quantized DRAFTER copy of the lm_head vocab shard.

Unlike the ``full-*`` REPLACE path (:func:`_rtn_quantize_vocab_head`),
this only READS the shard weight to pack a PRIVATE copy for the drafter
override; the real head is never mutated or freed. That makes it safe on a
TIED head: the copy reads the shared ``VocabParallelEmbedding`` weight, the
tie stays 100% intact (the input-embedding read and the bf16 verify path
both keep using the original tied tensor), and ONLY the drafter projects
through the private quantized copy.

A tied ``VocabParallelEmbedding`` exposes the SAME rank-local shard surface
as an untied ``VocabParallelLMHead`` — ``weight`` ``(per_shard, hidden)``,
a padded global vocab, and ``org_vocab_size`` — so the copy is built the
identical way and is bit-identical to slicing a globally quantized tied
tensor (with TP>1, each rank packs its own local shard, matching the untied
per-rank slice).

Apply the ``lm_head_quant`` mode to a fully-loaded model IN PLACE.

Returns the RESIDENT quantized-head bytes on this rank (0 for ``off``) —
the KV-floor refusal sizes its ``--lm-head-quant off`` escape from this
measured number rather than a quoted constant.

The lm_head RTN packs in bit-identical vocab-row chunks
(:func:`_rtn_pack_quantized_int4_rowchunked`) so its fp32 working set stays
small enough for the GPU default allocator's slack even on a warm boot
(where the cuMem weights pool has starved the default allocator's VA) — a
whole-matrix pack would OOM / gc-retry-stall there. No scratch pool needed.

``mtp_enabled`` (``cfg.mtp.enabled``) is the gate: head-quant drafter
modes require MTP/spec-decode ENABLED at launch, not the mere presence
of a bundled ``mtp.*`` head. With MTP off the drafter never runs, so
the INT4 head copy would be dead VRAM — an explicit drafter/full mode
then fails loud.

See the module docstring for mode semantics. Fail-loud contract: an
unknown mode, a drafter mode with MTP disabled (or no MTP head), an
already-quantized head, a ``full-*`` mode on a TIED head (its shard cannot
be freed — it is shared with the embedding), or a ``drafter-*`` mode on a
non-vocab-parallel head all raise — no silent fallback to bf16 (a
config asking for quant and silently not getting it would surface as
unexplained model drift, not an error). A ``drafter-*`` mode on a TIED head
(with MTP on) is SUPPORTED (a private read-only quantized copy; the tie
stays intact).

Yield the runtime head-quant AWQ-Marlin linears :func:`apply_head_quant`
installs: the drafter INT4 ``lm_head`` override, the drafter ``mtp.fc``, and
— full-int4 only — the swapped ``model.lm_head``.

These are the ONLY Marlin linears the EAGER MTP drafter-seed / verify
forward projects through directly; the transformer blocks serve through
captured decode graphs (their Marlin workspace VA + lock lifecycle are baked
into the replay). A residency-wake re-establish of per-member Marlin state
(:func:`~arbi_serve.engine.stable_va_controller._reset_head_quant_marlin_workspaces_on_wake`)
walks exactly this set. Filtered to modules that actually carry an ACTIVE
Marlin workspace, so a dense/tied ``lm_head`` (the drafter-int4 case, where
only the override + fc are quantized) is skipped — as is an FP8 head, which
has no Marlin workspace at all (its GEMM carries no cross-CTA lock state, so
there is nothing for a residency wake to re-establish).

Read the head-quant mode a fully-built model is ALREADY in.

Inverse of :func:`resolve_head_quant_mode` (which decides what to APPLY):
this reports what a model's head modules currently ARE, by inspecting the
swap :func:`apply_head_quant` performs. Used by the donor-share path, whose
member inherits its head modules from the donor (shared, never re-applied),
to write the truthful resolved mode back to ``cfg`` for downstream readers.

  * ``full-<scheme>``  — ``model.lm_head`` is quantized (verify + drafter
    share it; the bf16 shard was freed);
  * ``drafter-<scheme>`` — a drafter lm_head override is installed while
    ``lm_head`` stays dense (verify keeps the bf16 head);
  * ``off`` — dense heads (or a checkpoint-quantized head, which
    head-quant leaves alone).

``<scheme>`` is read off the installed module's class, so an fp8 donor is
reported as fp8 rather than silently as the int4 mode of the same
placement.

Reproduce a donor's RUNTIME head-quant on a donor-bound MEMBER by SHARING
the donor's already-built INT4 head modules — zero new weight VRAM, no
re-pack. Returns the number of head modules shared.

A build-required same-model swap constructs the member's module graph fresh
(dense ``mtp_head.fc``, dense ``lm_head``) and aliases every weight the donor
holds by qualified name (:class:`DonorWeightBinder`). Runtime head-quant
(:func:`apply_head_quant`) is a LATER build phase, so a donor booted with it
carries only the packed INT4 head buffers while the fresh member still
carries the dense ``mtp_head.fc.weight`` / ``lm_head.weight`` the donor no
longer has. Without this mirror those names are a donor COVERAGE GAP and the
member re-loads a dense head from the checkpoint — correct output, but a
second head's VRAM the share exists to avoid.

Head-quant weights are IMMUTABLE and read-only at serve time (identical
rationale to the shared dense / EXL3 / AWQ weights, and to the head's own
``_tied_modules`` embed/lm_head sharing), so the member SHARES the donor's
built head modules outright. This mirrors the checkpoint-quant weight
aliasing (same immutable-shared-storage premise) at MODULE granularity —
which the head needs because the drafter override is deliberately held OUT
of the module graph (the binder never visits it) and because a head module's
small non-zero state buffer (``_format_marker``) is not a shareable
zero-element placeholder the generic binder aliases. Runs during the
member's graph construction, BEFORE the binder walks it, so the member never
presents a dense head the donor cannot back.

Reproduces exactly what the donor's ``apply_head_quant`` produced:

  * ``mtp_head.fc`` → the donor's quantized fc (both drafter/full modes);
  * the drafter lm_head override (``_draft_lm_head_holder``) → installed from
    the donor's (a ``drafter-*`` mode);
  * ``model.lm_head`` → the donor's quantized head (a ``full-*`` mode), with
    the head's tied ``_lm_head`` rebound so the ``MtpDriver`` identity assert
    (``head._lm_head is model.lm_head``) still holds.

Scheme-agnostic: the sharing keys on the ``is_quantized`` marker every
:class:`QuantLinearBase` carries, so an INT4 and an FP8 donor share by the
identical path.

Only reproduces a divergence the fresh member does NOT already have: a
checkpoint-quantized head (EXL3 trellis / AWQ-lm_head export) is built
quantized by ``apply_quant_if_present`` and shared by the generic binder, so
``fc`` / ``lm_head`` are already quantized here and are left untouched. A
dense-head donor (head-quant off) shares nothing.

The ``lm_head`` mirror is independent of the MTP head: a ``full-*`` mode
swaps ``model.lm_head`` itself, so a member built without ``mtp_head`` still
needs it.

Generic quant loader: detect → swap → bind.

One entry point — :func:`apply_quant_if_present` — that any arch's
``from_safetensors`` calls once after constructing the model graph.
The loader walks the registry, picks the backend that owns the
checkpoint, enumerates quantized linear paths, replaces each
corresponding dense Linear with the backend's mirror class, and
delegates per-Linear tensor binding to the backend's :meth:`bind`.

Path resolution is backend-agnostic but understands the model's
``weight_map()``: when an arch renames safetensors keys (LFM2-MoE
``feed_forward.gate_proj`` ↔ HF ``feed_forward.w1``), the rename is
inverted from ``weight_map`` and applied during path resolution. Most
arches use HF-faithful naming and the remap is empty.

Repoint tied module references at their quant replacements.

A module that shares another module by reference (DeepSeek-MTP
convention: :class:`Qwen3_5MtpHead` reuses the main model's
``embed_tokens`` + ``lm_head``, held in a ``_tied_modules`` list so
they are not registered as the head's own submodules) keeps pointing
at the pre-swap dense object after the swap rebinds the owning
attribute via ``setattr``. Walk every module's ``_tied_modules`` and
replace any swapped entry with its quant mirror, so identity holds
(``MtpDriver`` asserts ``head._lm_head is model.lm_head``) and the
head runs on the loaded quant weights rather than an orphaned,
never-filled dense tensor.

No-op when nothing was swapped or the model has no tied modules.

Public alias for :func:`_detect_backends`. Every registered backend
that claims ``stc``, in ascending ``detection_priority`` order; ``[]``
for a dense checkpoint.

:func:`apply_quant_if_present` is the hook for an arch that owns a
``weight_map``; a consumer that drives the swap itself (the DFlash
drafter — its graph is built outside the arch table) resolves its
backends here and then walks ``quant_class_for`` / ``bind`` directly.

Ask ``backend`` whether it can serve ``expert_paths`` as a stacked
fused-MoE expert set, and with what spec.

``None`` for a dense checkpoint (the caller already takes the fused
path) and for a backend that cannot stack this pack. See
:meth:`~arbi_serve.weight_quant.base.QuantBackend.stacked_moe_spec`.

One-call quant hook for an arch's ``from_safetensors``.

Returns ``(backend, n_swapped)``. ``backend`` is ``None`` for
dense checkpoints (no-op); ``n_swapped`` is the number of linears
converted to their quantized mirrors.

``dtype`` is the engine activation/weight dtype (e.g. ``bf16``). It is
passed to every :meth:`QuantBackend.bind` as ``compute_dtype`` (the
dtype a backend's 16-bit side tensors must be materialized in) and
sizes the dense (embedding / norm) slab on the cold DIRECT-slab path.

The arch should call this AFTER constructing the model graph on
meta but BEFORE materializing meta tensors and BEFORE
:func:`load_model_weights`. Doing the swap before materialization
matters at scale: at 9B+ parameters, allocating dense CUDA tensors
The swap reads only ``in_features`` / ``out_features`` off the
Linear instances, so meta-device input is fine; backend-specific
tensor buffers are allocated on ``device`` directly.

The arch's ``weight_map`` should route through
:func:`arbi_serve.weight_quant.base.filter_weight_map_for_quant` so
swapped linears are excluded from the dense bind.

Every registered backend whose :meth:`detect` accepts ``stc``, in
ascending ``detection_priority`` order (as :func:`all_backends` returns).

A MIXED_PRECISION checkpoint returns more than one (e.g. NVFP4 + FP8); a
uniform quant pack returns exactly one; a dense checkpoint returns ``[]``.
One backend's optional-dependency import error must not hide the others,
so detection is best-effort per backend.

Swap + bind every linear ``backend`` owns that a higher-priority
backend has not already ``claimed``. Mutates ``claimed`` (each swapped
safetensors path) and ``tied_swaps`` (``(old, new)`` pairs the caller
rebinds once after all backends run). Returns the count swapped.

Called once per detected backend so a MIXED_PRECISION pack can dispatch
several backends over disjoint linear sets; for a uniform pack it runs
exactly once and ``claimed`` never fires.

Note that ``model_path``'s quant bind was DEFERRED by ``skip_weight_load``.

The record is ``(model_path, backend, safetensors_path)`` on
``model._deferred_quant_binds``. Written during the swap because that is the
only point where the checkpoint→model prefix remap is complete: the swap
marks each linear ``is_quantized``, and ``weight_map()`` filters quantized
``.weight`` entries out from then on, taking the rename with them.

Run the DEFERRED :meth:`QuantBackend.bind` for the named quant modules.

``module_paths`` are model-side dotted paths whose modules the swap already
replaced with their quant mirrors under ``skip_weight_load`` — their weight
buffers are still ``__init__`` placeholders. Binds each from the checkpoint
exactly as the cold path's inline bind does.

Returns ``(n_bound, unowned)`` — ``unowned`` lists the requested paths no
detected backend claims, i.e. modules this checkpoint cannot fill. The
caller decides what an unfillable module means; this function never
partially serves one silently.

Boot-time fire evidence: count the AWQ linears that ACTUALLY bound the
a8 (fp8-activation) kernel vs the a16 kernel, bucketed by role. ``_marlin_is_a8``
is the authoritative per-linear decision (set at bind, read in forward), so
this is a true fire-count on the served config, not a flag reading. Free —
a one-time boot walk, off the serving path.

OCP MXFP4 payload decoding.

On-disk format::

    blocks   uint8 [..., K // 32, 16]  2 E2M1 nibbles per byte
                                       (LOW = col 2j, HIGH = 2j+1)
    scales   uint8 [..., K // 32]      E8M0 biased exponent

The value is ``e2m1_level(nibble) * 2**(scale - 127)``. There is no
per-tensor global scale.

MXFP4 reaches this repo as MoE expert stacks only: gpt-oss quantizes the
experts and leaves attention, router, embeddings and lm_head bf16. Those
stacks are carried by :class:`~arbi_serve.models.moe.FusedMoE` under
``quant_kind="mxfp4"``, bound by
:func:`~arbi_serve.models._moe_stacked_loader.bind_stacked_expert_stacks`
and multiplied by ``fused_moe_kernel_mxfp4``. This module is the shared
decode used by the off-device reference and the tests; there is no
dense-``Linear`` MXFP4 backend because no exporter emits one.

Dequantize an MXFP4 pack to a dense tensor of ``dtype``.

Args:
    blocks: ``(..., K // 2)`` uint8 — two E2M1 nibbles per byte.
    scales: ``(..., K // block_size)`` uint8 — E8M0 biased exponents.

Returns ``(..., K)``.

NVFP4 weight-quantization backend for arbi-serve.

Importing this subpackage registers the backend with the global
quant-backend registry. Detection is arch-independent; the hardware
requirement is decided per layer at bind time from the checkpoint's own
modelopt ``quant_algo`` (see :class:`NVFP4Backend`).

The on-disk format is the modelopt / compressed-tensors NVFP4
``W4A16``-pack-quantized layout produced by ``nvidia-modelopt``::

    weight            uint8  [out, in // 2]   (2 FP4 E2M1 nibbles per byte;
                                              LOW = col 2j, HIGH = col 2j+1)
    weight_scale      fp8e4m3 [out, in // 16]  per-16-element block scale
    weight_scale_2    fp32   ()                per-tensor global scale
    input_scale       fp32   ()                per-tensor activation scale
                                              (informational; unused — we
                                              compute activation scales
                                              dynamically per forward)

W4A4 forward path is **NVFP4 × NVFP4 dynamic**: the activation is
quantized to NVFP4 on the fly with calibration-free per-row absmax (one
Triton kernel pass), then ``torch._scaled_mm`` runs the cuBLAS Blackwell
FP4 tensor-core MMA. The ``input_scale`` carried by the checkpoint is
informational only; dynamic per-row activation quantization does not
need the calibrated value. This path needs ``get_device_capability() >=
(10, 0)``; :meth:`NVFP4Backend.bind` refuses a W4A4 layer at load on
older silicon.

W4A16 (weight-only) packs
-------------------------
modelopt MIXED_PRECISION packs (e.g. ``Qwen3.6-27B-NVFP4``) tag their
MLP / lm_head layers ``W4A16_NVFP4`` — 4-bit weight, **16-bit
activation**. These must NOT go through the W4A4 path above: quantising
the activation to fp4 injects per-layer error that compounds into
degenerate output. The backend routes them to :mod:`weight_only_gemm`
— a fused Triton fp4-weight × bf16-activation GEMM that keeps the
weight packed and dequantises each tile in registers (no dense bf16
mirror). It runs a standard bf16 ``tl.dot`` (no fp4 tensor cores), so it
runs on any Triton-capable CUDA device including Ada sm_89. On Blackwell
a weight-only layer additionally repacks into the vendored fp4-Marlin
kernel at bind time (:mod:`marlin`), whose prebaked cubin is sm_120.

NVFP4 :class:`QuantBackend` implementation.

Detects the modelopt / compressed-tensors NVFP4 layout from the
safetensors keys, maps HF-canonical paths to the NVFP4Linear mirror
class, and reads the four-tensor bundle per linear.

Detection key set (must all be present for any prefix to count):

  - ``<path>.weight``           uint8 (packed FP4)
  - ``<path>.weight_scale``     fp8_e4m3 (per-16-block)
  - ``<path>.weight_scale_2``   fp32 (per-tensor global)
  - ``<path>.input_scale``      fp32 (per-tensor activation, optional;
                                informational — we use dynamic act-quant)

The ``.weight_scale_2`` suffix is the strongest discriminant — it's
unique to modelopt NVFP4 packs. The FP8 backend uses ``.weight_scale``
without ``.weight_scale_2``; AWQ pack-quantized uses ``.weight_packed``
(distinct from ``.weight``); legacy AWQ uses ``.qweight``. To prevent
the FP8 backend from claiming an NVFP4 checkpoint via ``.weight +
.weight_scale``, the NVFP4 backend is registered BEFORE FP8 in
:mod:`arbi_serve.weight_quant.registry`.

Hardware requirement is PER LAYER, decided by the checkpoint's own
modelopt ``quant_algo`` tag, not by the backend as a whole:

  * ``W4A16_NVFP4`` (weight-only) — the packed fp4 × bf16 GEMM in
    :mod:`weight_only_gemm` runs a plain bf16 ``tl.dot`` over
    register-dequantised weight tiles. No FP4 tensor cores; permitted on
    any CUDA device Triton supports.
  * ``NVFP4`` (W4A4) — the activation is quantised to FP4 and
    ``torch._scaled_mm`` runs the FP4 MMA, which exists only on
    Blackwell (``cuda.get_device_capability() >= (10, 0)``).
    :meth:`NVFP4Backend.bind` refuses such a layer at load on older
    silicon rather than degrading it to a weight-only path the pack was
    not calibrated for.

NVFP4 weight-quantization backend.

Dispatches per layer on the modelopt ``quant_algo`` in
``hf_quant_config.json`` (see :meth:`_weight_only_paths`):

  * ``NVFP4`` — **W4A4**: activation is dynamically quantised to fp4
    and the cuBLAS fp4 tensor-core MMA runs
    (:meth:`NVFP4LinearBase._nvfp4_forward`). Blackwell only; refused
    at :meth:`bind` on older silicon.
  * ``W4A16_NVFP4`` — **weight-only**: activation stays bf16 and the
    packed fp4 × bf16 GEMM runs (:mod:`weight_only_gemm`). Runs on any
    Triton-capable CUDA device. Running such a layer at A4
    over-quantises the activation into degenerate output, so the mode
    MUST follow the checkpoint.

Return True iff the safetensors carries the modelopt NVFP4
four-tensor layout for at least one prefix.

Arch-independent: the per-layer W4A4 requirement is enforced in
:meth:`bind`. Declining here on non-Blackwell silicon would hand
the checkpoint to the FP8 backend, whose ``.weight_scale``
signature is a strict superset — it then rejects the uint8 packed
weight with an error naming the wrong format.

Return the NVFP4 quantized Linear subclass mirroring
``dense_cls``.

``VocabParallelLMHead`` is the lm_head's dense class. Unlike the
0.8B fixture (modelopt ``exclude_modules`` keeps lm_head dense),
production modelopt packs such as ``Qwen3.6-27B-NVFP4`` quantize the
lm_head — it carries ``lm_head.weight_scale_2``. It is a
``ColumnParallelLinear`` subclass (``gather_output=True``), so an
exact-class table lookup misses it; it is mapped explicitly below,
mirroring :meth:`EXL3Backend.quant_class_for`.

``MergedColumnParallelLinear`` is supported at any TP degree — the
NVFP4 mirror shards each fused sub-block independently across ranks
(see :class:`NVFP4MergedColumnParallelLinear`). Any other
unregistered dense class raises ``RuntimeError``.

``(quantized_layers map, top-level quant_algo)``.

Read from ``hf_quant_config.json`` → ``quantization``; modelopt
also writes the same fields to ``config.json`` →
``quantization_config``, which is the only declaration some packs
ship. A MIXED_PRECISION pack carries the per-layer map; a uniform
pack carries only the top-level tag. Cached per model directory.

Checkpoint prefixes whose NVFP4 algo is *weight-only* (W4A16).

``W4A16_NVFP4`` layers keep bf16 activations (weight-only); bare
``NVFP4`` layers are W4A4 (activation → fp4). A uniform pack's
top-level ``quant_algo`` applies to every prefix.

Cached per model directory. Missing / unreadable config ⇒ empty set
⇒ every layer is treated as W4A4.

``nvfp4`` when EVERY routed-expert projection is a WEIGHT-ONLY
modelopt NVFP4 triple, else ``None``.

The fused kernel dequantizes the packed weight into the bf16
activation's type and never quantizes the activation, so it serves
``W4A16_NVFP4`` exactly and would over-quantize nothing — but a
``NVFP4`` (W4A4) projection is calibrated for an fp4 activation and
does not belong on it. A missing member of the triple, a
non-``(N, K // 2)`` uint8 weight, a block scale that is not
fp8_e4m3 on the 16-element grid, or a non-scalar
``weight_scale_2`` also answer ``None`` and keep the per-expert
modules.

Dynamic NVFP4 quantization for activations.

NVFP4 (modelopt convention) packs each value as a 4-bit E2M1 float
in a per-16-element-block scaling scheme:

    block_scale_fp8[m, b] = absmax(x[m, b*16 : (b+1)*16]) / 6.0
    global_scale_fp32     = max(block_scale_fp8) / 448.0
    block_scale_fp8[m, b] /= global_scale_fp32  ; cast to fp8_e4m3
    eff_scale[m, b]       = block_scale_fp8[m, b].dequantized * global_scale_fp32
    nibble[m, k]          = round_to_nearest_fp4_level(x[m, k] / eff_scale[m, k//16])

Two FP4 nibbles pack into one byte: LOW nibble (bits 0..3) holds the
column ``2j`` value, HIGH nibble (bits 4..7) holds the column ``2j+1``.
This matches the modelopt on-disk weight layout we read at load time.

The activation pipeline computes the per-row global_scale from the
running absmax — no calibration.

cuBLAS FP4 MMA (``torch._scaled_mm`` on sm_120+) consumes the per-block
fp8 scales in a **swizzled** layout — see :func:`swizzle_blockscale`.
Unswizzled flat layout produces output that LOOKS plausible but is
wrong; only the swizzled layout is correct.

The Triton kernel produces unswizzled scales; the swizzle is a
follow-up reshape/permute on the host (cheap relative to the GEMM).

Per-row dynamic NVFP4 quantization (torch reference).

Inputs:
  ``x`` — bf16 / fp16 / fp32 tensor of shape ``(M, K)`` (or higher
  rank; flattened to 2-D for quantization). ``K`` must be a multiple
  of ``block_size``.

Returns:
  ``packed_uint8`` ``(M, K//2)`` — LOW nibble = col 2j, HIGH = col 2j+1
  ``block_scale_fp8`` ``(M, K//block_size)`` fp8_e4m3
  ``global_scale_fp32`` ``(M, 1)`` per-row fp32 scalar

Bit-equivalent to :func:`nvfp4_quant_dynamic_triton` on GPU, slower.
Used as oracle in the parity unit test and as the CPU code path.

Dequantize an NVFP4 tensor bundle to its dense weight (torch
reference — build-time / test use only, NOT a hot path).

Exact inverse of the packing convention documented at module top
(and of the modelopt on-disk weight layout): LOW nibble (bits 0..3)
holds column ``2j``, HIGH nibble (bits 4..7) holds column ``2j+1``;
each 16-element block is scaled by ``block_scale_fp8 * global_scale``.

Inputs:
  ``packed`` — ``(rows, cols // 2)`` uint8 packed FP4 nibble pairs.
  ``block_scale_fp8`` — ``(rows, cols // block_size)`` fp8_e4m3.
  ``global_scale`` — scalar fp32 (modelopt ``weight_scale_2``).

Returns the dense ``(rows, cols)`` weight in ``dtype`` (contiguous).

Per-row dynamic NVFP4 quantization (Triton fast path on GPU).

Falls back to the torch reference on CPU or when Triton's not
importable. Returns the same triple as
:func:`nvfp4_quant_dynamic_torch`.

The Triton implementation is ~30 LOC and JIT-compiles on first
call; subsequent forwards are kernel-bound.

The actual Triton-launched implementation.

Arranged as a per-row kernel: each program handles one row of the
(M, K) input. K must be a multiple of ``block_size``; block_size
itself is fixed at 16 (NVFP4 spec). The kernel computes the row's
absmax (one tree-reduce), derives the global_scale, then walks
each 16-element block to round to FP4 + emit the fp8 block scale.

Packing strategy: load 16 fp32 values per block, classify each to a
4-bit FP4 code, and pack 2 codes/byte by loading the input twice
with even/odd lane offsets — one lane writes the LOW nibble, the
paired lane writes the HIGH nibble. We avoid atomic-or (which
Triton doesn't support on uint8) by computing both nibbles
arithmetically into a single 8-element vector and storing it as
one contiguous byte slice.

Reshape the (M, K_blocks) fp8 block scales into the swizzled
layout cuBLAS Blackwell FP4 MMA expects.

Passing the flat unswizzled layout to ``torch._scaled_mm`` yields
output that LOOKS plausible but is bit-wrong; the swizzled layout
below is correct.

Padding: M is rounded up to multiples of 128, K_blocks to multiples
of 4. The padding rows/cols hold zero scales (decoded as zero, no
contribution). The output is a 1-D contiguous tensor of length
``M_padded * K_padded`` per cuBLAS's expected stride.

Layout follows vLLM / TRT-LLM convention:
    (M_pad / 128, K_pad / 4, 32, 4, 4)
permuted to (0, 3, 2, 1, 4) before flattening.

This is a port of vLLM's ``swizzle_blockscale``; the math is the
on-wire cuBLAS contract.

NVFP4 weight-quantization Linear classes — concrete impl of
:mod:`arbi_serve.weight_quant.base.QuantLinearBase`.

Each NVFP4 projection is stored on disk as the four-tensor bundle
``{weight (uint8 packed), weight_scale (fp8e4m3 per-16-block),
weight_scale_2 (fp32 global), input_scale (fp32 calibrated act)}``.

A W4A16 (weight-only) layer runs the packed fp4 × bf16 GEMM and needs no
FP4 tensor cores. The W4A4 forward path below does; the backend refuses
to bind a W4A4 layer on a device without them. The W4A4 path:

  1. Dynamic-NVFP4-quantize the input activation in a single Triton
     pass — emits ``(packed_fp4, per-block fp8 scale, per-row fp32
     global scale)``. No calibration: the row global is ``absmax /
     (FP4_MAX * FP8_MAX)``. Cost is one row-pass over the activation.
  2. Swizzle both the activation and the (cached) weight block scales
     into the cuBLAS Blackwell layout (math in :mod:`dynamic_quant`).
  3. Call ``torch._scaled_mm`` with the packed FP4 inputs and the
     swizzled fp8 scales. Returns the unscaled bf16 output.
  4. Multiply by ``activation_global * weight_global_scale_2`` to
     undo the per-row + per-tensor global scaling.

The result is mathematically equivalent to dequant-then-matmul but
runs through the FP4 tensor cores on sm_120 (RTX 50-series, B200).

dtype handling. Activation enters at bf16/fp16; the dynamic quant
produces fp4 + fp8 + fp32 metadata. Output is bf16 (the standard
serving dtype for arbi-serve). fp16 callers see an internal upcast.

TP. NVFP4 lays out the weight as ``(out, in // 2)`` packed FP4:

  - **Column-parallel** slices the OUTPUT axis (weight dim 0,
    weight_scale dim 0). The packed-FP4 input axis is unsliced so
    nibble-pair boundaries are preserved trivially.
  - **Row-parallel** slices the INPUT axis. Each row holds two FP4
    values per byte, so the slicing must align to even-column
    boundaries — the per-rank ``in_features`` must be divisible by 2.
    The block scale (per-16-element block) imposes a tighter
    constraint: per-rank ``in_features`` divisible by 16. The
    constructor checks both.
  - **MergedColumnParallelLinear** (e.g. GDN ``in_proj_qkv``): the
    safetensors ships the merged tensor at the merged path; at TP=1
    we copy it through as a single Linear with summed out_features.
    Per-shard fused-split TP>1 is **not** wired (mirrors AWQ Marlin).

LoRA. NVFP4 + LoRA composes correctly — the BGMV correction lands on
top of the FP4-MMA output regardless of how it was computed
(:class:`QuantLinearBase` inherits the hook from :class:`LinearBase`).

Common surface for the NVFP4 parallel variants.

Carries the four bound buffers (``weight_packed``,
``weight_scale_fp8``, ``weight_scale_2``, ``input_scale``), the
pre-swizzled weight scale used inside forward, and the local TP
in/out feature counts.

Forward implementation lives in :meth:`_nvfp4_forward`.
Subclasses do TP slicing in their override of :meth:`nvfp4_load`
and call back into ``super().nvfp4_load`` once the per-rank slice
is computed.

NVFP4 mirror of :class:`ColumnParallelLinear`.

No special pack-boundary alignment on the output axis
(:attr:`tp_block_align` = 1). ``__init__`` + ``forward`` come from
:class:`ColumnParallelMixin`.

TP slicing math:

  - ``weight: (out_features, in_features // 2)`` packed-FP4 →
    slice dim 0: ``[rank * local_out : (rank+1) * local_out, :]``
    Pack-axis (dim 1) is on the INPUT axis, untouched.
  - ``weight_scale: (out_features, in_features // 16)`` →
    slice dim 0: same as weight.
  - ``weight_scale_2: ()`` scalar → unchanged.

NVFP4 mirror of :class:`RowParallelLinear`.

Each rank's input shard must align to a full per-block scale
(:attr:`tp_block_align` = :data:`NVFP4_BLOCK_SIZE`). ``__init__`` +
``forward`` come from :class:`RowParallelMixin`; the constructor's
divisibility error keeps the NVFP4-specific wording via
:meth:`_tp_divisibility_message`.

TP slicing math:

  - ``weight: (out_features, in_features // 2)`` packed-FP4 →
    slice dim 1 in BYTES: each rank gets
    ``[:, (rank * local_in) // 2 : ((rank+1) * local_in) // 2]``.
    Pack-axis alignment requires ``local_in`` divisible by 2 (one
    byte = 2 nibbles); the per-block scale tightens this to
    ``local_in % NVFP4_BLOCK_SIZE == 0`` (16). The constructor
    checks the tighter bound.
  - ``weight_scale: (out_features, in_features // 16)`` →
    slice dim 1 in BLOCKS: ``[:, (rank * local_in) // 16 :
    ((rank+1) * local_in) // 16]``.
  - ``weight_scale_2: ()`` scalar → unchanged.

NVFP4 mirror of :class:`MergedColumnParallelLinear` (TP-aware).

Used for the GDN ``in_proj_qkv`` projection in Qwen 3.5 / 3.6, which
HF stores as a single fused tensor stacking several column-parallel
sub-projections along the OUTPUT axis with potentially-distinct
widths (``[Q(key_dim) | K(key_dim) | V(value_dim)]``; gate_up_proj:
``[gate(intermediate) | up(intermediate)]``). The safetensors stores
ONE fused tensor whose output dim is ``sum(_shard_out)``.

At TP=1 the fused tensor maps to a single column-parallel NVFP4 Linear
with ``out_features = sum(_shard_out)``; consumers ``.split()`` the
output downstream and see correct per-shard slices.

At TP>1 a NAIVE contiguous output-dim cut splits mid-sub-block (rank 0
receives ``[full Q + half K]`` at TP=2 with QKV widths
``[key_dim, key_dim, value_dim]``), corrupting the per-shard semantics
the GDN block depends on (it splits the ``in_proj_qkv`` output by
``[key_dim_local, key_dim_local, value_dim_local]``). The fix — mirrors
:class:`AWQMergedColumnParallelLinear` — is to shard EACH sub-block
independently across TP ranks, then concatenate the per-rank slices in
declared order.

NVFP4 output-axis slicing is clean at any ``tp_size`` granularity:
the fp4 pack (2 nibbles/byte) is on the INPUT axis (weight dim 1) and
the 16-element block scale is likewise along the INPUT axis
(``weight_scale`` dim 1), so the OUTPUT shard (dim 0 of both
``weight`` and ``weight_scale``) never straddles a pack or block
boundary. The only per-sub-block constraint is therefore
``sub_dim % tp_size == 0`` (:attr:`tp_block_align` = 1). The
constructor asserts this so a bad model shape fails loudly rather than
silently truncating a sub-block.

Bit-exactness: because each rank slices whole output rows out of each
sub-block and the block/global scales index the untouched input axis,
``dequant(shard) == dequant(full)[this rank's rows]`` exactly — the
concatenation of the per-rank dequantized shards reproduces the
single-GPU dequantized merged weight row-for-row.

Restore the ``_is_bound`` flag after a warm flat-dump reload.

The forward reads only persistent buffers (``weight_packed`` /
``weight_scale_2`` / ``_weight_scale_swizzled`` / ``_weight_only_flag``),
all restored by the mmap fill; the non-buffer state is ``_is_bound``
and ``_weight_only`` (plain attrs the meta graph reset because the warm
path skipped ``nvfp4_load``). Re-derive both from the restored buffers
so the forward's bound-check and W4A16/W4A4 routing survive a warm
boot. Also fired post-compaction on the cold path, where it is a
harmless no-op. No-op while the buffers are still empty placeholders.

Bind the NVFP4 tensor bundle and pre-compute the swizzled
weight scale.

``weight_only`` marks a W4A16 layer (activations stay bf16 — the
packed fp4 × bf16 GEMM path); the cuBLAS-fp4 swizzle is then skipped
since that layer never runs the W4A4 tensor-core path.

Subclasses override to apply per-rank slicing first; at TP=1
every subclass calls straight through.

Repack the bound fp4 weight + fp8 block scale into fp4-Marlin's
layout, if the extension is available on this (Blackwell) GPU.

Best-effort: a missing extension (no nvcc + no prebaked ``.so``) or
any repack failure leaves ``_marlin_ready`` False, so the forward
transparently falls back to the Triton weight-only GEMM. The
non-negative-scale requirement of Marlin's S1E4M3→S0E5M3 recode is
asserted per layer inside :func:`prepare_fp4_layer_for_marlin`.

Dequantize to the dense ``(out_features, in_features)`` weight in
``F.linear`` layout (row-major over output).

Mirrors :meth:`AWQLinearBase.dequantized_weight` — the contract the
GDN block's fused dense ``[B|A]`` prebuild relies on
(:meth:`GDNBlock.prepare_fused_ba_dequant`): sub-tile GDN
projections (``in_proj_b`` / ``in_proj_a``) are concatenated into
one dense weight ONCE at build time, so this is a load-path helper,
never a hot path. Default dtype is bf16 (the serving dtype).

Raises (never silently degrades) when the raw fp4 buffers are gone —
a Marlin-repacked W4A16 layer released ``weight_packed`` at bind time
(see :meth:`_maybe_prepare_marlin`), so the dense weight is not
recoverable from this surface.

Slice this rank's output rows for one sub-block out of ``tensor``.

Both ``weight`` (``(out, in//2)`` packed fp4) and ``weight_scale``
(``(out, in//16)`` fp8 block scale) are output-major on dim 0 with
NO packing along dim 0, so the same row-range slice applies to
both. The rank's chunk within the sub-block is
``[rank*local_sub, (rank+1)*local_sub)``, shifted by the
sub-block's offset in the fused output dim.

Per-sub-block + per-rank slice of the fused NVFP4 tensors along
the OUTPUT axis (dim 0), concatenated in declared shard order, then
bound via the base load.

Delegates to :meth:`NVFP4LinearBase.nvfp4_load` directly (NOT
``super()``, which would re-run :class:`NVFP4ColumnParallelLinear`'s
contiguous output-axis slice and double-shard); the per-sub-block
slicing here already produced this rank's local weights.

Vendored fp4-Marlin GEMM for NVFP4 W4A16 — Python entry points.

The Blackwell (sm_120+) fast path for weight-only NVFP4 linears: a
fused fp4-weight × bf16-activation Marlin GEMM. It replaces the
arch-agnostic Triton kernel in
:mod:`arbi_serve.weight_quant.nvfp4.weight_only_gemm` (kept as the
fallback) with the same kernel vLLM/SGLang ship for weight-only NVFP4.

Two torch ops are exposed (both register into the vendored extension's
``_C`` namespace on first :func:`_ensure_loaded`; see
:mod:`arbi_serve.weight_quant.nvfp4.marlin.loader`):

  - :func:`gptq_marlin_repack` — repack a standard-order int32 qweight
    (packed along the INPUT axis, sequential bit order) into Marlin's
    interleaved layout. Used once per layer at load time.
  - :func:`marlin_gemm` — fused dequant + matmul on Marlin-laid-out fp4
    weights with per-16-block fp8 scales + a per-tensor global scale.
    Output dtype follows ``a.dtype`` (bf16).

The pure-PyTorch layout transforms (:func:`prepare_fp4_layer_for_marlin`,
:func:`apply_fp4_marlin_linear` and their helpers) are vendored verbatim
from vLLM's ``marlin_utils.py`` / ``marlin_utils_fp4.py`` — they are the
proven repack + scale-recode contract. Do NOT alter their math without
re-validating on Blackwell.

Vanilla bf16 / W4A4 runs never touch this path; it is imported lazily
(only when a W4A16 layer binds on Blackwell).

Bit-pack a ``vllm::ScalarType`` id — mirrors the C++
``ScalarType::id`` field layout (same helper the AWQ Marlin wrapper
uses). Field order LSB→MSB: exponent u8, mantissa u8, signed 1b,
bias i32, finite 1b, nan_repr u8.

True if the fp4-Marlin extension can be loaded (built or prebaked).

A single guarded, memoized probe: any failure (no nvcc + no prebaked
``.so``, header/arch mismatch, …) latches False so the caller falls
back to the Triton weight-only kernel without re-attempting the build.

Run fp4-Marlin's fused dequant + GEMM on NVFP4 weights.

Arg order matches the vendored ``_C::marlin_gemm`` schema; the unused
W4A4 / GPTQ / bias parameters (``c``, ``b_bias``, ``a_scales``,
``b_zeros``, ``g_idx``, ``perm``, ``is_zp_float``) are passed as
``None`` / ``False``.

Recode the per-block fp8 weight scales into Marlin's NVFP4 layout.

This is the mandatory S1E4M3 → S0E5M3 recode: the fp8_e4m3 block
scales are shifted 1 bit (dropping the sign) and re-viewed as
fp8_e4m3fn, keeping only the odd bytes. Marlin's NVFP4 path REQUIRES
the input scales to be non-negative (they already are for a valid
modelopt pack); the check makes a corrupt pack fail loud at load
rather than produce silent garbage.

Repack a bound NVFP4 weight bundle into fp4-Marlin's layout.

Args:
    weight_packed: ``(N, K // 2)`` uint8 packed fp4 (as bound by
        :meth:`NVFP4LinearBase.nvfp4_load`).
    weight_scale_fp8: ``(N, K // 16)`` fp8_e4m3fn per-block scale.
    weight_scale_2: fp32 scalar per-tensor global scale.
    out_features / in_features: N / K (post-TP-slice, per rank).

Returns ``(marlin_qweight, marlin_scales, marlin_global_scale,
padded_n, padded_k)``. The scale recode asserts non-negativity.

``y = x @ dequant(W).T`` via the fp4-Marlin GEMM.

``x`` is ``(..., K)`` bf16; returns ``(..., N)`` bf16. The K axis is
padded to Marlin's tile before the call and the N axis is trimmed
after — mirrors vLLM ``apply_fp4_marlin_linear``.

JIT / prebake loader for arbi-serve's vendored fp4-Marlin GEMM.

The fp4-Marlin .cu sources live under
:mod:`arbi_serve.weight_quant.nvfp4.marlin.csrc`. On first call to
:func:`_ensure_loaded` we hand the (small) file list to
:func:`torch.utils.cpp_extension.load`, which spawns nvcc, produces a
``.so``, and registers two Torch ops into the ``_C`` op namespace (the
schemas are in ``csrc/bindings.cpp``):

  - ``_C::gptq_marlin_repack``
  - ``_C::marlin_gemm``

Same three-file structure the AWQ Marlin backend uses
(:mod:`arbi_serve.weight_quant.awq.marlin.loader`), with two differences
that make it the NVFP4 W4A16 fast path rather than the INT4 AWQ one:

  1. **Stable-ABI sources.** The vendored kernels are vLLM's
     ``libtorch_stable`` Marlin variant (``torch::stable`` / stable-ABI
     bindings). This needs ``-DUSE_CUDA`` (else
     ``libtorch_stable/torch_utils.h`` can't find
     ``aoti_torch_get_current_cuda_stream``, which torch declares under
     ``#ifdef USE_CUDA``).
  2. **Only the bf16 NVFP4 W4A16 kernel** (``fe2m1f`` weight /
     ``fe4m3fn`` scale) is compiled in — 15 template instantiations,
     generated GPU-free by ``csrc/gen_nvfp4.py``.

Both extensions set ``-static-global-template-stub=false`` (nvcc-13
whole-program mode otherwise turns the cross-TU ``__global__`` Marlin
template instantiations into hidden stubs → link error).

Prebake. The runtime slim image ships this ``.so`` prebaked into
``/opt/cache-baked/`` (built in the -devel builder stage with
``TORCH_CUDA_ARCH_LIST=12.0``); :mod:`arbi_serve._prebake_loader`
short-circuits the ``cpp_extension.load`` below straight to
``torch.ops.load_library`` on the prebaked file, so the runtime needs
no nvcc for it. Unlike AWQ Marlin, the compile is GPU-free (nvcc emits
sm_120 cubins device-less from ``TORCH_CUDA_ARCH_LIST``), so it bakes in
the GPU-less builder exactly like the xgrammar / tkv kernels.

Memory hygiene. Never imported at boot. Only
:meth:`NVFP4LinearBase.nvfp4_load` (weight-only + Blackwell) and its
forward reach here, so CPU-only smoke + import-time tests stay fast.

JIT-compile (or short-circuit to the prebaked / cached) fp4-Marlin
extension and attach the Python-side fakes.

Unlike AWQ Marlin this does NOT require a live GPU: the compile is
device-free (nvcc emits sm_120 cubins from ``TORCH_CUDA_ARCH_LIST``),
so the same code path bakes the ``.so`` in the GPU-less builder.

Registering into ``_C`` via ``torch.ops.load_library`` is a side
effect of ``load(is_python_module=False)``; then the fakes let
``torch.compile`` / Inductor plan through the GEMM as a black box.

Weight-only NVFP4 GEMM (``W4A16``): fp4 weight × bf16 activation.

modelopt ships two NVFP4 flavours (see ``hf_quant_config.json``):

  * ``NVFP4`` — **W4A4**: both the weight and the activation are NVFP4.
    The activation is quantised to fp4 per forward and the cuBLAS
    Blackwell fp4 tensor-core MMA runs. This is the path in
    :meth:`NVFP4LinearBase._nvfp4_forward` and it is correct for packs
    the model was *trained/calibrated* for at A4.

  * ``W4A16_NVFP4`` — **weight-only**: the weight is 4-bit NVFP4 but the
    activation stays 16-bit. Running such a pack through the W4A4 path
    over-quantises the activation to 4 bits, which compounds across a
    deep stack into degenerate output. The activation MUST stay bf16.

Blackwell exposes no mixed fp4×bf16 tensor-core op (``torch._scaled_mm``
is symmetric-dtype), and the stack bundles no cutlass/marlin fp4 kernel,
so this module provides the weight-only GEMM directly, on top of the
format-agnostic kernel in
:mod:`arbi_serve.weight_quant.packed_fp4_gemm` (16-element blocks,
fp8_e4m3 block scale).

Layout (modelopt NVFP4, as bound by :meth:`NVFP4LinearBase.nvfp4_load`)::

    weight_packed    uint8   [N, K // 2]    2 E2M1 nibbles / byte
                                            LOW nibble  = real col 2j
                                            HIGH nibble = real col 2j+1
    weight_scale_fp8 fp8e4m3 [N, K // 16]   per-16-element block scale
    weight_scale_2   fp32    ()             per-tensor global scale

The real weight is ``level(nibble) * block_scale_fp8 * weight_scale_2``.
Because byte ``j`` holds real cols ``2j`` and ``2j+1`` — and
``(2j)//16 == (2j+1)//16 == j//8`` — both nibbles of a byte share block
scale column ``j // 8``. The kernel accumulates
``dot(x_even, W_low) + dot(x_odd, W_high)`` where the even/odd real-K
activation streams are read STRIDED from the single ``(M, K)`` tensor
(no pre-split copy). The scalar global scale is folded once into the
host-side epilogue, so the weight is never dequantised to a dense bf16
tensor — peak VRAM is just the ``(M, N)`` output.

The GEMM body itself is the format-agnostic kernel in
:mod:`arbi_serve.weight_quant.packed_fp4_gemm`, specialised here on
``SCALE_BLOCK=16`` / ``SCALE_KIND_FP8``; MXFP4 specialises the same kernel
on its own 32-element E8M0 grid.

The launcher is a ``torch.library.custom_op`` with a ``register_fake``
that mutates a caller-allocated fp32 accumulator, so the op itself
allocates nothing on the device and Inductor plans the accumulator.

Accumulate ``x @ dequant(W).T`` (block scale only) into ``out``.

Args:
    out: ``(M, N)`` fp32, pre-allocated and pre-zeroed (``SPLIT_K > 1``
        combines per-slice partials with ``atomic_add``).
    x: ``(M, K)`` bf16 activation.
    weight_packed: ``(N, K // 2)`` uint8 packed fp4.
    weight_scale_fp8: ``(N, K // 16)`` fp8e4m3 per-block scale.

``y = x @ dequant(W).T`` with a fp4-packed, bf16-activation GEMM.

Args:
    x: ``(..., K)`` bf16 activation (K = ``in_features``).
    weight_packed: ``(N, K // 2)`` uint8 packed fp4 (as bound).
    weight_scale_fp8: ``(N, K // 16)`` fp8e4m3 per-block scale (the
        UNswizzled buffer — swizzling is only for the cuBLAS fp4 path).
    weight_scale_2: fp32 scalar global scale.

Returns ``(..., N)`` bf16.

Weight-only packed-FP4 × bf16 GEMM shared by NVFP4 and MXFP4.

The weight stays FP4-packed in HBM (two E2M1 nibbles per ``uint8``, LOW
nibble = even real column) and each tile is decoded in registers before a
standard bf16 ``tl.dot``. No fp4 tensor cores are used, so the kernel runs
on any Triton-capable GPU.

Two ``tl.constexpr`` axes specialise it per format:

  * ``SCALE_BLOCK`` — real elements covered by one block scale along K
    (NVFP4 16, MXFP4 32). One packed byte holds real columns ``2j`` and
    ``2j+1``, which share block ``j // (SCALE_BLOCK // 2)``.
  * ``SCALE_KIND`` — :data:`SCALE_KIND_FP8` reads an ``fp8_e4m3`` block
    scale directly; :data:`SCALE_KIND_E8M0` reads a ``uint8`` biased
    exponent and decodes it to ``2**(byte - 127)``.

The per-tensor fp32 global scale NVFP4 carries is not applied here — it is
a scalar the caller folds into its epilogue. MXFP4 has none.

(BLOCK_M, BLOCK_N, BLOCK_KP, SPLIT_K, num_stages, num_warps).

``BLOCK_KP`` is the tile size along the PACKED-K (byte) axis, rounded up
to a multiple of ``scale_bytes`` (``SCALE_BLOCK // 2``) so a tile covers
whole scale blocks.

``SPLIT_K`` partitions the K reduction across CTAs: at decode (M≈1) a
narrow-N / deep-K Linear yields too few column tiles to fill the SMs, so
the reduction fans out to ``n_tiles × SPLIT_K`` CTAs whose partials are
combined via ``atomic_add``. It is sized to roughly fill the SMs and
capped so each split still owns at least two ``BLOCK_KP`` tiles.

Quant-backend registry.

Backends are registered at import time via :func:`register`; the
loader iterates the registry in :attr:`QuantBackend.detection_priority`
order (ascending — lower priority value runs first) and picks the
first backend whose :meth:`QuantBackend.detect` returns True.

The priority sort makes ``all_backends()`` deterministic regardless of
which subpackage imported what first. This matters when one backend's
detection key set is a strict subset of another's:

  - **NVFP4** (priority 10) — needs ``.weight`` + ``.weight_scale`` +
    ``.weight_scale_2``. Strictly tighter than FP8.
  - **AWQ** (priority 20) — needs ``.qweight`` + ``.qzeros`` +
    ``.scales`` (legacy) or ``.weight_packed`` + ``.weight_scale`` +
    ``.weight_shape`` (compressed-tensors pack-quantized).
  - **EXL3** (priority 30) — needs ``.trellis``.
  - **FP8** (priority 40) — needs ``.weight`` + ``.weight_scale``.
    Strictly looser than NVFP4; must run AFTER NVFP4 or it would
    misclaim NVFP4 packs.

To add a new backend:

  1. Implement :class:`arbi_serve.weight_quant.base.QuantBackend` with
     a ``detection_priority`` class attribute.
  2. Call :func:`register` from your subpackage's ``__init__``
     (e.g. ``arbi_serve/quant/awq/__init__.py``).
  3. Add ``arbi_serve.weight_quant.<name>`` to the
     :func:`_load_default_backends` import list below so a fresh
     ``import arbi_serve.weight_quant`` discovers it.

Add a backend to the registry. Idempotent on (name + class).

``detection_priority`` is required. Detection order decides which
backend claims a checkpoint whose key set is a subset of another's, so a
backend that leaves the priority to a default would be ordered by
something other than its own detection signature — and the wrong claim
surfaces as mis-decoded weights, not as an error. Refusing the
registration is what keeps that decision explicit.

Return registered backends sorted by ``detection_priority``
(ascending; lower runs first), with ``name`` as a tie-breaker for a
fully deterministic order.

The sort makes the returned order independent of the import
sequence: whichever subpackage registered first does not influence
detection priority. Read-only view — callers should not mutate.

Import-side-effect: ensure the bundled backends register.

Each backend's ``__init__`` calls :func:`register`. We import them
lazily — heavy upstream deps (exllamav3, awq_kernels, etc.) only
load when their backend is actually used. The import below is just
enough to run their registration call.

Detection ordering does NOT depend on this import order: each
backend declares its own ``detection_priority`` and
:func:`all_backends` sorts on that.
